[Mlir-commits] [mlir] [mlir][OpenACC] Ignore non-constant worker and vector sizes on kernels loops (PR #228494)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-openacc

Author: Delaram Talaashrafi (delaram-talaashrafi)

<details>
<summary>Changes</summary>

A non-constant worker and vector clause on `acc.loop` inside acc.kernels is skipped, instead of failing as not yet implemented. The pass reports the skip.

---
Full diff: https://github.com/llvm/llvm-project/pull/228494.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp (+10-1) 
- (modified) mlir/test/Dialect/OpenACC/acc-compute-lowering-compute.mlir (+62) 


``````````diff
diff --git a/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp b/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
index 65d69f447cb2c8..84867cf58d5e62 100644
--- a/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
+++ b/mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp
@@ -162,7 +162,8 @@ struct SizedLevel {
 };
 using SizedLevelMap = DenseMap<Operation *, SmallVector<SizedLevel>>;
 
-/// Record a sized clause if `size` is a constant; NYI otherwise.
+/// Record a sized clause if `size` is a constant. A non-constant vector or
+/// worker clause is ignored. Other non-constant sizes are NYI.
 static LogicalResult tryAddSizedLevel(SizedLevelMap &sizedLevelMap,
                                       Operation *computeOp, LoopOp loopOp,
                                       ParLevel level, Value size,
@@ -171,6 +172,14 @@ static LogicalResult tryAddSizedLevel(SizedLevelMap &sizedLevelMap,
     return success();
   std::optional<int64_t> constSize = getConstantIntValue(size);
   if (!constSize) {
+    if (level == ParLevel::vector || level == ParLevel::worker) {
+      accSupport.emitRemark(loopOp,
+                            level == ParLevel::vector
+                                ? "ignoring non-constant vector clause"
+                                : "ignoring non-constant worker clause",
+                            DEBUG_TYPE);
+      return success();
+    }
     accSupport.emitNYI(loopOp.getLoc(),
                        "non-constant sized parallelism clause");
     return failure();
diff --git a/mlir/test/Dialect/OpenACC/acc-compute-lowering-compute.mlir b/mlir/test/Dialect/OpenACC/acc-compute-lowering-compute.mlir
index f8719325874950..82e87179347a1d 100644
--- a/mlir/test/Dialect/OpenACC/acc-compute-lowering-compute.mlir
+++ b/mlir/test/Dialect/OpenACC/acc-compute-lowering-compute.mlir
@@ -1,4 +1,5 @@
 // RUN: mlir-opt %s -acc-compute-lowering | FileCheck %s
+// RUN: mlir-opt %s -acc-compute-lowering --remarks-filter="(open)?acc.*" 2>&1 | FileCheck %s --check-prefix=REMARK
 
 // CHECK-LABEL: func.func @parallel_gang_loop
 func.func @parallel_gang_loop(%buf: memref<1xi32>) {
@@ -301,3 +302,64 @@ func.func @kernels_loop_sized_vector(%buf: memref<4xi32>) {
   acc.copyout accPtr(%dev : memref<4xi32>) to varPtr(%buf : memref<4xi32>)
   return
 }
+
+// -----
+
+// A sized `worker(n)` clause on a loop inside acc.kernels supplies the worker
+// launch width when the construct itself has no num_workers clause.
+
+// CHECK-LABEL: func.func @kernels_loop_sized_worker
+func.func @kernels_loop_sized_worker(%buf: memref<4xi32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  %c8_i32 = arith.constant 8 : i32
+
+  %dev = acc.copyin varPtr(%buf : memref<4xi32>) -> memref<4xi32>
+  // CHECK: %[[NW:.*]] = arith.constant 8 : index
+  // CHECK: acc.kernel_environment
+  // CHECK: %[[PW:.*]] = acc.par_width %[[NW]] par_dim(#acc.par_dim<thread_y>)
+  // CHECK: acc.compute_region launch(%{{.*}} = %[[PW]])
+  // CHECK: scf.parallel
+  // CHECK: acc.par_dims = #acc<par_dims[thread_y]>
+  acc.kernels dataOperands(%dev : memref<4xi32>) {
+    acc.loop worker(%c8_i32 : i32) control(%i : index) = (%c0 : index) to (%c4 : index) step (%c1 : index) {
+      %vi = arith.index_cast %i : index to i32
+      memref.store %vi, %dev[%i] : memref<4xi32>
+      acc.yield
+    } independent
+    acc.terminator
+  }
+  acc.copyout accPtr(%dev : memref<4xi32>) to varPtr(%buf : memref<4xi32>)
+  return
+}
+
+// -----
+
+// A non-constant worker(n) on a loop inside acc.kernels is ignored. The loop
+// stays a worker loop, and the launch does not take a worker width from it.
+
+// REMARK: remark:{{.*}}Category:acc-compute-lowering{{.*}}Function=kernels_loop_nonconst_worker{{.*}}Remark="ignoring non-constant worker clause"
+// CHECK-LABEL: func.func @kernels_loop_nonconst_worker
+func.func @kernels_loop_nonconst_worker(%buf: memref<4xi32>, %workers: i32) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+
+  %dev = acc.copyin varPtr(%buf : memref<4xi32>) -> memref<4xi32>
+  // CHECK: acc.kernel_environment
+  // CHECK-NOT: par_dim(#acc.par_dim<thread_y>)
+  // CHECK: acc.compute_region
+  // CHECK: scf.parallel
+  // CHECK: acc.par_dims = #acc<par_dims[thread_y]>
+  acc.kernels dataOperands(%dev : memref<4xi32>) {
+    acc.loop worker(%workers : i32) control(%i : index) = (%c0 : index) to (%c4 : index) step (%c1 : index) {
+      %vi = arith.index_cast %i : index to i32
+      memref.store %vi, %dev[%i] : memref<4xi32>
+      acc.yield
+    } independent
+    acc.terminator
+  }
+  acc.copyout accPtr(%dev : memref<4xi32>) to varPtr(%buf : memref<4xi32>)
+  return
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/228494