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

[Dmitriy Smirnov]

https://github.com/d-smirnov updated https://github.com/llvm/llvm-project/pull/79317

>From 5aa03966af855d0ce13ea8ee07ba99d106b8a1cc Mon Sep 17 00:00:00 2001
From: Dmitriy Smirnov <dmitriy.smirnov at arm.com>
Date: Wed, 24 Jan 2024 14:38:08 +0000
Subject: [PATCH] [MLIR] Added check for IsTerminator trait

Add check for IsTerminator trait to prevent deletion of
ops like gpu.terminator as a "simple op" by RemoveDeadValues pass.
---
 mlir/lib/Transforms/RemoveDeadValues.cpp     |  3 +++
 mlir/test/Transforms/remove-dead-values.mlir | 22 ++++++++++++++++++++
 2 files changed, 25 insertions(+)

diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index d80e514aed36f88..6ff1cc40d43a994 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -598,6 +598,9 @@ void RemoveDeadValues::runOnOperation() {
       // 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 (op->hasTrait<::mlir::OpTrait::IsTerminator>()) {
+      // Nothing to do here because this a terminator op and it should be
+      // honored with respect to its parent
     } 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.
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