[mlir][scf] Fix scf.forall to scf.parallel pass walker #95385
Conversation
Adds proper walk results to the pass body to prevent runtime crashes on transformation failure.
|
@llvm/pr-subscribers-mlir Author: Adam Siemieniuk (adam-smnk) ChangesAdds proper walk results to the pass body to prevent runtime crashes on transformation failure. Full diff: https://github.com/llvm/llvm-project/pull/95385.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/SCF/Transforms/ForallToParallel.cpp b/mlir/lib/Dialect/SCF/Transforms/ForallToParallel.cpp
index 44e6840b03a3d..925d4a3c0a085 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ForallToParallel.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ForallToParallel.cpp
@@ -71,8 +71,9 @@ struct ForallToParallelLoop final
parentOp->walk([&](scf::ForallOp forallOp) {
if (failed(scf::forallToParallelLoop(rewriter, forallOp))) {
- return signalPassFailure();
+ return WalkResult::skip();
}
+ return WalkResult::advance();
});
}
};
diff --git a/mlir/test/Dialect/SCF/forall-to-parallel.mlir b/mlir/test/Dialect/SCF/forall-to-parallel.mlir
index acde601d47259..21e816956a094 100644
--- a/mlir/test/Dialect/SCF/forall-to-parallel.mlir
+++ b/mlir/test/Dialect/SCF/forall-to-parallel.mlir
@@ -78,3 +78,21 @@ func.func @mapping_attr() -> () {
return
}
+
+// -----
+
+// CHECK-LABEL: @forall_with_outputs
+// CHECK-SAME: %[[ARG0:.+]]: tensor<32x32xf32>
+func.func @forall_with_outputs(%arg0: tensor<32x32xf32>) -> tensor<8x112x32x32xf32> {
+ // CHECK-NOT: scf.parallel
+ // CHECK: %[[RES:.+]] = scf.forall{{.*}}shared_outs
+ // CHECK: return %[[RES]] : tensor<8x112x32x32xf32>
+
+ %0 = tensor.empty() : tensor<8x112x32x32xf32>
+ %1 = scf.forall (%arg1, %arg2) in (8, 112) shared_outs(%arg3 = %0) -> (tensor<8x112x32x32xf32>) {
+ scf.forall.in_parallel {
+ tensor.parallel_insert_slice %arg0 into %arg3[%arg1, %arg2, 0, 0] [1, 1, 32, 32] [1, 1, 1, 1] : tensor<32x32xf32> into tensor<8x112x32x32xf32>
+ }
+ }
+ return %1 : tensor<8x112x32x32xf32>
+}
|
|
@llvm/pr-subscribers-mlir-scf Author: Adam Siemieniuk (adam-smnk) ChangesAdds proper walk results to the pass body to prevent runtime crashes on transformation failure. Full diff: https://github.com/llvm/llvm-project/pull/95385.diff 2 Files Affected:
diff --git a/mlir/lib/Dialect/SCF/Transforms/ForallToParallel.cpp b/mlir/lib/Dialect/SCF/Transforms/ForallToParallel.cpp
index 44e6840b03a3d..925d4a3c0a085 100644
--- a/mlir/lib/Dialect/SCF/Transforms/ForallToParallel.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/ForallToParallel.cpp
@@ -71,8 +71,9 @@ struct ForallToParallelLoop final
parentOp->walk([&](scf::ForallOp forallOp) {
if (failed(scf::forallToParallelLoop(rewriter, forallOp))) {
- return signalPassFailure();
+ return WalkResult::skip();
}
+ return WalkResult::advance();
});
}
};
diff --git a/mlir/test/Dialect/SCF/forall-to-parallel.mlir b/mlir/test/Dialect/SCF/forall-to-parallel.mlir
index acde601d47259..21e816956a094 100644
--- a/mlir/test/Dialect/SCF/forall-to-parallel.mlir
+++ b/mlir/test/Dialect/SCF/forall-to-parallel.mlir
@@ -78,3 +78,21 @@ func.func @mapping_attr() -> () {
return
}
+
+// -----
+
+// CHECK-LABEL: @forall_with_outputs
+// CHECK-SAME: %[[ARG0:.+]]: tensor<32x32xf32>
+func.func @forall_with_outputs(%arg0: tensor<32x32xf32>) -> tensor<8x112x32x32xf32> {
+ // CHECK-NOT: scf.parallel
+ // CHECK: %[[RES:.+]] = scf.forall{{.*}}shared_outs
+ // CHECK: return %[[RES]] : tensor<8x112x32x32xf32>
+
+ %0 = tensor.empty() : tensor<8x112x32x32xf32>
+ %1 = scf.forall (%arg1, %arg2) in (8, 112) shared_outs(%arg3 = %0) -> (tensor<8x112x32x32xf32>) {
+ scf.forall.in_parallel {
+ tensor.parallel_insert_slice %arg0 into %arg3[%arg1, %arg2, 0, 0] [1, 1, 32, 32] [1, 1, 1, 1] : tensor<32x32xf32> into tensor<8x112x32x32xf32>
+ }
+ }
+ return %1 : tensor<8x112x32x32xf32>
+}
|
| parentOp->walk([&](scf::ForallOp forallOp) { | ||
| if (failed(scf::forallToParallelLoop(rewriter, forallOp))) { | ||
| return signalPassFailure(); | ||
| return WalkResult::skip(); |
There was a problem hiding this comment.
Is this a failure? Why remove the signalPassFailure()?
There is also likely a missing error message before signalPassFailure() here (we shouldn't fail silently).
There was a problem hiding this comment.
The result is not a real failure as it only occurs on match failure.
The main motivation for the change is that simply calling signalPassFailure() produces no output when the pass is called (at least from CLI). I'd expect the IR to remain unchanged in such case.
I think I should've captures the walk result and added some error on interruption. But there is no reason to interrupt on this error.
Perhaps a greedy rewriter could be better here instead of walking the graph manually.
There was a problem hiding this comment.
The scf::forallToParallelLoop function internally calls notifyMatchFailure, so some diagnostic should occur. That may not mean much if the pass terminates successfully though.
There was a problem hiding this comment.
notifyMatchFailure is a debug function. This is a question of semantics for the pass though, and unfortunately this pass does not even have a description!
Can we start here and document the pass behavior before changing it?
(is the pass promising to turn all ForAll to scf.parallel? Or it is opportunistically doing it? Under which conditions? etc)
There was a problem hiding this comment.
This is a question of semantics for the pass though, and unfortunately this pass does not even have a description!
The lack of description is my fault. My original intention was to error out when scf.forall cannot be lowered. I don't think it makes sense to run this transform before bufferization, and after bufferization all scf.forall operations should produce no results.
There was a problem hiding this comment.
Good points all together. My change was too eager too.
@sabauma My view on the pass is that indicating full failure (through signalPassFailure) is a bit heavy handed in this case (and viewed it as "error") but if that is the intention, it is equally valid approach.
I'll leave the pass as is. Perhaps the description could be explicit about the intended behavior.
|
Misinterpreted pass' intention which works as intended. No changes needed. |
Adds proper walk results to the pass body to prevent runtime crashes on transformation failure.