[Mlir-commits] [mlir] [OpenACC] set CollapseCountAttr for scf.parallel (PR #211676)

[Scott Manley]

https://github.com/rscottmanley created https://github.com/llvm/llvm-project/pull/211676

During compute lowering, also set the collapse count attribute when creating scf.parallel ops from acc.loop ops. While this collapse count is implied by the number of IVs, later canonicalizations may use this information to decide whether or not to collapse the loop nest further

>From 57131a55c852fea238e768c13ee0596b35e3e336 Mon Sep 17 00:00:00 2001
From: Scott Manley <scmanley at nvidia.com>
Date: Thu, 23 Jul 2026 14:48:32 -0700
Subject: [PATCH] [OpenACC] set CollapseCountAttr for scf.parallel

During compute lowering, also set the collapse count attribute when
creating scf.parallel ops from acc.loop ops. While this collapse count
is implied by the number of IVs, later canonicalizations may use this
information to decide whether or not to collapse the loop nest further
---
 .../OpenACC/Utils/OpenACCUtilsLoop.cpp        |  1 +
 .../OpenACC/acc-compute-lowering-loop.mlir    | 28 +++++++++++++++++++
 2 files changed, 29 insertions(+)

diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
index fd46a2800e03d..52f97e0867ee6 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
@@ -320,6 +320,7 @@ scf::ParallelOp convertACCLoopToSCFParallel(LoopOp loopOp,
       normalizeIVUses(rewriter, loc, iv, loopOp.getLowerbound()[idx],
                       loopOp.getStep()[idx]);
 
+  setCollapseCountAttr(parallelOp, parallelOp.getNumLoops());
   return parallelOp;
 }
 
diff --git a/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir b/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
index 0207157383c49..0772e6a2a536d 100644
--- a/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
@@ -83,6 +83,34 @@ func.func @parallel_loop_auto_collapse(%buf: memref<1xi32>, %lb0 : index, %ub0 :
 
 // -----
 
+// CHECK-LABEL: func.func @parallel_loop_collapse
+func.func @parallel_loop_collapse(%buf: memref<1xi32>, %lb0 : index, %ub0 : index, %lb1 : index, %ub1 : index) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+
+  %dev = acc.copyin varPtr(%buf : memref<1xi32>) -> memref<1xi32>
+  // CHECK-NOT: acc.parallel
+  // CHECK: acc.kernel_environment
+  // CHECK-NOT: acc.par_width
+  // CHECK: acc.compute_region
+  // CHECK: scf.parallel (%{{.*}}, %{{.*}})
+  // 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
+      memref.store %vi, %dev[%c0] : memref<1xi32>
+      acc.yield
+    } attributes {independent = [#acc.device_type<none>]}
+    acc.yield
+  }
+  acc.copyout accPtr(%dev : memref<1xi32>) to varPtr(%buf : memref<1xi32>)
+  return
+}
+
+// -----
+
 // CHECK-LABEL: func.func @serial_loop_normalized
 func.func @serial_loop_normalized(%buf: memref<1xi32>) {
   %c0 = arith.constant 0 : index