[mlir][openacc] Add gang dim operand to acc.loop operation

[clementval]

Summary

OpenACC 3.3 introduces a dim argument on the gang clause. This patch
adds a new operand for it on the acc.loop and update the custom
gang clause parser/printer for it.

Depends on D151970

[Anastasia]

mlir/test/Dialect/OpenACC/ops.mlir

269

Since there are a number of restrictions on gang dim e.g. it can't appear together with num argument, should this be an invalid IR and some validations rules could catch and report this?

[clementval]

mlir/test/Dialect/OpenACC/ops.mlir

269

I’ll add check in the verifier.

[clementval]

address review comments

[clementval]

clang-format

[jeanPerier]

Looks good

[razvanlupusoru]

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

549

I am feeling conflicted about this check because I don't understand why spec specifically restricts this case. To me it feels like something acc dialect should allow but not clause semantic verifier.

[clementval]

Remove num/dim check

[clementval]

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

549

Remove here.

flang/lib/Lower/OpenACC.cpp

createLoopOp(Fortran::lower::AbstractConverter &converter,
              Fortran::semantics::SemanticsContext &semanticsContext,
              Fortran::lower::StatementContext &stmtCtx,
              const Fortran::parser::AccClauseList &accClauseList) {
   fir::FirOpBuilder &firOpBuilder = converter.getFirOpBuilder();
 
   mlir::Value workerNum;
   mlir::Value vectorNum;
   mlir::Value gangNum;
+  mlir::Value gangDim;
   mlir::Value gangStatic;
   llvm::SmallVector<mlir::Value, 2> tileOperands, privateOperands,
       reductionOperands;
   bool hasGang = false, hasVector = false, hasWorker = false;
 
   for (const Fortran::parser::AccClause &clause : accClauseList.v) {
     mlir::Location clauseLocation = converter.genLocation(clause.source);
     if (const auto *gangClause =
if (const auto *gangClause =
       TODO(clauseLocation, "OpenACC compute construct reduction lowering");
     }
   }
 
   // Prepare the operand segment size attribute and the operands value range.
   llvm::SmallVector<mlir::Value> operands;
   llvm::SmallVector<int32_t> operandSegments;
   addOperand(operands, operandSegments, gangNum);
+  addOperand(operands, operandSegments, gangDim);
   addOperand(operands, operandSegments, gangStatic);
   addOperand(operands, operandSegments, workerNum);
   addOperand(operands, operandSegments, vectorNum);
   addOperands(operands, operandSegments, tileOperands);
   addOperands(operands, operandSegments, privateOperands);
   addOperands(operands, operandSegments, reductionOperands);
 
   auto loopOp = createRegionOp<mlir::acc::LoopOp, mlir::acc::YieldOp>(

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

acc.loop gang vector {
       }
       acc.yield
     } attributes { collapse = 3 }
     ```
   }];
 
   let arguments = (ins OptionalAttr<I64Attr>:$collapse,
                        Optional<IntOrIndex>:$gangNum,
+                       Optional<IntOrIndex>:$gangDim,
                        Optional<IntOrIndex>:$gangStatic,
                        Optional<IntOrIndex>:$workerNum,
                        Optional<IntOrIndex>:$vectorLength,
                        UnitAttr:$seq,
                        UnitAttr:$independent,
                        UnitAttr:$auto_,
                        UnitAttr:$hasGang,
                        UnitAttr:$hasWorker,
                        UnitAttr:$hasVector,
                        Variadic<IntOrIndex>:$tileOperands,
                        Variadic<AnyType>:$privateOperands,
                        OptionalAttr<OpenACC_ReductionOpAttr>:$reductionOp,
                        Variadic<AnyType>:$reductionOperands);
 
   let results = (outs Variadic<AnyType>:$results);
 
   let regions = (region AnyRegion:$region);
 
   let extraClassDeclaration = [{
     static StringRef getAutoAttrStrName() { return "auto"; }
     static StringRef getGangNumKeyword() { return "num"; }
+    static StringRef getGangDimKeyword() { return "dim"; }
     static StringRef getGangStaticKeyword() { return "static"; }
   }];
 
   let hasCustomAssemblyFormat = 1;
   let assemblyFormat = [{
     oilist(
-        `gang` `` custom<GangClause>($gangNum, type($gangNum), $gangStatic, type($gangStatic), $hasGang)
+        `gang` `` custom<GangClause>($gangNum, type($gangNum), $gangDim, type($gangDim), $gangStatic, type($gangStatic), $hasGang)
       | `worker` `` custom<WorkerClause>($workerNum, type($workerNum), $hasWorker)
       | `vector` `` custom<VectorClause>($vectorLength, type($vectorLength), $hasVector)
       | `private` `(` $privateOperands `:` type($privateOperands) `)`
       | `tile` `(` $tileOperands `:` type($tileOperands) `)`
       | `reduction` `(` $reductionOperands `:` type($reductionOperands) `)`
     )
     $region
     ( `(` type($results)^ `)` )?

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

if (succeeded(parser.parseOptionalKeyword(keyword))) {
     if (parser.parseOperand(*value) || parser.parseColonType(valueType))
       return failure();
     needComa = true;
     newValue = true;
   }
   return success();
 }
 
-static ParseResult
-parseGangClause(OpAsmParser &parser,
-                std::optional<OpAsmParser::UnresolvedOperand> &gangNum,
-                Type &gangNumType,
-                std::optional<OpAsmParser::UnresolvedOperand> &gangStatic,
-                Type &gangStaticType, UnitAttr &hasGang) {
+static ParseResult parseGangClause(
+    OpAsmParser &parser, std::optional<OpAsmParser::UnresolvedOperand> &gangNum,
+    Type &gangNumType, std::optional<OpAsmParser::UnresolvedOperand> &gangDim,
+    Type &gangDimType,
   hasGang = UnitAttr::get(parser.getBuilder().getContext());
   gangNum = std::nullopt;
+  gangDim = std::nullopt;
   gangStatic = std::nullopt;
   bool needComa = false;
 
   // optional gang operands
   if (succeeded(parser.parseOptionalLParen())) {
     while (true) {
       bool newValue = false;
       bool needValue = false;
       if (needComa) {
         if (succeeded(parser.parseOptionalComma()))
           needValue = true; // expect a new value after comma.
         else
           break;
       }
 
       if (failed(parseGangValue(parser, LoopOp::getGangNumKeyword(), gangNum,
                                 gangNumType, needComa, newValue)))
         return failure();
+      if (failed(parseGangValue(parser, LoopOp::getGangDimKeyword(), gangDim,
+                                gangDimType, needComa, newValue)))
+        return failure();
       if (failed(parseGangValue(parser, LoopOp::getGangStaticKeyword(),
                                 gangStatic, gangStaticType, needComa,
                                 newValue)))
         return failure();
 
       if (!newValue && needValue) {
         parser.emitError(parser.getCurrentLocation(),
                          "new value expected after comma");
         return failure();
       }
 
       if (!newValue)
         break;
     }
 
-    if (!gangNum && !gangStatic) {
+    if (!gangNum && !gangDim && !gangStatic) {
       parser.emitError(parser.getCurrentLocation(),
-                       "expect num and/or static value(s)");
+                       "expect at least one of num, dim or static values");
       return failure();
     }
 
     if (failed(parser.parseRParen()))
       return failure();
   }
   return success();
 }
 
 void printGangClause(OpAsmPrinter &p, Operation *op, Value gangNum,
-                     Type gangNumType, Value gangStatic, Type gangStaticType,
-                     UnitAttr hasGang) {
-  if (gangNum || gangStatic) {
+                     Type gangNumType, Value gangDim, Type gangDimType,
+                     Value gangStatic, Type gangStaticType, UnitAttr hasGang) {
+  if (gangNum || gangStatic || gangDim) {
     p << "(";
     if (gangNum) {
       p << LoopOp::getGangNumKeyword() << "=" << gangNum << " : "
         << gangNumType;
+      if (gangStatic || gangDim)
+        p << ", ";
+    }
+    if (gangDim) {
+      p << LoopOp::getGangDimKeyword() << "=" << gangDim << " : "
+        << gangDimType;
       if (gangStatic)
         p << ", ";
     }
     if (gangStatic)
       p << LoopOp::getGangStaticKeyword() << "=" << gangStatic << " : "
         << gangStaticType;
     p << ")";
   }
LogicalResult acc::LoopOp::verify() {
   if (getSeq() && (getHasGang() || getHasWorker() || getHasVector()))
     return emitError("gang, worker or vector cannot appear with the seq attr");
 
   // Check non-empty body().
   if (getRegion().empty())
     return emitError("expected non-empty body.");
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // DataOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult acc::DataOp::verify() {
   // 2.6.5. Data Construct restriction
   // At least one copy, copyin, copyout, create, no_create, present, deviceptr,

mlir/test/Dialect/OpenACC/invalid.mlir

 // expected-error@+1 {{new value expected after comma}}
 acc.loop gang(static=%i64Value: i64, ) {
   "test.openacc_dummy_op"() : () -> ()
   acc.yield
 }
 
 // -----
 
-// expected-error@+1 {{expect num and/or static value(s)}}
+// expected-error@+1 {{expect at least one of num, dim or static values}}
 acc.loop gang() {
   "test.openacc_dummy_op"() : () -> ()
   acc.yield
 }

mlir/test/Dialect/OpenACC/ops.mlir

func.func @testloopop() -> () {
   acc.loop tile(%i32Value, %i32Value : i32, i32) {
     "test.openacc_dummy_op"() : () -> ()
     acc.yield
   }
   acc.loop gang(static=%i64Value: i64, num=%i64Value: i64) {
     "test.openacc_dummy_op"() : () -> ()
     acc.yield
   }
+  acc.loop gang(num=%i64Value: i64, dim=%i64Value : i64, static=%i64Value: i64) {
+    "test.openacc_dummy_op"() : () -> ()
+    acc.yield
+  }
   return
 }
 
 // CHECK:      [[I64VALUE:%.*]] = arith.constant 1 : i64
 // CHECK-NEXT: [[I32VALUE:%.*]] = arith.constant 128 : i32
 // CHECK-NEXT: [[IDXVALUE:%.*]] = arith.constant 8 : index
 // CHECK:      acc.loop gang worker vector {
 // CHECK-NEXT:   "test.openacc_dummy_op"() : () -> ()

 // CHECK:      acc.loop tile([[I32VALUE]], [[I32VALUE]] : i32, i32) {
 // CHECK-NEXT:   "test.openacc_dummy_op"() : () -> ()
 // CHECK-NEXT:   acc.yield
 // CHECK-NEXT: }
 // CHECK:      acc.loop gang(num=[[I64VALUE]] : i64, static=[[I64VALUE]] : i64) {
 // CHECK-NEXT:   "test.openacc_dummy_op"() : () -> ()
 // CHECK-NEXT:   acc.yield
 // CHECK-NEXT: }
+// CHECK:      acc.loop gang(num=[[I64VALUE]] : i64, dim=[[I64VALUE]] : i64, static=[[I64VALUE]] : i64) {
+// CHECK-NEXT:   "test.openacc_dummy_op"() : () -> ()
+// CHECK-NEXT:   acc.yield
+// CHECK-NEXT: }
 
 // -----
 
 func.func @acc_loop_multiple_block() {
   acc.parallel {
     acc.loop {
       %c1 = arith.constant 1 : index
       cf.br ^bb1(%c1 : index)