Skip to content

Commit a7c9fe7

Browse files
fschlimbfschlimb
committed
removing source from UpdateHaloOp, because not required for destination passing style
1 parent 9578fc4 commit a7c9fe7

File tree

5 files changed

+23
-30
lines changed

5 files changed

+23
-30
lines changed

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

Lines changed: 7 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1095,8 +1095,7 @@ def Mesh_UpdateHaloOp : Mesh_Op<"update_halo", [
10951095
TypesMatchWith<
10961096
"result has same type as destination",
10971097
"result", "destination", "$_self">,
1098-
DeclareOpInterfaceMethods<SymbolUserOpInterface>,
1099-
AttrSizedOperandSegments
1098+
DeclareOpInterfaceMethods<SymbolUserOpInterface>
11001099
]> {
11011100
let summary = "Update halo data.";
11021101
let description = [{
@@ -1105,7 +1104,7 @@ def Mesh_UpdateHaloOp : Mesh_Op<"update_halo", [
11051104
on the remote devices. Changes might be caused by mutating operations
11061105
and/or if the new halo regions are larger than the existing ones.
11071106

1108-
Source and destination might have different halo sizes.
1107+
Destination is supposed to be initialized with the local data (not halos).
11091108

11101109
Assumes all devices hold tensors with same-sized halo data as specified
11111110
by `source_halo_sizes/static_source_halo_sizes` and
@@ -1117,25 +1116,21 @@ def Mesh_UpdateHaloOp : Mesh_Op<"update_halo", [
11171116

11181117
}];
11191118
let arguments = (ins
1120-
AnyTypeOf<[AnyNon0RankedMemRef, AnyNon0RankedTensor]>:$source,
11211119
AnyTypeOf<[AnyNon0RankedMemRef, AnyNon0RankedTensor]>:$destination,
11221120
FlatSymbolRefAttr:$mesh,
11231121
Mesh_MeshAxesArrayAttr:$split_axes,
1124-
Variadic<I64>:$source_halo_sizes,
1125-
DefaultValuedAttr<DenseI64ArrayAttr, "{}">:$static_source_halo_sizes,
1126-
Variadic<I64>:$destination_halo_sizes,
1127-
DefaultValuedAttr<DenseI64ArrayAttr, "{}">:$static_destination_halo_sizes
1122+
Variadic<I64>:$halo_sizes,
1123+
DefaultValuedAttr<DenseI64ArrayAttr, "{}">:$static_halo_sizes
11281124
);
11291125
let results = (outs
11301126
AnyTypeOf<[AnyNon0RankedMemRef, AnyNon0RankedTensor]>:$result
11311127
);
11321128
let assemblyFormat = [{
1133-
$source `into` $destination
1129+
$destination
11341130
`on` $mesh
11351131
`split_axes` `=` $split_axes
1136-
(`source_halo_sizes` `=` custom<DynamicIndexList>($source_halo_sizes, $static_source_halo_sizes)^)?
1137-
(`destination_halo_sizes` `=` custom<DynamicIndexList>($destination_halo_sizes, $static_destination_halo_sizes)^)?
1138-
attr-dict `:` type($source) `->` type($result)
1132+
(`halo_sizes` `=` custom<DynamicIndexList>($halo_sizes, $static_halo_sizes)^)?
1133+
attr-dict `:` type($result)
11391134
}];
11401135
let extraClassDeclaration = [{
11411136
MutableOperandRange getDpsInitsMutable() { return getDestinationMutable(); }

‎mlir/lib/Conversion/MeshToMPI/MeshToMPI.cpp‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -70,13 +70,13 @@ struct ConvertUpdateHaloOp
7070
cast<IntegerAttr>(v.get<Attribute>()).getInt()));
7171
};
7272

73-
auto array = op.getInput();
74-
auto rank = array.getType().getRank();
73+
auto array = op.getDestination();
74+
auto rank = cast<ShapedType>(array.getType()).getRank();
7575
auto opSplitAxes = op.getSplitAxes().getAxes();
7676
auto mesh = op.getMesh();
7777
auto meshOp = getMesh(op, symbolTableCollection);
7878
auto haloSizes = getMixedValues(op.getStaticHaloSizes(),
79-
op.getDynamicHaloSizes(), rewriter);
79+
op.getHaloSizes(), rewriter);
8080
// subviews need Index values
8181
for (auto &sz : haloSizes) {
8282
if (sz.is<Value>()) {
@@ -94,7 +94,7 @@ struct ConvertUpdateHaloOp
9494
auto currHaloDim = -1; // halo sizes are provided for split dimensions only
9595
// we need the actual shape to compute offsets and sizes
9696
for (auto i = 0; i < rank; ++i) {
97-
auto s = array.getType().getShape()[i];
97+
auto s = cast<ShapedType>(array.getType()).getShape()[i];
9898
if (ShapedType::isDynamic(s)) {
9999
shape[i] = rewriter.create<memref::DimOp>(loc, array, s).getResult();
100100
} else {
@@ -176,7 +176,7 @@ struct ConvertUpdateHaloOp
176176
auto hasTo = rewriter.create<arith::CmpIOp>(
177177
loc, arith::CmpIPredicate::sge, to, zero);
178178
auto buffer = rewriter.create<memref::AllocOp>(
179-
loc, dimSizes, array.getType().getElementType());
179+
loc, dimSizes, cast<ShapedType>(array.getType()).getElementType());
180180
// if has neighbor: copy halo data from array to buffer and send
181181
rewriter.create<scf::IfOp>(
182182
loc, hasTo, [&](OpBuilder &builder, Location loc) {

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

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -496,11 +496,9 @@ tryUpdateHaloInResharding(ImplicitLocOpBuilder &builder, MeshOp mesh,
496496
sourceShard.getLoc(),
497497
RankedTensorType::get(outShape,
498498
sourceShard.getType().getElementType()),
499-
sourceShard, initOprnd, mesh.getSymName(),
499+
initOprnd, mesh.getSymName(),
500500
MeshAxesArrayAttr::get(builder.getContext(),
501501
sourceSharding.getSplitAxes()),
502-
sourceSharding.getDynamicHaloSizes(),
503-
sourceSharding.getStaticHaloSizes(),
504502
targetSharding.getDynamicHaloSizes(),
505503
targetSharding.getStaticHaloSizes())
506504
.getResult();

‎mlir/test/Dialect/Mesh/ops.mlir‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -615,16 +615,16 @@ func.func @update_halo(
615615
// CHECK-SAME: %[[ARG:.*]]: memref<12x12xi8>
616616
%arg0 : memref<12x12xi8>) {
617617
// CHECK-NEXT: %[[C2:.*]] = arith.constant 2 : i64
618-
// CHECK-NEXT: %[[UH1:.*]] = mesh.update_halo %[[ARG]] into %[[ARG]] on @mesh0
618+
// CHECK-NEXT: %[[UH1:.*]] = mesh.update_halo %[[ARG]] on @mesh0
619619
// CHECK-SAME: split_axes = {{\[\[}}0]]
620-
// CHECK-SAME: halo_sizes = [2, %c2_i64] : memref<12x12xi8> -> memref<12x12xi8>
620+
// CHECK-SAME: halo_sizes = [2, %c2_i64] : memref<12x12xi8>
621621
%c2 = arith.constant 2 : i64
622-
%uh1 = mesh.update_halo %arg0 into %arg0 on @mesh0 split_axes = [[0]]
623-
source_halo_sizes = [2, %c2] : memref<12x12xi8> -> memref<12x12xi8>
624-
// CHECK-NEXT: %[[UH2:.*]] = mesh.update_halo %[[ARG]] into %[[UH1]] on @mesh0
622+
%uh1 = mesh.update_halo %arg0 on @mesh0 split_axes = [[0]]
623+
halo_sizes = [2, %c2] : memref<12x12xi8>
624+
// CHECK-NEXT: %[[UH2:.*]] = mesh.update_halo %[[UH1]] on @mesh0
625625
// CHECK-SAME: split_axes = {{\[\[}}0], [1]]
626-
// CHECK-SAME: halo_sizes = [2, 2, %[[C2]], 2] : memref<12x12xi8> -> memref<12x12xi8>
627-
%uh2 = mesh.update_halo %arg0 into %uh1 on @mesh0 split_axes = [[0], [1]]
628-
source_halo_sizes = [2, 2, %c2, 2] : memref<12x12xi8> -> memref<12x12xi8>
626+
// CHECK-SAME: halo_sizes = [2, 2, %[[C2]], 2] : memref<12x12xi8>
627+
%uh2 = mesh.update_halo %uh1 on @mesh0 split_axes = [[0], [1]]
628+
halo_sizes = [2, 2, %c2, 2] : memref<12x12xi8>
629629
return
630630
}

‎mlir/test/Dialect/Mesh/spmdization.mlir‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -226,7 +226,7 @@ func.func @test_shard_update_halo(%arg0: tensor<1200x1200xi64>) -> tensor<1200x1
226226
%sharding = mesh.sharding @mesh_1d_4 split_axes = [[0]] : !mesh.sharding
227227
// CHECK: %[[T:.*]] = tensor.empty() : tensor<304x1200xi64>
228228
// CHECK: %[[inserted_slice:.*]] = tensor.insert_slice %[[IN1]] into %[[T]][2, 0] [300, 1200] [1, 1] : tensor<300x1200xi64> into tensor<304x1200xi64>
229-
// CHECK: %[[UH:.*]] = mesh.update_halo %[[IN1]] into %[[inserted_slice]] on @mesh_1d_4 split_axes = {{\[\[0]]}} destination_halo_sizes = [2, 2] : tensor<300x1200xi64> -> tensor<304x1200xi64>
229+
// CHECK: %[[UH:.*]] = mesh.update_halo %[[inserted_slice]] on @mesh_1d_4 split_axes = {{\[\[0]]}} halo_sizes = [2, 2] : tensor<304x1200xi64>
230230
%sharding_annotated = mesh.shard %arg0 to %sharding : tensor<1200x1200xi64>
231231
%sharding_0 = mesh.sharding @mesh_1d_4 split_axes = [[0]] halo_sizes = [2, 2] : !mesh.sharding
232232
%sharding_annotated_1 = mesh.shard %sharding_annotated to %sharding_0 : tensor<1200x1200xi64>
@@ -242,7 +242,7 @@ func.func @test_shard_update_halo2d(%arg0: tensor<1200x1200xi64>) -> tensor<1200
242242
%sharding = mesh.sharding @mesh4x4 split_axes = [[0], [1]] : !mesh.sharding
243243
// CHECK: %[[T:.*]] = tensor.empty() : tensor<303x307xi64>
244244
// CHECK: %[[inserted_slice:.*]] = tensor.insert_slice %[[IN1]] into %[[T]][1, 3] [300, 300] [1, 1] : tensor<300x300xi64> into tensor<303x307xi64>
245-
// CHECK: %[[UH:.*]] = mesh.update_halo %[[IN1]] into %[[inserted_slice]] on @mesh4x4 split_axes = {{\[\[}}0], [1]] destination_halo_sizes = [1, 2, 3, 4] : tensor<300x300xi64> -> tensor<303x307xi64>
245+
// CHECK: %[[UH:.*]] = mesh.update_halo %[[inserted_slice]] on @mesh4x4 split_axes = {{\[\[}}0], [1]] halo_sizes = [1, 2, 3, 4] : tensor<303x307xi64>
246246
%sharding_annotated = mesh.shard %arg0 to %sharding : tensor<1200x1200xi64>
247247
%sharding_0 = mesh.sharding @mesh4x4 split_axes = [[0], [1]] halo_sizes = [1, 2, 3, 4] : !mesh.sharding
248248
%sharding_annotated_1 = mesh.shard %sharding_annotated to %sharding_0 : tensor<1200x1200xi64>

0 commit comments

Comments
 (0)