[Mlir-commits] [mlir] 65fd055 - [MLIR] Added check for IsTerminator trait (#79317)

[llvmlistbot at llvm.org]

Author: Dmitriy Smirnov
Date: 2024-01-26T14:14:54+01:00
New Revision: 65fd05517fe19bb436432ca45c03b7e576a328b5

URL: https://github.com/llvm/llvm-project/commit/65fd05517fe19bb436432ca45c03b7e576a328b5
DIFF: https://github.com/llvm/llvm-project/commit/65fd05517fe19bb436432ca45c03b7e576a328b5.diff

LOG: [MLIR] Added check for IsTerminator trait (#79317)

This PR adds a check for IsTerminator trait to prevent deletion of ops
like gpu.terminator as a "simple op" by RemoveDeadValues pass.

Added: 
    

Modified: 
    mlir/lib/Transforms/RemoveDeadValues.cpp
    mlir/test/Transforms/remove-dead-values.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index d80e514aed36f88..055256903a15228 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -594,13 +594,9 @@ void RemoveDeadValues::runOnOperation() {
       cleanFuncOp(funcOp, module, la);
     } else if (auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op)) {
       cleanRegionBranchOp(regionBranchOp, la);
-    } else if (op->hasTrait<OpTrait::ReturnLike>()) {
-      // Nothing to do because this terminator is associated with either a
-      // function op or a region branch op and gets cleaned when these ops are
-      // cleaned.
-    } else if (isa<RegionBranchTerminatorOpInterface>(op)) {
-      // Nothing to do because this terminator is associated with a region
-      // branch op and gets cleaned when the latter is cleaned.
+    } else if (op->hasTrait<::mlir::OpTrait::IsTerminator>()) {
+      // Nothing to do here because this is a terminator op and it should be
+      // honored with respect to its parent
     } else if (isa<CallOpInterface>(op)) {
       // Nothing to do because this op is associated with a function op and gets
       // cleaned when the latter is cleaned.

diff  --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 600810a785b1f9e..69426fdb6208329 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -335,3 +335,25 @@ func.func @main(%arg3 : i32, %arg4 : i1) {
   %non_live_0 = func.call @clean_region_branch_op_erase_it(%arg3, %arg4) : (i32, i1) -> (i32)
   return
 }
+
+// -----
+
+#map = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)>
+func.func @kernel(%arg0: memref<18xf32>) {
+  %c1 = arith.constant 1 : index
+  %c18 = arith.constant 18 : index
+  gpu.launch blocks(%arg3, %arg4, %arg5) in (%arg9 = %c18, %arg10 = %c18, %arg11 = %c18) threads(%arg6, %arg7, %arg8) in (%arg12 = %c1, %arg13 = %c1, %arg14 = %c1) {
+    %c1_0 = arith.constant 1 : index
+    %c0_1 = arith.constant 0 : index
+    %cst_2 = arith.constant 25.4669495 : f32
+    %6 = affine.apply #map(%arg3)[%c1_0, %c0_1]
+    memref.store %cst_2, %arg0[%6] : memref<18xf32>
+    gpu.terminator
+  } {SCFToGPU_visited}
+  return
+}
+
+// CHECK-LABEL: func.func @kernel(%arg0: memref<18xf32>) {
+// CHECK: gpu.launch blocks
+// CHECK: memref.store
+// CHECK-NEXT: gpu.terminator