[Mlir-commits] [mlir] [OpenACC] Add an attribute to record number of loops collapsed. (PR #201969)

[Moazin K.]

https://github.com/moazin updated https://github.com/llvm/llvm-project/pull/201969

>From c2fe9eccb5c53b7db43756651a21e2e5117072ed Mon Sep 17 00:00:00 2001
From: Moazin Khatti <mkhatti at nvidia.com>
Date: Fri, 5 Jun 2026 14:50:18 -0700
Subject: [PATCH] [OpenACC] Add an attribute to record number of loops
 collapsed.

Attach an attribute indicating the number of collapsed loops if
`convertACCLoopToSCFFor` collapses the loops.

Assisted by Claude Code.
---
 mlir/include/mlir/Dialect/OpenACC/OpenACC.h     |  6 ++++++
 .../mlir/Dialect/OpenACC/OpenACCUtilsLoop.h     |  7 +++++++
 .../Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp  | 17 ++++++++++++++++-
 .../OpenACC/acc-compute-lowering-loop.mlir      |  1 +
 .../Dialect/OpenACC/OpenACCUtilsLoopTest.cpp    | 10 ++++++++++
 5 files changed, 40 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACC.h b/mlir/include/mlir/Dialect/OpenACC/OpenACC.h
index d5cd04949c751..650f6f780273f 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACC.h
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACC.h
@@ -206,6 +206,12 @@ static constexpr StringLiteral getFromDefaultClauseAttrName() {
   return StringLiteral("acc.from_default");
 }
 
+/// Name for an attribute attached to a loop indicating the number of loops
+/// collapsed to create that loop
+static constexpr StringLiteral getCollapseCountAttrName() {
+  return StringLiteral("acc.collapse_count");
+}
+
 static constexpr StringLiteral getVarNameAttrName() {
   return VarNameAttr::name;
 }
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsLoop.h b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsLoop.h
index faaa8231e696e..a08d11a7ec8ab 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsLoop.h
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCUtilsLoop.h
@@ -75,6 +75,13 @@ scf::ExecuteRegionOp
 convertUnstructuredACCLoopToSCFExecuteRegion(LoopOp loopOp,
                                              RewriterBase &rewriter);
 
+/// Record on a collapsed loop how many original loops were folded into it.
+void setCollapseCountAttr(Operation *op, uint64_t count);
+
+/// Number of original loops collapsed into op, or 1 when op carries no
+/// collapse_count attribute.
+uint64_t getCollapseCount(Operation *op);
+
 } // namespace acc
 } // namespace mlir
 
diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
index d9afba7c2daaf..fd46a2800e03d 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
@@ -249,9 +249,13 @@ scf::ForOp convertACCLoopToSCFFor(LoopOp loopOp, RewriterBase &rewriter,
   }
 
   // Optionally collapse nested loops
-  if (enableCollapse && forOps.size() > 1)
+  if (enableCollapse && forOps.size() > 1) {
+    unsigned numCollapsed = forOps.size();
     if (failed(coalesceLoops(rewriter, forOps)))
       loopOp.emitError("failed to collapse acc.loop");
+    else
+      setCollapseCountAttr(forOps.front(), numCollapsed);
+  }
 
   return forOps.front();
 }
@@ -334,5 +338,16 @@ convertUnstructuredACCLoopToSCFExecuteRegion(LoopOp loopOp,
                                                   loopOp->getLoc(), rewriter);
 }
 
+void setCollapseCountAttr(Operation *op, uint64_t count) {
+  op->setAttr(getCollapseCountAttrName(),
+              IntegerAttr::get(IntegerType::get(op->getContext(), 64), count));
+}
+
+uint64_t getCollapseCount(Operation *op) {
+  if (auto attr = op->getAttrOfType<IntegerAttr>(getCollapseCountAttrName()))
+    return attr.getValue().getZExtValue();
+  return 1;
+}
+
 } // namespace acc
 } // namespace mlir
diff --git a/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir b/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
index 358d1328fa028..2f9276e8525bb 100644
--- a/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
@@ -68,6 +68,7 @@ func.func @parallel_loop_auto_collapse(%buf: memref<1xi32>, %lb0 : index, %ub0 :
   // CHECK: scf.for
   // CHECK-NOT: scf.for
   // CHECK-NOT: scf.parallel
+  // CHECK: acc.collapse_count = 2 : i64
   acc.parallel dataOperands(%dev : memref<1xi32>) {
     acc.loop control(%i : index, %j : index) = (%lb0, %lb1 : index, index) to (%ub0, %ub1 : index, index) step (%c1, %c1 : index, index) {
       %vi = arith.index_cast %i : index to i32
diff --git a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp
index 4296d56dfa4d9..b71affa8ded11 100644
--- a/mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp
+++ b/mlir/unittests/Dialect/OpenACC/OpenACCUtilsLoopTest.cpp
@@ -271,6 +271,13 @@ TEST_F(OpenACCUtilsLoopTest, ConvertLoopToSCFForWithCollapse) {
   forOp.getBody()->walk([&](scf::ForOp) { hasNestedFor = true; });
   EXPECT_FALSE(hasNestedFor);
 
+  // Ensure the collapsed loop has an attribute indicating the number
+  // of collapsed loops
+  auto collapseAttr =
+      forOp->getAttrOfType<IntegerAttr>(getCollapseCountAttrName());
+  ASSERT_TRUE(collapseAttr);
+  EXPECT_EQ(collapseAttr.getInt(), 2);
+
   // The collapsed loop should iterate over the product of dimensions
   // lb=0, step=1 (after collapsing two 0..10 inclusive loops)
   auto lbConst = getConstantIndex(forOp.getLowerBound());
@@ -302,6 +309,9 @@ TEST_F(OpenACCUtilsLoopTest, ConvertLoopToSCFForNoCollapse) {
   bool hasNestedFor = false;
   forOp.getBody()->walk([&](scf::ForOp) { hasNestedFor = true; });
   EXPECT_TRUE(hasNestedFor);
+
+  // No collapse happened, so no collapse_count attribute is expected
+  EXPECT_FALSE(forOp->hasAttr(getCollapseCountAttrName()));
 }
 
 TEST_F(OpenACCUtilsLoopTest, ConvertLoopToSCFForWithCollapseAndDynamicBounds) {