Skip to content

Commit 0cb024b

Browse files
[mlir][Mesh] Fix invalid IR in rewrite pattern (#78094)
This commit fixes `test/Dialect/Mesh/folding.mlir` when running with `MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS`. ``` /usr/local/google/home/springerm/mlir_public/llvm-project/mlir/test/Dialect/Mesh/folding.mlir:19:10: error: Unexpected number of results 0. Expected 2. %0:2 = mesh.cluster_shape @mesh1 : index, index ^ /usr/local/google/home/springerm/mlir_public/llvm-project/mlir/test/Dialect/Mesh/folding.mlir:19:10: note: see current operation: "mesh.cluster_shape"() <{axes = array<i16>, mesh = @mesh1}> : () -> () mlir-asm-printer: Verifying operation: builtin.module Unexpected number of results 0. Expected 2. mlir-asm-printer: 'builtin.module' failed to verify and will be printed in generic form "builtin.module"() ({ "mesh.cluster"() <{dim_sizes = array<i64: 2, 3>, rank = 2 : i64, sym_name = "mesh1"}> : () -> () "func.func"() <{function_type = () -> (index, index), sym_name = "cluster_shape_op_folding_all_axes_static_mesh"}> ({ %0 = "arith.constant"() <{value = 2 : index}> : () -> index %1 = "arith.constant"() <{value = 3 : index}> : () -> index "mesh.cluster_shape"() <{axes = array<i16>, mesh = @mesh1}> : () -> () %2:2 = "mesh.cluster_shape"() <{axes = array<i16>, mesh = @mesh1}> : () -> (index, index) "func.return"(%0, %1) : (index, index) -> () }) : () -> () }) : () -> () LLVM ERROR: IR failed to verify after pattern application ``` If `axes` is empty, the op verifier assumes that all dimensions are queried. (Expected 2 results.)
1 parent 844f833 commit 0cb024b

File tree

1 file changed

+7
-6
lines changed

1 file changed

+7
-6
lines changed

mlir/lib/Dialect/Mesh/Transforms/Simplifications.cpp

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -103,13 +103,14 @@ struct ClusterShapeFolder : OpRewritePattern<ClusterShapeOp> {
103103
}
104104

105105
// Leave only the dynamic mesh axes to be queried.
106-
ClusterShapeOp newShapeOp =
107-
builder.create<ClusterShapeOp>(mesh.getSymName(), newShapeOpMeshAxes);
108-
for (size_t i = 0; i < newShapeOp->getResults().size(); ++i) {
109-
newResults[newToOldResultsIndexMap[i]] = newShapeOp->getResults()[i];
106+
if (!newShapeOpMeshAxes.empty()) {
107+
ClusterShapeOp newShapeOp =
108+
builder.create<ClusterShapeOp>(mesh.getSymName(), newShapeOpMeshAxes);
109+
for (size_t i = 0; i < newShapeOp->getResults().size(); ++i) {
110+
newResults[newToOldResultsIndexMap[i]] = newShapeOp->getResults()[i];
111+
}
110112
}
111-
112-
rewriter.replaceAllUsesWith(op.getResults(), newResults);
113+
rewriter.replaceOp(op, newResults);
113114

114115
return success();
115116
}

0 commit comments

Comments
 (0)