[mlir][acc] ACCComputeLowering needs to account for device_type par (#201267)

When assigning parallelism for compute constructs or loops, device_type
parallelism must be first considered as a group for all available (gang,
worker, vector) - if any of these have device_type setting, then those
are the only ones that should be considered. Only if the loop has no
device_type specific parallelism then default parallelism should be
assigned.

Author: razvanlupusoru

mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td

@@ -1829,6 +1829,10 @@ def OpenACC_ParallelOp
     mlir::Operation::operand_range
     getNumGangsValues(mlir::acc::DeviceType deviceType);
 
+    /// Return true if the op has any num_gangs, num_workers, or vector_length
+    /// clause for the given device_type.
+    bool hasAnyGangWorkerVector(mlir::acc::DeviceType deviceType);
+
     /// Return true if the op has the wait attribute for the
     /// mlir::acc::DeviceType::None device_type.
     bool hasWaitOnly();
@@ -2156,6 +2160,10 @@ def OpenACC_KernelsOp
     mlir::Operation::operand_range
     getNumGangsValues(mlir::acc::DeviceType deviceType);
 
+    /// Return true if the op has any num_gangs, num_workers, or vector_length
+    /// clause for the given device_type.
+    bool hasAnyGangWorkerVector(mlir::acc::DeviceType deviceType);
+
     /// Return true if the op has the wait attribute for the
     /// mlir::acc::DeviceType::None device_type.
     bool hasWaitOnly();
@@ -2808,6 +2816,10 @@ def OpenACC_LoopOp
     // 'default'/None device-type.
     bool hasDefaultGangWorkerVector();
 
+    // Return whether this LoopOp has a gang, worker, or vector for the given
+    // device-type.
+    bool hasAnyGangWorkerVector(DeviceType deviceType);
+
     // Used to obtain the parallelism mode for the requested device type.
     // This first checks if the mode is set for the device_type requested.
     // And if not, it returns the non-device_type mode.

mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp

@@ -2173,6 +2173,31 @@ ParallelOp::getNumGangsValues(mlir::acc::DeviceType deviceType) {
                                getNumGangsSegments(), deviceType);
 }
 
+static bool hasAnyGangWorkerVectorForDeviceType(
+    std::optional<mlir::ArrayAttr> numGangsDeviceType,
+    mlir::Operation::operand_range numGangs,
+    std::optional<llvm::ArrayRef<int32_t>> numGangsSegments,
+    std::optional<mlir::ArrayAttr> numWorkersDeviceType,
+    mlir::Operation::operand_range numWorkers,
+    std::optional<mlir::ArrayAttr> vectorLengthDeviceType,
+    mlir::Operation::operand_range vectorLength,
+    mlir::acc::DeviceType deviceType) {
+  return !getValuesFromSegments(numGangsDeviceType, numGangs, numGangsSegments,
+                                deviceType)
+              .empty() ||
+         getValueInDeviceTypeSegment(numWorkersDeviceType, numWorkers,
+                                     deviceType) ||
+         getValueInDeviceTypeSegment(vectorLengthDeviceType, vectorLength,
+                                     deviceType);
+}
+
+bool acc::ParallelOp::hasAnyGangWorkerVector(mlir::acc::DeviceType deviceType) {
+  return hasAnyGangWorkerVectorForDeviceType(
+      getNumGangsDeviceType(), getNumGangs(), getNumGangsSegments(),
+      getNumWorkersDeviceType(), getNumWorkers(), getVectorLengthDeviceType(),
+      getVectorLength(), deviceType);
+}
+
 bool acc::ParallelOp::hasWaitOnly() {
   return hasWaitOnly(mlir::acc::DeviceType::None);
 }
@@ -3039,6 +3064,13 @@ KernelsOp::getNumGangsValues(mlir::acc::DeviceType deviceType) {
                                getNumGangsSegments(), deviceType);
 }
 
