[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.

Author: d-smirnov

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.

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