[Mlir-commits] [mlir] [mlir][acc] Improve acc.loop support as a container (PR #137887)

[Razvan Lupusoru]

https://github.com/razvanlupusoru created https://github.com/llvm/llvm-project/pull/137887

Dialects which have their own loop representation not representable with numeric bounds + steps cannot be represented cleanly with acc.loop. In such a case, we permit the dialects representation with acc.loop merely encompasing its loop representation. This limitation became obvious when looking at range / random iterator C++ loops.

The API of acc.loop was updated to test for this differentiation. Additionally, the verifier was updated to check for consistent bounds and whether inner-loops are contained when it works as a container.

>From dac2419d76dd40f61479acee829d500d75f2945f Mon Sep 17 00:00:00 2001
From: Razvan Lupusoru <rlupusoru at nvidia.com>
Date: Tue, 29 Apr 2025 15:18:06 -0700
Subject: [PATCH] [mlir][acc] Improve acc.loop support as a container

Dialects which have their own loop representation not representable
with numeric bounds + steps cannot be represented cleanly with
acc.loop. In such a case, we permit the dialects representation with
acc.loop merely encompasing its loop representation. This limitation
became obvious when looking at range / random iterator C++ loops.

The API of acc.loop was updated to test for this differentiation.
Additionally, the verifier was updated to check for consistent bounds
and whether inner-loops are contained when it works as a container.
---
 mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td |  5 +++++
 mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp         | 17 +++++++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
index df7ff28ca5926..3ad8e4f9ccbeb 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
@@ -2127,6 +2127,11 @@ def OpenACC_LoopOp : OpenACC_Op<"loop",
     /// Used to retrieve the block inside the op's region.
     Block &getBody() { return getLoopRegions().front()->front(); }
 
+    /// Used to determine if this operation is merely a container for a loop
+    /// operation instead of being loop-like itself.
+    bool isLoopLike() { return !getLowerbound().empty(); }
+    bool isContainerLike() { return !isLoopLike(); }
+
     /// Return true if the op has the auto attribute for the
     /// mlir::acc::DeviceType::None device_type.
     bool hasAuto();
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
index 91025e90b8e76..a7b01abe34e03 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
@@ -2298,6 +2298,14 @@ LogicalResult checkDeviceTypes(mlir::ArrayAttr deviceTypes) {
 }
 
 LogicalResult acc::LoopOp::verify() {
+  if (getUpperbound().size() != getStep().size())
+    return emitError() << "number of upperbounds expected to be the same as "
+                          "number of steps";
+
+  if (getUpperbound().size() != getLowerbound().size())
+    return emitError() << "number of upperbounds expected to be the same as "
+                          "number of lowerbounds";
+
   if (!getUpperbound().empty() && getInclusiveUpperbound() &&
       (getUpperbound().size() != getInclusiveUpperbound()->size()))
     return emitError() << "inclusiveUpperbound size is expected to be the same"
@@ -2415,6 +2423,15 @@ LogicalResult acc::LoopOp::verify() {
   if (getRegion().empty())
     return emitError("expected non-empty body.");
 
+  // When it is container-like - it is expected to hold a loop-like operation.
+  // TODO: Get the collapse attribute into account.
+  if (isContainerLike()) {
+    // TODO: Ensure there is a single loop-like operation at any one level.
+    auto loopLikeOps = getRegion().getOps<LoopLikeOpInterface>();
+    if (loopLikeOps.empty())
+      return emitError("expected to hold a loop-like operation.");
+  }
+
   return success();
 }