+bool acc::KernelsOp::hasAnyGangWorkerVector(mlir::acc::DeviceType deviceType) {
+  return hasAnyGangWorkerVectorForDeviceType(
+      getNumGangsDeviceType(), getNumGangs(), getNumGangsSegments(),
+      getNumWorkersDeviceType(), getNumWorkers(), getVectorLengthDeviceType(),
+      getVectorLength(), deviceType);
+}
+
 bool acc::KernelsOp::hasWaitOnly() {
   return hasWaitOnly(mlir::acc::DeviceType::None);
 }
@@ -3968,9 +4000,15 @@ bool acc::LoopOp::hasParallelismFlag(DeviceType dt) {
 }
 
 bool acc::LoopOp::hasDefaultGangWorkerVector() {
-  return hasVector() || getVectorValue() || hasWorker() || getWorkerValue() ||
-         hasGang() || getGangValue(GangArgType::Num) ||
-         getGangValue(GangArgType::Dim) || getGangValue(GangArgType::Static);
+  return hasAnyGangWorkerVector(DeviceType::None);
+}
+
+bool acc::LoopOp::hasAnyGangWorkerVector(DeviceType deviceType) {
+  return hasVector(deviceType) || getVectorValue(deviceType) ||
+         hasWorker(deviceType) || getWorkerValue(deviceType) ||
+         hasGang(deviceType) || getGangValue(GangArgType::Num, deviceType) ||
+         getGangValue(GangArgType::Dim, deviceType) ||
+         getGangValue(GangArgType::Static, deviceType);
 }
 
 acc::LoopParMode

mlir/lib/Dialect/OpenACC/Transforms/ACCComputeLowering.cpp

@@ -185,11 +185,32 @@ static void insertParDim(SmallVectorImpl<GPUParallelDimAttr> &parDims,
     parDims.insert(lb, parDim);
 }
 
