[mlir][acc] Improve acc.loop support as a container (#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.

Author: razvanlupusoru

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();

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,49 @@ 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.
+  if (isContainerLike()) {
+    // Obtain the maximum collapse count - we use this to check that there
+    // are enough loops contained.
+    uint64_t collapseCount = getCollapseValue().value_or(1);
+    if (getCollapseAttr()) {
+      for (auto collapseEntry : getCollapseAttr()) {
+        auto intAttr = mlir::dyn_cast<IntegerAttr>(collapseEntry);
+        if (intAttr.getValue().getZExtValue() > collapseCount)
+          collapseCount = intAttr.getValue().getZExtValue();
+      }
+    }
+
+    // We want to check that we find enough loop-like operations inside.
+    // PreOrder walk allows us to walk in a breadth-first manner at each nesting
+    // level.
+    mlir::Operation *expectedParent = this->getOperation();
+    bool foundSibling = false;
+    getRegion().walk<WalkOrder::PreOrder>([&](mlir::Operation *op) {
+      if (mlir::isa<mlir::LoopLikeOpInterface>(op)) {
+        // This effectively checks that we are not looking at a sibling loop.
+        if (op->getParentOfType<mlir::LoopLikeOpInterface>() !=
+            expectedParent) {
+          foundSibling = true;
+          return mlir::WalkResult::interrupt();
+        }
+
+        collapseCount--;
+        expectedParent = op;
+      }
+      // We found enough contained loops.
+      if (collapseCount == 0)
+        return mlir::WalkResult::interrupt();
+      return mlir::WalkResult::advance();
+    });
+
+    if (foundSibling)
+      return emitError("found sibling loops inside container-like acc.loop");
+    if (collapseCount != 0)
+      return emitError("failed to find enough loop-like operations inside "
+                       "container-like acc.loop");
+  }
+
   return success();
 }
 

mlir/test/Dialect/OpenACC/invalid.mlir

@@ -752,7 +752,6 @@ func.func @acc_combined() {
   // expected-error @below {{expected 'loop'}}
   acc.parallel combined() {
   }
-
   return
 }
 
@@ -762,7 +761,6 @@ func.func @acc_combined() {
   // expected-error @below {{expected compute construct name}}
   acc.loop combined(loop) {
   }
-
   return
 }
 
@@ -772,7 +770,6 @@ func.func @acc_combined() {
   // expected-error @below {{expected 'loop'}}
   acc.parallel combined(parallel loop) {
   }
-
   return
 }
 
@@ -782,6 +779,43 @@ func.func @acc_combined() {
   // expected-error @below {{expected ')'}}
   acc.loop combined(parallel loop) {
   }
+  return
+}
+
+// -----
 
+func.func @acc_loop_container() {
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+  // expected-error @below {{found sibling loops inside container-like acc.loop}}
+  acc.loop {
+    scf.for %arg4 = %c0 to %c10 step %c1 {
+      scf.yield
+    }
+    scf.for %arg5 = %c0 to %c10 step %c1 {
+        scf.yield
+    }
+    acc.yield
+  } attributes { collapse = [2], collapseDeviceType = [#acc.device_type<none>] }
+  return
+}
+
+// -----
+
+func.func @acc_loop_container() {
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+  // expected-error @below {{failed to find enough loop-like operations inside container-like acc.loop}}
+  acc.loop {
+    scf.for %arg4 = %c0 to %c10 step %c1 {
+      scf.for %arg5 = %c0 to %c10 step %c1 {
+          scf.yield
+      }
+      scf.yield
+    }
+    acc.yield
+  } attributes { collapse = [3], collapseDeviceType = [#acc.device_type<none>] }
   return
 }

mlir/test/Dialect/OpenACC/ops.mlir

@@ -1862,22 +1862,26 @@ func.func @acc_num_gangs() {
 
 // CHECK-LABEL: func.func @acc_combined
 func.func @acc_combined() {
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+
   acc.parallel combined(loop) {
-    acc.loop combined(parallel) {
+    acc.loop combined(parallel) control(%arg3 : index) = (%c0 : index) to (%c10 : index) step (%c1 : index) {
       acc.yield
     }
     acc.terminator
   }
 
   acc.kernels combined(loop) {
-    acc.loop combined(kernels) {
+    acc.loop combined(kernels) control(%arg3 : index) = (%c0 : index) to (%c10 : index) step (%c1 : index) {
       acc.yield
     }
     acc.terminator
   }
 
   acc.serial combined(loop) {
-    acc.loop combined(serial) {
+    acc.loop combined(serial) control(%arg3 : index) = (%c0 : index) to (%c10 : index) step (%c1 : index) {
       acc.yield
     }
     acc.terminator
@@ -1933,3 +1937,45 @@ acc.private.recipe @privatization_memref_i32 : memref<i32> init {
 
 // CHECK-LABEL: acc.private.recipe @privatization_memref_i32
 // CHECK:       memref.alloca
+
+// -----
+
+func.func @acc_loop_container() {
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+  acc.loop {
+    scf.for %arg4 = %c0 to %c10 step %c1 {
+      scf.yield
+    }
+    acc.yield
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @acc_loop_container
+// CHECK:       acc.loop
+// CHECK:       scf.for
+
+// -----
+
+func.func @acc_loop_container() {
+  %c0 = arith.constant 0 : index
+  %c10 = arith.constant 10 : index
+  %c1 = arith.constant 1 : index
+  acc.loop {
+    scf.for %arg4 = %c0 to %c10 step %c1 {
+      scf.for %arg5 = %c0 to %c10 step %c1 {
+        scf.yield
+      }
+      scf.yield
+    }
+    acc.yield
+  } attributes { collapse = [2], collapseDeviceType = [#acc.device_type<none>] }
+  return
+}
+
+// CHECK-LABEL: func.func @acc_loop_container
+// CHECK:       acc.loop
+// CHECK:       scf.for
+// CHECK:       scf.for