[Mlir-commits] [mlir] [OpenACC] propagate llvm loop annotations (PR #217414)

[Scott Manley]

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

When lowering acc.loop to scf.for or scf.parallel, copy any llvm loop annotations that are present

>From c1f7ceb9c8ee138cef5e2b9d41c4c4cf83725c9a Mon Sep 17 00:00:00 2001
From: Scott Manley <scmanley at nvidia.com>
Date: Wed, 19 Aug 2026 10:56:21 -0700
Subject: [PATCH] [OpenACC] propagate llvm loop annotations

When lowering acc.loop to scf.for or scf.parallel, copy any llvm loop
annotations that are present
---
 .../OpenACC/Utils/OpenACCUtilsLoop.cpp        | 11 +++++++++
 .../OpenACC/acc-compute-lowering-loop.mlir    | 24 +++++++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
index 52f97e0867ee6..963fbe29de6fc 100644
--- a/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
+++ b/mlir/lib/Dialect/OpenACC/Utils/OpenACCUtilsLoop.cpp
@@ -14,6 +14,7 @@
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/OpenACC/OpenACC.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SCF/Utils/Utils.h"
@@ -97,6 +98,14 @@ static Block::iterator cloneACCRegionIntoForLoop(Region *src, Block *dest,
   return ip;
 }
 
+/// Copy the discardable LLVM loop annotation attribute from an acc.loop to the
+/// lowered SCF op so later SCF to CFG/LLVM lowering can emit !llvm.loop
+/// metadata.
+static void copyLoopAnnotationAttr(Operation *from, Operation *to) {
+  if (Attribute ann = from->getDiscardableAttr(LLVM::LoopAnnotationAttr::name))
+    to->setDiscardableAttr(LLVM::LoopAnnotationAttr::name, ann);
+}
+
 } // namespace
 
 namespace mlir {
@@ -257,6 +266,7 @@ scf::ForOp convertACCLoopToSCFFor(LoopOp loopOp, RewriterBase &rewriter,
       setCollapseCountAttr(forOps.front(), numCollapsed);
   }
 
+  copyLoopAnnotationAttr(loopOp, forOps.front());
   return forOps.front();
 }
 
@@ -321,6 +331,7 @@ scf::ParallelOp convertACCLoopToSCFParallel(LoopOp loopOp,
                       loopOp.getStep()[idx]);
 
   setCollapseCountAttr(parallelOp, parallelOp.getNumLoops());
+  copyLoopAnnotationAttr(loopOp, parallelOp);
   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 0772e6a2a536d..2258f6bce49e5 100644
--- a/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-compute-lowering-loop.mlir
@@ -1,5 +1,8 @@
 // RUN: mlir-opt %s -acc-compute-lowering | FileCheck %s
 
+// CHECK: [[UNROLL:#.*]] = #llvm.loop_unroll<disable = false, full = true>
+// CHECK: #loop_annotation = #llvm.loop_annotation<unroll = [[UNROLL]]>
+
 // CHECK-LABEL: func.func @parallel_independent_loop
 func.func @parallel_independent_loop(%buf: memref<16xi32>) {
   %c0 = arith.constant 0 : index
@@ -232,3 +235,24 @@ func.func @parallel_loop_auto_gang(%buf: memref<1xi32>) {
   acc.copyout accPtr(%dev : memref<1xi32>) to varPtr(%buf : memref<1xi32>)
   return
 }
+
+// -----
+
+// Preserve llvm.loop_annotation (e.g. from !dir$ unroll) when lowering acc.loop.
+// CHECK-LABEL: func.func @orphan_loop_unroll_annotation
+// CHECK-NOT: acc.loop
+// CHECK: scf.for {{.*}} {
+// CHECK: } {llvm.loop_annotation = #loop_annotation}
+func.func @orphan_loop_unroll_annotation(%buf: memref<8xi32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c8 = arith.constant 8 : index
+  %c0_i32 = arith.constant 0 : i32
+
+  acc.loop control(%i : index) = (%c0 : index) to (%c8 : index) step (%c1 : index) {
+    memref.store %c0_i32, %buf[%i] : memref<8xi32>
+    acc.yield
+  } attributes {independent = [#acc.device_type<none>],
+                llvm.loop_annotation = #llvm.loop_annotation<unroll = <disable = false, full = true>>} 
+  return
+}