[Mlir-commits] [mlir] Keep attribute when bufferizing `scf.forall` op (PR #91236)

[llvmlistbot at llvm.org]

https://github.com/srcarroll created https://github.com/llvm/llvm-project/pull/91236

None

>From e3537e7d52bec2f6d0ffbaad5827929c62d7a2ec Mon Sep 17 00:00:00 2001
From: Sam <srcarroll314 at gmail.com>
Date: Mon, 6 May 2024 11:39:24 -0500
Subject: [PATCH] Keep attribute when bufferizing `scf.forall` op

---
 .../lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp | 3 +++
 mlir/test/Dialect/SCF/one-shot-bufferize.mlir                  | 3 ++-
 2 files changed, 5 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
index 2a16b10bbaf8ef..cf40443ff38390 100644
--- a/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp
@@ -1267,6 +1267,9 @@ struct ForallOpInterface
         forallOp.getMixedUpperBound(), forallOp.getMixedStep(),
         /*outputs=*/ValueRange(), forallOp.getMapping());
 
+    // Keep discardable attributes from the original op.
+    newForallOp->setDiscardableAttrs(op->getDiscardableAttrDictionary());
+
     rewriter.eraseOp(newForallOp.getBody()->getTerminator());
 
     // Move over block contents of the old op.
diff --git a/mlir/test/Dialect/SCF/one-shot-bufferize.mlir b/mlir/test/Dialect/SCF/one-shot-bufferize.mlir
index 485fdd9b0e5932..bb9f7dfdba83f7 100644
--- a/mlir/test/Dialect/SCF/one-shot-bufferize.mlir
+++ b/mlir/test/Dialect/SCF/one-shot-bufferize.mlir
@@ -499,7 +499,8 @@ func.func @parallel_insert_slice_no_conflict(
         tensor.parallel_insert_slice %8 into %o[5] [%idx] [%c1] :
           tensor<?xf32> into tensor<?xf32>
       }
-  }
+  } {keep_this_attribute}
+  // CHECK: keep_this_attribute
 
   // CHECK: %[[load:.*]] = memref.load %[[arg2]]
   %f = tensor.extract %2[%c0] : tensor<?xf32>