Skip to content

Commit 43e1a5a

Browse files
fschlimbfschlimb
authored
[mlir][mesh] adding option for traversal order in sharding propagation (#144079)
The traversal order in sharding propagation was hard-coded. This PR provides options to the pass to select a suitable order - forward-only - backward-only - forward-backward - backward-forward Default is the previous behavior (backward-forward).
1 parent 669627d commit 43e1a5a

File tree

8 files changed

+173
-28
lines changed

8 files changed

+173
-28
lines changed

‎mlir/include/mlir/Dialect/Mesh/IR/MeshOps.h

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -206,9 +206,6 @@ Type shardType(Type type, MeshOp mesh, MeshSharding sharding);
206206
// Use newShardOp if it is not null. Otherwise create a new one.
207207
// May insert resharding if required.
208208
// Potentially updates newShardOp.
209-
void maybeInsertTargetShardingAnnotation(MeshSharding sharding,
210-
OpOperand &operand, OpBuilder &builder,
211-
ShardOp &newShardOp);
212209
void maybeInsertTargetShardingAnnotation(MeshSharding sharding, OpResult result,
213210
OpBuilder &builder);
214211
void maybeInsertSourceShardingAnnotation(MeshSharding sharding,

‎mlir/include/mlir/Dialect/Mesh/Transforms/Passes.h

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,18 @@ class FuncOp;
1919

2020
namespace mesh {
2121

22+
/// This enum controls the traversal order for the sharding propagation.
23+
enum class TraversalOrder {
24+
/// Forward traversal.
25+
Forward,
26+
/// Backward traversal.
27+
Backward,
28+
/// Forward then backward traversal.
29+
ForwardBackward,
30+
/// Backward then forward traversal.
31+
BackwardForward
32+
};
33+
2234
//===----------------------------------------------------------------------===//
2335
// Passes
2436
//===----------------------------------------------------------------------===//

‎mlir/include/mlir/Dialect/Mesh/Transforms/Passes.td

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,21 @@ def ShardingPropagation : InterfacePass<"sharding-propagation", "mlir::FunctionO
2424
operation, and the operations themselves are added with sharding option
2525
attributes.
2626
}];
27+
let options = [
28+
Option<"traversal", "traversal",
29+
"mlir::mesh::TraversalOrder", /*default=*/"mlir::mesh::TraversalOrder::BackwardForward",
30+
"Traversal order to use for sharding propagation:",
31+
[{::llvm::cl::values(
32+
clEnumValN(mlir::mesh::TraversalOrder::Forward, "forward",
33+
"Forward only traversal."),
34+
clEnumValN(mlir::mesh::TraversalOrder::Backward, "backward",
35+
"backward only traversal."),
36+
clEnumValN(mlir::mesh::TraversalOrder::ForwardBackward, "forward-backward",
37+
"forward-backward traversal."),
38+
clEnumValN(mlir::mesh::TraversalOrder::BackwardForward, "backward-forward",
39+
"backward-forward traversal.")
40+
)}]>,
41+
];
2742
let dependentDialects = [
2843
"mesh::MeshDialect"
2944
];

‎mlir/lib/Dialect/Mesh/IR/MeshOps.cpp

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -275,13 +275,12 @@ Type mesh::shardType(Type type, MeshOp mesh, MeshSharding sharding) {
275275
return type;
276276
}
277277

278-
void mlir::mesh::maybeInsertTargetShardingAnnotation(MeshSharding sharding,
279-
OpOperand &operand,
280-
OpBuilder &builder,
281-
ShardOp &newShardOp) {
278+
static void maybeInsertTargetShardingAnnotationImpl(MeshSharding sharding,
279+
Value &operandValue,
280+
Operation *operandOp,
281+
OpBuilder &builder,
282+
ShardOp &newShardOp) {
282283
OpBuilder::InsertionGuard insertionGuard(builder);
283-
Value operandValue = operand.get();
284-
Operation *operandOp = operand.getOwner();
285284
builder.setInsertionPointAfterValue(operandValue);
286285
ShardOp shardOp = dyn_cast<ShardOp>(operandOp);
287286
if (shardOp && sharding == shardOp.getSharding() &&
@@ -300,9 +299,8 @@ void mlir::mesh::maybeInsertTargetShardingAnnotation(MeshSharding sharding,
300299
builder.create<ShardOp>(operandValue.getLoc(), operandValue, shardingOp,
301300
/*annotate_for_users*/ false);
302301
}
303-
IRRewriter rewriter(builder);
304-
rewriter.replaceUsesWithIf(
305-
operandValue, newShardOp, [operandOp, operandValue](OpOperand &use) {
302+
operandValue.replaceUsesWithIf(
303+
newShardOp, [operandOp, operandValue](OpOperand &use) {
306304
return use.getOwner() == operandOp && use.get() == operandValue;
307305
});
308306

@@ -313,15 +311,20 @@ void mlir::mesh::maybeInsertTargetShardingAnnotation(MeshSharding sharding,
313311
auto newShardOp2 = builder.create<ShardOp>(operandValue.getLoc(), newShardOp,
314312
newShardOp.getSharding(),
315313
/*annotate_for_users*/ true);
316-
rewriter.replaceAllUsesExcept(newShardOp, newShardOp2, newShardOp2);
314+
newShardOp.getResult().replaceAllUsesExcept(newShardOp2, newShardOp2);
317315
}
318316

319317
void mlir::mesh::maybeInsertTargetShardingAnnotation(MeshSharding sharding,
320318
OpResult result,
321319
OpBuilder &builder) {
322320
ShardOp newShardOp;
323-
for (auto &use : llvm::make_early_inc_range(result.getUses())) {
324-
maybeInsertTargetShardingAnnotation(sharding, use, builder, newShardOp);
321+
SmallVector<std::pair<Value, Operation *>> uses;
322+
for (auto &use : result.getUses()) {
323+
uses.emplace_back(use.get(), use.getOwner());
324+
}
325+
for (auto &[operandValue, operandOp] : uses) {
326+
maybeInsertTargetShardingAnnotationImpl(sharding, operandValue, operandOp,
327+
builder, newShardOp);
325328
}
326329
}
327330

‎mlir/lib/Dialect/Mesh/Transforms/ShardingPropagation.cpp

Lines changed: 29 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -362,6 +362,9 @@ static LogicalResult visitOp(Operation *op, OpBuilder &builder) {
362362
//===----------------------------------------------------------------------===//
363363
struct ShardingPropagation
364364
: public mesh::impl::ShardingPropagationBase<ShardingPropagation> {
365+
366+
using ShardingPropagationBase<ShardingPropagation>::ShardingPropagationBase;
367+
365368
void runOnOperation() override {
366369
FunctionOpInterface funcOp = getOperation();
367370
MLIRContext *ctx = funcOp.getContext();
@@ -382,18 +385,31 @@ struct ShardingPropagation
382385
shardingOp.printLoopTypesAndIndexingMaps(llvm::dbgs());
383386
});
384387

385-
// 1. propagate in reversed order
386-
for (Operation &op : llvm::make_early_inc_range(llvm::reverse(block)))
387-
if (failed(visitOp(&op, builder)))
388-
return signalPassFailure();
389-
390-
LLVM_DEBUG(DBGS() << "After reversed order propagation:\n"
391-
<< funcOp << "\n");
392-
LLVM_DEBUG(assert(succeeded(mlir::verify(funcOp))));
393-
394-
// 2. propagate in original order
395-
for (Operation &op : llvm::make_early_inc_range(block))
396-
if (failed(visitOp(&op, builder)))
397-
return signalPassFailure();
388+
auto traverse = [&](auto &&range, OpBuilder &builder,
389+
const char *order) -> bool {
390+
for (Operation &op : range) {
391+
if (failed(visitOp(&op, builder))) {
392+
signalPassFailure();
393+
return true;
394+
}
395+
}
396+
LLVM_DEBUG(DBGS() << "After " << order << " order propagation:\n"
397+
<< funcOp << "\n");
398+
LLVM_DEBUG(assert(succeeded(mlir::verify(funcOp))));
399+
return false;
400+
};
401+
402+
// 1. Propagate in reversed order.
403+
if (traversal == TraversalOrder::Backward ||
404+
traversal == TraversalOrder::BackwardForward)
405+
traverse(llvm::reverse(block), builder, "backward");
406+
407+
// 2. Propagate in original order.
408+
if (traversal != TraversalOrder::Backward)
409+
traverse(block, builder, "forward");
410+
411+
// 3. Propagate in backward order if needed.
412+
if (traversal == TraversalOrder::ForwardBackward)
413+
traverse(llvm::reverse(block), builder, "backward");
398414
}
399415
};

‎mlir/test/Dialect/Mesh/backward-sharding-propagation.mlir

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
// RUN: mlir-opt --pass-pipeline="builtin.module(func.func(sharding-propagation{traversal=backward}))" %s | FileCheck %s
2+
3+
#map = affine_map<(d0, d1) -> (d0, d1)>
4+
module {
5+
mesh.mesh @mesh(shape = 1) {sym_visibility = "private"}
6+
func.func @test_forward() -> tensor<6x6xi32> {
7+
%c1_i32 = arith.constant 1 : i32
8+
// CHECK: tensor.empty()
9+
%0 = tensor.empty() : tensor<6x6xi32>
10+
%sharding = mesh.sharding @mesh split_axes = [[0]] : !mesh.sharding
11+
// CHECK-COUNT-2: mesh.shard
12+
%sharding_annotated = mesh.shard %0 to %sharding : tensor<6x6xi32>
13+
%1 = linalg.fill ins(%c1_i32 : i32) outs(%sharding_annotated : tensor<6x6xi32>) -> tensor<6x6xi32>
14+
// CHECK: tensor.empty()
15+
// CHECK-NOT: mesh.shard @
16+
%2 = tensor.empty() : tensor<6x6xi32>
17+
%3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%1, %1
18+
: tensor<6x6xi32>, tensor<6x6xi32>) outs(%2 : tensor<6x6xi32>) {
19+
^bb0(%in: i32, %in_2: i32, %out: i32):
20+
%9 = arith.addi %in, %in_2 : i32
21+
linalg.yield %9 : i32
22+
} -> tensor<6x6xi32>
23+
// CHECK: return
24+
return %3 : tensor<6x6xi32>
25+
}
26+
}

‎mlir/test/Dialect/Mesh/forward-backward-sharding-propagation.mlir

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,27 @@
1+
// RUN: mlir-opt --pass-pipeline="builtin.module(func.func(sharding-propagation{traversal=forward-backward}))" %s | FileCheck %s
2+
3+
#map = affine_map<(d0, d1) -> (d0, d1)>
4+
module {
5+
mesh.mesh @mesh(shape = 1) {sym_visibility = "private"}
6+
func.func @test_forward() -> tensor<6x6xi32> {
7+
%c1_i32 = arith.constant 1 : i32
8+
// CHECK: tensor.empty()
9+
%0 = tensor.empty() : tensor<6x6xi32>
10+
// CHECK-COUNT-3: mesh.sharding @mesh split_axes = {{\[\[0}}]]
11+
%sharding_row = mesh.sharding @mesh split_axes = [[0]] : !mesh.sharding
12+
%annotated_row = mesh.shard %0 to %sharding_row : tensor<6x6xi32>
13+
%1 = linalg.fill ins(%c1_i32 : i32) outs(%annotated_row : tensor<6x6xi32>) -> tensor<6x6xi32>
14+
%2 = tensor.empty() : tensor<6x6xi32>
15+
// CHECK-COUNT-4: mesh.sharding @mesh split_axes = {{\[\[1}}]]
16+
%3 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%2, %1
17+
: tensor<6x6xi32>, tensor<6x6xi32>) outs(%2 : tensor<6x6xi32>) {
18+
^bb0(%in: i32, %in_2: i32, %out: i32):
19+
%9 = arith.addi %in, %in_2 : i32
20+
linalg.yield %9 : i32
21+
} -> tensor<6x6xi32>
22+
%sharding_col = mesh.sharding @mesh split_axes = [[1]] : !mesh.sharding
23+
%annotated_col = mesh.shard %3 to %sharding_col : tensor<6x6xi32>
24+
// CHECK: return
25+
return %annotated_col : tensor<6x6xi32>
26+
}
27+
}

‎mlir/test/Dialect/Mesh/forward-sharding-propagation.mlir

Lines changed: 49 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
// RUN: mlir-opt --pass-pipeline="builtin.module(func.func(sharding-propagation{traversal=forward}))" %s | FileCheck %s
2+
3+
#map = affine_map<(d0, d1) -> (d0, d1)>
4+
module attributes {dlti.map = #dlti.map<"MPI:Implementation" = "mpich", "MPI:comm_world_rank" = 0 : i32>} {
5+
mesh.mesh @mesh(shape = 1) {sym_visibility = "private"}
6+
func.func @test_forward() -> (tensor<6x6xi32>, tensor<6x6xi32>, tensor<i32>) attributes {llvm.emit_c_interface} {
7+
%c1_i32 = arith.constant 1 : i32
8+
// CHECK: [[v3:%.*]] = tensor.empty() : tensor<6x6xi32>
9+
%0 = tensor.empty() : tensor<6x6xi32>
10+
// CHECK: [[v1:%.*]] = linalg.fill ins
11+
// CHECK: [[vsharding_0:%.*]] = mesh.sharding @mesh split_axes = {{\[\[}}0]] : !mesh.sharding
12+
// CHECK: [[vsharding_annotated_1:%.*]] = mesh.shard [[v1]] to [[vsharding_0]] : tensor<6x6xi32>
13+
%1 = linalg.fill ins(%c1_i32 : i32) outs(%0 : tensor<6x6xi32>) -> tensor<6x6xi32>
14+
%sharding = mesh.sharding @mesh split_axes = [[0]] : !mesh.sharding
15+
%sharding_annotated = mesh.shard %1 to %sharding : tensor<6x6xi32>
16+
// CHECK: [[v2:%.*]] = tensor.empty() : tensor<6x6xi32>
17+
// CHECK: [[vsharding_2:%.*]] = mesh.sharding @mesh split_axes = {{\[\[}}0]] : !mesh.sharding
18+
// CHECK: [[vsharding_annotated_3:%.*]] = mesh.shard [[vsharding_annotated_1]] to [[vsharding_2]] annotate_for_users : tensor<6x6xi32>
19+
%3 = tensor.empty() : tensor<6x6xi32>
20+
// CHECK: [[vsharding_4:%.*]] = mesh.sharding @mesh split_axes = {{\[\[}}0]] : !mesh.sharding
21+
// CHECK: [[vsharding_annotated_5:%.*]] = mesh.shard [[v2]] to [[vsharding_4]] annotate_for_users : tensor<6x6xi32>
22+
// CHECK: [[v3:%.*]] = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]}
23+
// CHECK-SAME: ins([[vsharding_annotated_3]], [[vsharding_annotated_3]] : tensor<6x6xi32>, tensor<6x6xi32>) outs([[vsharding_annotated_5]] : tensor<6x6xi32>) {
24+
// CHECK: [[vsharding_6:%.*]] = mesh.sharding @mesh split_axes = {{\[\[}}0]] : !mesh.sharding
25+
// CHECK: [[vsharding_annotated_7:%.*]] = mesh.shard [[v3]] to [[vsharding_6]] : tensor<6x6xi32>
26+
%4 = linalg.generic {indexing_maps = [#map, #map, #map], iterator_types = ["parallel", "parallel"]} ins(%sharding_annotated, %sharding_annotated
27+
: tensor<6x6xi32>, tensor<6x6xi32>) outs(%3 : tensor<6x6xi32>) {
28+
^bb0(%in: i32, %in_2: i32, %out: i32):
29+
%9 = arith.addi %in, %in_2 : i32
30+
linalg.yield %9 : i32
31+
} -> tensor<6x6xi32>
32+
%c0_i32 = arith.constant 0 : i32
33+
%6 = tensor.empty() : tensor<i32>
34+
%7 = linalg.fill ins(%c0_i32 : i32) outs(%6 : tensor<i32>) -> tensor<i32>
35+
// CHECK: [[vreduced:%.*]] = linalg.reduce ins
36+
// CHECK: [[vsharding_12:%.*]] = mesh.sharding @mesh split_axes = [] partial = sum [0] : !mesh.sharding
37+
// CHECK: [[vsharding_annotated_13:%.*]] = mesh.shard [[vreduced]] to [[vsharding_12]] : tensor<i32>
38+
%reduced = linalg.reduce ins(%4 : tensor<6x6xi32>) outs(%7 : tensor<i32>) dimensions = [0, 1]
39+
(%in: i32, %init: i32) {
40+
%9 = arith.addi %in, %init : i32
41+
linalg.yield %9 : i32
42+
}
43+
// CHECK: [[vsharding_14:%.*]] = mesh.sharding @mesh split_axes = {{\[\[}}]] : !mesh.sharding
44+
%sharding_0 = mesh.sharding @mesh split_axes = [[]] : !mesh.sharding
45+
// CHECK: [[vsharding_annotated_15:%.*]] = mesh.shard [[vsharding_annotated_13]] to [[vsharding_14]] annotate_for_users : tensor<i32>
46+
%sharding_annotated_1 = mesh.shard %reduced to %sharding_0 annotate_for_users : tensor<i32>
47+
return %sharding_annotated, %4, %sharding_annotated_1 : tensor<6x6xi32>, tensor<6x6xi32>, tensor<i32>
48+
}
49+
}

0 commit comments

Comments
 (0)