+/// Return the device type from which gang/worker/vector clauses should be read.
+/// If the requested device type has any such clauses, use that exclusively;
+/// otherwise fall back to the default (DeviceType::None).
+static DeviceType getGangWorkerVectorDeviceType(LoopOp loopOp,
+                                                DeviceType deviceType) {
+  if (deviceType != DeviceType::None &&
+      loopOp.hasAnyGangWorkerVector(deviceType))
+    return deviceType;
+  return DeviceType::None;
+}
+
+template <typename ComputeConstructT>
+static DeviceType getParDimsDeviceType(ComputeConstructT computeOp,
+                                       DeviceType deviceType) {
+  if (deviceType != DeviceType::None &&
+      computeOp.hasAnyGangWorkerVector(deviceType))
+    return deviceType;
+  return DeviceType::None;
+}
+
 /// Map loop parallelism clauses (gang/worker/vector) to GPU parallel
 /// dimensions using the given mapping policy.
 static SmallVector<GPUParallelDimAttr>
 getParallelDimensions(LoopOp loopOp, const ACCToGPUMappingPolicy &policy,
                       DeviceType deviceType) {
+  deviceType = getGangWorkerVectorDeviceType(loopOp, deviceType);
   SmallVector<GPUParallelDimAttr> parDims;
   auto *ctx = loopOp->getContext();
 
@@ -229,12 +250,12 @@ assignKnownLaunchArgs(ComputeConstructT computeOp, DeviceType deviceType,
     if (isEffectivelySerial(computeOp))
       return {ParWidthOp::create(rewriter, loc, Value(), policy.seqDim(ctx))};
 
+    deviceType = getParDimsDeviceType(computeOp, deviceType);
+
     SmallVector<Value> values;
     auto indexTy = rewriter.getIndexType();
 
     auto numGangs = computeOp.getNumGangsValues(deviceType);
-    if (numGangs.empty())
-      numGangs = computeOp.getNumGangsValues();
     for (auto [gangDimIdx, gangSize] : llvm::enumerate(numGangs)) {
       auto gangLevel = getGangParLevel(gangDimIdx + 1);
       values.push_back(ParWidthOp::create(
@@ -245,8 +266,6 @@ assignKnownLaunchArgs(ComputeConstructT computeOp, DeviceType deviceType,
     }
 
     Value numWorkers = computeOp.getNumWorkersValue(deviceType);
-    if (!numWorkers)
-      numWorkers = computeOp.getNumWorkersValue();
     if (numWorkers) {
       values.push_back(ParWidthOp::create(
           rewriter, loc,
@@ -256,8 +275,6 @@ assignKnownLaunchArgs(ComputeConstructT computeOp, DeviceType deviceType,
     }
 
     Value vectorLength = computeOp.getVectorLengthValue(deviceType);
-    if (!vectorLength)
-      vectorLength = computeOp.getVectorLengthValue();
     if (vectorLength) {
       values.push_back(ParWidthOp::create(
           rewriter, loc,

mlir/test/Dialect/OpenACC/acc-compute-lowering-compute-device-type.mlir

@@ -0,0 +1,25 @@
+// RUN: mlir-opt %s -acc-compute-lowering=device-type=nvidia | FileCheck %s
+
+// Default num_gangs, nvidia vector_length: with device-type=nvidia only vector applies.
+// CHECK-LABEL: func.func @parallel_default_gangs_nvidia_vector_length
+func.func @parallel_default_gangs_nvidia_vector_length(%buf: memref<4xi32>) {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  %c4_i32 = arith.constant 4 : i32
+  %c32_i32 = arith.constant 32 : i32
+
+  %dev = acc.copyin varPtr(%buf : memref<4xi32>) -> memref<4xi32>
+  // CHECK-NOT: acc.par_width {{.*}} {par_dim = #acc.par_dim<block_x>}
+  // CHECK: acc.par_width {{.*}} {par_dim = #acc.par_dim<thread_x>}
+  acc.parallel num_gangs({%c4_i32 : i32}) vector_length(%c32_i32 : i32 [#acc.device_type<nvidia>]) dataOperands(%dev : memref<4xi32>) {
+    acc.loop 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
+    } attributes {independent = [#acc.device_type<none>]}
+    acc.yield
+  }
+  acc.copyout accPtr(%dev : memref<4xi32>) to varPtr(%buf : memref<4xi32>)
+  return
+}

mlir/test/Dialect/OpenACC/acc-compute-lowering-loop-device-type.mlir

@@ -0,0 +1,23 @@
+// RUN: mlir-opt %s -acc-compute-lowering=device-type=nvidia | FileCheck %s
+
+// Gang on default, vector on nvidia: with device-type=nvidia only vector applies.
+// CHECK-LABEL: func.func @parallel_loop_gang_default_vector_nvidia
+func.func @parallel_loop_gang_default_vector_nvidia(%buf: memref<1xi32>) {
+  %c0 = arith.constant 0 : index
+  %c1_i32 = arith.constant 1 : i32
+  %c10_i32 = arith.constant 10 : i32
+  %c100_i32 = arith.constant 100 : i32
+
+  %dev = acc.copyin varPtr(%buf : memref<1xi32>) -> memref<1xi32>
+  // CHECK-NOT: acc.par_dims = #acc<par_dims[block_x]>
+  // CHECK: acc.par_dims = #acc<par_dims[thread_x]>
+  acc.parallel num_gangs({%c10_i32 : i32}) dataOperands(%dev : memref<1xi32>) {
+    acc.loop gang control(%arg0 : i32) = (%c1_i32 : i32) to (%c100_i32 : i32) step (%c1_i32 : i32) {
+      memref.store %arg0, %dev[%c0] : memref<1xi32>
+      acc.yield
+    } attributes {auto_ = [#acc.device_type<none>], gang = [#acc.device_type<none>], vector = [#acc.device_type<nvidia>]}
+    acc.yield
+  }
+  acc.copyout accPtr(%dev : memref<1xi32>) to varPtr(%buf : memref<1xi32>)
+  return
+}