[mlir][openacc] Switch numGangs to a variadic operand

[clementval]

Summary

In the latest spec, the num_gangs clause accepts up to three
arguments. Update the dialect to swicth numGangs operands from
optional single operand to a variadic operand. The verifier limits
the number of operands to three as specified in the spec.

[razvanlupusoru]

This looks great! Thank you!

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

let description = [{
     }
     ```
   }];
 
   let arguments = (ins Optional<IntOrIndex>:$async,
                        UnitAttr:$asyncAttr,
                        Variadic<IntOrIndex>:$waitOperands,
                        UnitAttr:$waitAttr,
-                       Optional<IntOrIndex>:$numGangs,
+                       Variadic<IntOrIndex>:$numGangs,
                        Optional<IntOrIndex>:$numWorkers,
                        Optional<IntOrIndex>:$vectorLength,
                        Optional<I1>:$ifCond,
                        Optional<I1>:$selfCond,
                        UnitAttr:$selfAttr,
                        Variadic<AnyType>:$reductionOperands,
                        OptionalAttr<SymbolRefArrayAttr>:$reductionRecipes,
                        Variadic<OpenACC_PointerLikeTypeInterface>:$gangPrivateOperands,
let description = [{
     }
     ```
   }];
 
   let arguments = (ins Optional<IntOrIndex>:$async,
                        UnitAttr:$asyncAttr,
                        Variadic<IntOrIndex>:$waitOperands,
                        UnitAttr:$waitAttr,
-                       Optional<IntOrIndex>:$numGangs,
+                       Variadic<IntOrIndex>:$numGangs,
                        Optional<IntOrIndex>:$numWorkers,
                        Optional<IntOrIndex>:$vectorLength,
                        Optional<I1>:$ifCond,
                        Optional<I1>:$selfCond,
                        UnitAttr:$selfAttr,
                        Variadic<OpenACC_PointerLikeTypeInterface>:$dataClauseOperands,
                        OptionalAttr<DefaultValueAttr>:$defaultAttr);
 

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

 
 unsigned ParallelOp::getNumDataOperands() {
   return getReductionOperands().size() + getGangPrivateOperands().size() +
          getGangFirstPrivateOperands().size() + getDataClauseOperands().size();
 }
 
 Value ParallelOp::getDataOperand(unsigned i) {
   unsigned numOptional = getAsync() ? 1 : 0;
-  numOptional += getNumGangs() ? 1 : 0;
+  numOptional += getNumGangs().size();
   numOptional += getNumWorkers() ? 1 : 0;
   numOptional += getVectorLength() ? 1 : 0;
   numOptional += getIfCond() ? 1 : 0;
   numOptional += getSelfCond() ? 1 : 0;
   return getOperand(getWaitOperands().size() + numOptional + i);
 }
 
 LogicalResult acc::ParallelOp::verify() {
   if (failed(checkSymOperandList<mlir::acc::PrivateRecipeOp>(
           *this, getPrivatizations(), getGangPrivateOperands(), "private",
           "privatizations")))
     return failure();
   if (failed(checkSymOperandList<mlir::acc::ReductionRecipeOp>(
           *this, getReductionRecipes(), getReductionOperands(), "reduction",
           "reductions", false)))
     return failure();
+  if (getNumGangs().size() > 3)
+    return emitOpError() << "num_gangs expects a maximum of 3 values";
   return checkDataOperands<acc::ParallelOp>(*this, getDataClauseOperands());
 }
 
 //===----------------------------------------------------------------------===//
 // SerialOp
 //===----------------------------------------------------------------------===//
 
 unsigned SerialOp::getNumDataOperands() {

 //===----------------------------------------------------------------------===//
 
 unsigned KernelsOp::getNumDataOperands() {
   return getDataClauseOperands().size();
 }
 
 Value KernelsOp::getDataOperand(unsigned i) {
   unsigned numOptional = getAsync() ? 1 : 0;
+  numOptional += getWaitOperands().size();
+  numOptional += getNumGangs().size();
+  numOptional += getNumWorkers() ? 1 : 0;
+  numOptional += getVectorLength() ? 1 : 0;
   numOptional += getIfCond() ? 1 : 0;
   numOptional += getSelfCond() ? 1 : 0;
-  return getOperand(getWaitOperands().size() + numOptional + i);
+  return getOperand(numOptional + i);
 }
 
 LogicalResult acc::KernelsOp::verify() {
+  if (getNumGangs().size() > 3)
+    return emitOpError() << "num_gangs expects a maximum of 3 values";
   return checkDataOperands<acc::KernelsOp>(*this, getDataClauseOperands());
 }
 
 //===----------------------------------------------------------------------===//
 // HostDataOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult acc::HostDataOp::verify() {

mlir/test/Dialect/OpenACC/invalid.mlir

 
 // -----
 
 // expected-error@+1 {{expect at least one of num, dim or static values}}
 acc.loop gang() {
   "test.openacc_dummy_op"() : () -> ()
   acc.yield
 }
+
+// -----
+
+%i64value = arith.constant 1 : i64
+// expected-error@+1 {{num_gangs expects a maximum of 3 values}}
+acc.parallel num_gangs(%i64value, %i64value, %i64value, %i64value : i64, i64, i64, i64) {
+}

mlir/test/Dialect/OpenACC/ops.mlir

func.func @testparallelop(%a: memref<10xf32>, %b: memref<10xf32>, %c: memref<10x10xf32>) -> () {
   acc.parallel wait(%i64value, %i32value, %idxValue : i64, i32, index) {
   }
   acc.parallel num_gangs(%i64value: i64) {
   }
   acc.parallel num_gangs(%i32value: i32) {
   }
   acc.parallel num_gangs(%idxValue: index) {
   }
+  acc.parallel num_gangs(%i64value, %i64value, %idxValue : i64, i64, index) {
+  }
   acc.parallel num_workers(%i64value: i64) {
   }
   acc.parallel num_workers(%i32value: i32) {
   }
   acc.parallel num_workers(%idxValue: index) {
   }
   acc.parallel vector_length(%i64value: i64) {
   }

 // CHECK:      acc.parallel wait([[I64VALUE]], [[I32VALUE]], [[IDXVALUE]] : i64, i32, index) {
 // CHECK-NEXT: }
 // CHECK:      acc.parallel num_gangs([[I64VALUE]] : i64) {
 // CHECK-NEXT: }
 // CHECK:      acc.parallel num_gangs([[I32VALUE]] : i32) {
 // CHECK-NEXT: }
 // CHECK:      acc.parallel num_gangs([[IDXVALUE]] : index) {
 // CHECK-NEXT: }
+// CHECK:      acc.parallel num_gangs([[I64VALUE]], [[I64VALUE]], [[IDXVALUE]] : i64, i64, index) {
+// CHECK-NEXT: }
 // CHECK:      acc.parallel num_workers([[I64VALUE]] : i64) {
 // CHECK-NEXT: }
 // CHECK:      acc.parallel num_workers([[I32VALUE]] : i32) {
 // CHECK-NEXT: }
 // CHECK:      acc.parallel num_workers([[IDXVALUE]] : index) {
 // CHECK-NEXT: }
 // CHECK:      acc.parallel vector_length([[I64VALUE]] : i64) {
 // CHECK-NEXT: }