[mlir][openacc] Change operand type from index to AnyInteger in parallel op

[clementval]

Summary

This patch change the type of operands async, wait, numGangs, numWorkers and vectorLength from index
to AnyInteger to fit with acc.loop and the OpenACC specification.

[ftynse]

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

84–88

Make sure index is also supported. It was intended as a platform-specific opaque size type. While it's reasonable to adhere to the spec and allow any integer, most lowerings coming from higher levels and many standard ops (e.g. taking the size of some container) will use index.

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

156–157

This will now parse every type as i64, regardless of what type is actually used. Roundtripping will thus erase type information, which is incorrect. If you need to support different types, you need to specify which type is used somehow.

[clementval]

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

156–157

Right, I was suspecting smith was not quite right here. I guess I have to add the type and parse it as well like I did for the variadic operands in this op. I'll update the patch today.

[clementval]

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

84–88

I update this to use AnyTypeOf<[AnyInteger, Index]>

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

156–157

This has been change with parsing the actual type as well.

[clementval]

Address review comments

[clementval]

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

84–88

If that's fine I'll make a follow up patch to use the same type in acc.loop

[clementval]

Remove unrelated commits

[ftynse]

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

84–88

Yes, thanks!

You may want to define def IntOrIndex : AnyTypeOf<[...]>; and use it instead of repeating every time.

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

110

We have OptionalParseResult for this.

[clementval]

Address review comments

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

84–88

Sure makes sense. Patch is updated.

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

110

Thanks for the pointer. I was not aware of it. I updated this function and I'll update the reset in a follow up patch.

[clementval]

Rebase

flang/lib/Semantics/check-omp-structure.h

#include "llvm/Frontend/OpenMP/OMP.cpp.inc"
   void Leave(const parser::OpenMPSectionsConstruct &);
   void Enter(const parser::OmpEndSectionsDirective &);
 
   void Enter(const parser::OpenMPDeclareSimdConstruct &);
   void Leave(const parser::OpenMPDeclareSimdConstruct &);
   void Enter(const parser::OpenMPDeclareTargetConstruct &);
   void Leave(const parser::OpenMPDeclareTargetConstruct &);
 
+  void Enter(const parser::OpenMPAtomicConstruct &);
+  void Leave(const parser::OpenMPAtomicConstruct &);
   void Enter(const parser::OpenMPSimpleStandaloneConstruct &);
   void Leave(const parser::OpenMPSimpleStandaloneConstruct &);
   void Enter(const parser::OpenMPFlushConstruct &);
   void Leave(const parser::OpenMPFlushConstruct &);
   void Enter(const parser::OpenMPCancelConstruct &);
   void Leave(const parser::OpenMPCancelConstruct &);
   void Enter(const parser::OpenMPCancellationPointConstruct &);
   void Leave(const parser::OpenMPCancellationPointConstruct &);
#include "llvm/Frontend/OpenMP/OMP.cpp.inc"
   void Enter(const parser::OmpClause::Lastprivate &);
   void Enter(const parser::OmpClause::NumTasks &);
   void Enter(const parser::OmpClause::NumTeams &);
   void Enter(const parser::OmpClause::NumThreads &);
   void Enter(const parser::OmpClause::Ordered &);
   void Enter(const parser::OmpClause::Priority &);
   void Enter(const parser::OmpClause::Private &);
   void Enter(const parser::OmpClause::Safelen &);
+  void Enter(const parser::OmpClause::SeqCst &);
   void Enter(const parser::OmpClause::Shared &);
   void Enter(const parser::OmpClause::Simdlen &);
   void Enter(const parser::OmpClause::ThreadLimit &);
   void Enter(const parser::OmpClause::To &);
   void Enter(const parser::OmpClause::Link &);
   void Enter(const parser::OmpClause::Uniform &);
   void Enter(const parser::OmpClause::UseDevicePtr &);
   void Enter(const parser::OmpClause::IsDevicePtr &);

flang/lib/Semantics/check-omp-structure.cpp

case llvm::omp::Directive::OMPD_do_simd:
     SetClauseSets(dir.v);
     break;
   default:
     // no clauses are allowed
     break;
   }
 }
 
+void OmpStructureChecker::Enter(const parser::OpenMPAtomicConstruct &x) {
+  if (const auto &atomicRead{std::get<parser::OmpAtomicRead>(x.u)) {
+      const auto &dir{std::get<parser::Verbatim>(atomicRead.t)};
+      PushContextAndClauseSets(dir.source, llvm::omp::OMPD_atomic);
+    }
+}
+
+void OmpStructureChecker::Leave(const parser::OpenMPAtomicConstruct &) {
+    dirContext_.pop_back();
+}
+
 void OmpStructureChecker::Enter(const parser::OpenMPBlockConstruct &x) {
   const auto &beginBlockDir{std::get<parser::OmpBeginBlockDirective>(x.t)};
   const auto &endBlockDir{std::get<parser::OmpEndBlockDirective>(x.t)};
   const auto &beginDir{std::get<parser::OmpBlockDirective>(beginBlockDir.t)};
   const auto &endDir{std::get<parser::OmpBlockDirective>(endBlockDir.t)};
   CheckMatching<parser::OmpBlockDirective>(beginDir, endDir);
 
   PushContextAndClauseSets(beginDir.source, beginDir.v);

llvm/include/llvm/Frontend/OpenMP/OMP.td

let allowedClauses = [
     VersionedClause<OMPC_Depend>
   ];
 }
 def OMP_Atomic : Directive<"atomic"> {
   let allowedClauses = [
     VersionedClause<OMPC_Read>,
     VersionedClause<OMPC_Write>,
     VersionedClause<OMPC_Update>,
-    VersionedClause<OMPC_Capture>,
+    VersionedClause<OMPC_Capture>
+  ];
+  let allowedOnceClauses = [
     VersionedClause<OMPC_SeqCst>,
     VersionedClause<OMPC_AcqRel, 50>,
     VersionedClause<OMPC_Acquire, 50>,
     VersionedClause<OMPC_Release, 50>,
     VersionedClause<OMPC_Relaxed, 50>,
     VersionedClause<OMPC_Hint, 50>
   ];
 }

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

let description = [{
     ```mlir
     acc.parallel num_gangs(%c10) num_workers(%c10)
         private(%c : memref<10xf32>) {
       // parallel region
     }
     ```
   }];
 
-  let arguments = (ins Optional<Index>:$async,
-                       Variadic<Index>:$waitOperands,
-                       Optional<Index>:$numGangs,
-                       Optional<Index>:$numWorkers,
-                       Optional<Index>:$vectorLength,
+  let arguments = (ins Optional<AnyTypeOf<[AnyInteger, Index]>>:$async,
+                       Variadic<AnyTypeOf<[AnyInteger, Index]>>:$waitOperands,
+                       Optional<AnyTypeOf<[AnyInteger, Index]>>:$numGangs,
+                       Optional<AnyTypeOf<[AnyInteger, Index]>>:$numWorkers,
+                       Optional<AnyTypeOf<[AnyInteger, Index]>>:$vectorLength,
                        Optional<I1>:$ifCond,
                        Optional<I1>:$selfCond,
                        OptionalAttr<OpenACC_ReductionOpAttr>:$reductionOp,
                        Variadic<AnyType>:$reductionOperands,
                        Variadic<AnyType>:$copyOperands,
                        Variadic<AnyType>:$copyinOperands,
                        Variadic<AnyType>:$copyoutOperands,
                        Variadic<AnyType>:$createOperands,

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

if (succeeded(parser.parseOptionalKeyword(keyword))) {
     if (parser.parseLParen() || parser.parseOperand(operand) ||
         parser.resolveOperand(operand, type, result.operands) ||
         parser.parseRParen())
       return failure();
   }
   return success();
 }
 
+static ParseResult parseOptionalOperandAndType(OpAsmParser &parser,
+                                               StringRef keyword,
+                                               bool &hasOptional,
+                                               OperationState &result) {
+  OpAsmParser::OperandType operand;
+  Type type;
+  hasOptional = false;
+  if (succeeded(parser.parseOptionalKeyword(keyword))) {
+    hasOptional = true;
+    if (parser.parseLParen() || parser.parseOperand(operand) ||
+        parser.parseColonType(type) ||
+        parser.resolveOperand(operand, type, result.operands) ||
+        parser.parseRParen())
+      return failure();
+  }
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 /// Parse acc.parallel operation
 /// operation := `acc.parallel` `async` `(` index `)`?
 ///                             `wait` `(` index-list `)`?
 ///                             `num_gangs` `(` value `)`?

 static ParseResult parseParallelOp(OpAsmParser &parser,
                                    OperationState &result) {
   Builder &builder = parser.getBuilder();
   SmallVector<OpAsmParser::OperandType, 8> privateOperands,
       firstprivateOperands, createOperands, copyOperands, copyinOperands,
       copyoutOperands, noCreateOperands, presentOperands, devicePtrOperands,
       attachOperands, waitOperands, reductionOperands;
   SmallVector<Type, 8> operandTypes;
-  OpAsmParser::OperandType async, numGangs, numWorkers, vectorLength, ifCond,
-      selfCond;
+  OpAsmParser::OperandType ifCond, selfCond;
   bool hasAsync = false, hasNumGangs = false, hasNumWorkers = false;
   bool hasVectorLength = false, hasIfCond = false, hasSelfCond = false;
-
-  Type indexType = builder.getIndexType();
   Type i1Type = builder.getI1Type();
 
   // async()?
-  if (failed(parseOptionalOperand(parser, ParallelOp::getAsyncKeyword(), async,
-                                  indexType, hasAsync, result)))
+  if (failed(parseOptionalOperandAndType(parser, ParallelOp::getAsyncKeyword(),
+                                         hasAsync, result)))
     return failure();
 
   // wait()?
   if (failed(parseOperandList(parser, ParallelOp::getWaitKeyword(),
                               waitOperands, operandTypes, result)))
     return failure();
 
   // num_gangs(value)?
-  if (failed(parseOptionalOperand(parser, ParallelOp::getNumGangsKeyword(),
-                                  numGangs, indexType, hasNumGangs, result)))
+  if (failed(parseOptionalOperandAndType(
+          parser, ParallelOp::getNumGangsKeyword(), hasNumGangs, result)))
     return failure();
 
   // num_workers(value)?
-  if (failed(parseOptionalOperand(parser, ParallelOp::getNumWorkersKeyword(),
-                                  numWorkers, indexType, hasNumWorkers,
-                                  result)))
+  if (failed(parseOptionalOperandAndType(
+          parser, ParallelOp::getNumWorkersKeyword(), hasNumWorkers, result)))
     return failure();
 
   // vector_length(value)?
-  if (failed(parseOptionalOperand(parser, ParallelOp::getVectorLengthKeyword(),
-                                  vectorLength, indexType, hasVectorLength,
-                                  result)))
+  if (failed(parseOptionalOperandAndType(parser,
+                                         ParallelOp::getVectorLengthKeyword(),
+                                         hasVectorLength, result)))
     return failure();
 
   // if()?
   if (failed(parseOptionalOperand(parser, ParallelOp::getIfKeyword(), ifCond,
                                   i1Type, hasIfCond, result)))
     return failure();
 
   // self()?
static ParseResult parseParallelOp(OpAsmParser &parser,
   return success();
 }
 
 static void print(OpAsmPrinter &printer, ParallelOp &op) {
   printer << ParallelOp::getOperationName();
 
   // async()?
   if (Value async = op.async())
-    printer << " " << ParallelOp::getAsyncKeyword() << "(" << async << ")";
+    printer << " " << ParallelOp::getAsyncKeyword() << "(" << async << ": "
+            << async.getType() << ")";
 
   // wait()?
   printOperandList(op.waitOperands(), ParallelOp::getWaitKeyword(), printer);
 
   // num_gangs()?
   if (Value numGangs = op.numGangs())
     printer << " " << ParallelOp::getNumGangsKeyword() << "(" << numGangs
-            << ")";
+            << ": " << numGangs.getType() << ")";
 
   // num_workers()?
   if (Value numWorkers = op.numWorkers())
     printer << " " << ParallelOp::getNumWorkersKeyword() << "(" << numWorkers
-            << ")";
+            << ": " << numWorkers.getType() << ")";
 
   // vector_length()?
   if (Value vectorLength = op.vectorLength())
     printer << " " << ParallelOp::getVectorLengthKeyword() << "("
-            << vectorLength << ")";
+            << vectorLength << ": " << vectorLength.getType() << ")";
 
   // if()?
   if (Value ifCond = op.ifCond())
     printer << " " << ParallelOp::getIfKeyword() << "(" << ifCond << ")";
 
   // self()?
   if (Value selfCond = op.selfCond())
     printer << " " << ParallelOp::getSelfKeyword() << "(" << selfCond << ")";

mlir/test/Dialect/OpenACC/ops.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s
 // Verify the printed output can be parsed.
 // RUN: mlir-opt -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect  | FileCheck %s
 // Verify the generic form can be parsed.
 // RUN: mlir-opt -allow-unregistered-dialect -mlir-print-op-generic %s | mlir-opt -allow-unregistered-dialect | FileCheck %s
 
 func @compute1(%A: memref<10x10xf32>, %B: memref<10x10xf32>, %C: memref<10x10xf32>) -> memref<10x10xf32> {
   %c0 = constant 0 : index
   %c10 = constant 10 : index
   %c1 = constant 1 : index
+  %async = constant 1 : i64
 
-  acc.parallel async(%c1) {
+  acc.parallel async(%async: i64) {
     acc.loop gang vector {
       scf.for %arg3 = %c0 to %c10 step %c1 {
         scf.for %arg4 = %c0 to %c10 step %c1 {
           scf.for %arg5 = %c0 to %c10 step %c1 {
             %a = load %A[%arg3, %arg5] : memref<10x10xf32>
             %b = load %B[%arg5, %arg4] : memref<10x10xf32>
             %cij = load %C[%arg3, %arg4] : memref<10x10xf32>
             %p = mulf %a, %b : f32
func @compute1(%A: memref<10x10xf32>, %B: memref<10x10xf32>, %C: memref<10x10xf32>) -> memref<10x10xf32> {
 
   return %C : memref<10x10xf32>
 }
 
 // CHECK-LABEL: func @compute1(
 //  CHECK-NEXT:   %{{.*}} = constant 0 : index
 //  CHECK-NEXT:   %{{.*}} = constant 10 : index
 //  CHECK-NEXT:   %{{.*}} = constant 1 : index
-//  CHECK-NEXT:   acc.parallel async(%{{.*}}) {
+//  CHECK-NEXT:   [[ASYNC:%.*]] = constant 1 : i64
+//  CHECK-NEXT:   acc.parallel async([[ASYNC]]: i64) {
 //  CHECK-NEXT:     acc.loop gang vector {
 //  CHECK-NEXT:       scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
 //  CHECK-NEXT:         scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
 //  CHECK-NEXT:           scf.for %{{.*}} = %{{.*}} to %{{.*}} step %{{.*}} {
 //  CHECK-NEXT:             %{{.*}} = load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 //  CHECK-NEXT:             %{{.*}} = load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 //  CHECK-NEXT:             %{{.*}} = load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 //  CHECK-NEXT:             %{{.*}} = mulf %{{.*}}, %{{.*}} : f32

 //  CHECK-NEXT:   return %{{.*}} : memref<10x10xf32>
 //  CHECK-NEXT: }
 
 
 func @compute3(%a: memref<10x10xf32>, %b: memref<10x10xf32>, %c: memref<10xf32>, %d: memref<10xf32>) -> memref<10xf32> {
   %lb = constant 0 : index
   %st = constant 1 : index
   %c10 = constant 10 : index
+  %numGangs = constant 10 : i64
+  %numWorkers = constant 10 : i64
 
   acc.data present(%a: memref<10x10xf32>, %b: memref<10x10xf32>, %c: memref<10xf32>, %d: memref<10xf32>) {
-    acc.parallel num_gangs(%c10) num_workers(%c10) private(%c : memref<10xf32>) {
+    acc.parallel num_gangs(%numGangs: i64) num_workers(%numWorkers: i64) private(%c : memref<10xf32>) {
       acc.loop gang {
         scf.for %x = %lb to %c10 step %st {
           acc.loop worker {
             scf.for %y = %lb to %c10 step %st {
               %axy = load %a[%x, %y] : memref<10x10xf32>
               %bxy = load %b[%x, %y] : memref<10x10xf32>
               %tmp = addf %axy, %bxy : f32
               store %tmp, %c[%y] : memref<10xf32>
func @compute3(%a: memref<10x10xf32>, %b: memref<10x10xf32>, %c: memref<10xf32>, %d: memref<10xf32>) -> memref<10xf32> {
 
   return %d : memref<10xf32>
 }
 
 // CHECK:      func @compute3({{.*}}: memref<10x10xf32>, {{.*}}: memref<10x10xf32>, [[ARG2:%.*]]: memref<10xf32>, {{.*}}: memref<10xf32>) -> memref<10xf32> {
 // CHECK-NEXT:   [[C0:%.*]] = constant 0 : index
 // CHECK-NEXT:   [[C1:%.*]] = constant 1 : index
 // CHECK-NEXT:   [[C10:%.*]] = constant 10 : index
+// CHECK-NEXT:   [[NUMGANG:%.*]] = constant 10 : i64
+// CHECK-NEXT:   [[NUMWORKERS:%.*]] = constant 10 : i64
 // CHECK-NEXT:   acc.data present(%{{.*}}: memref<10x10xf32>, %{{.*}}: memref<10x10xf32>, %{{.*}}: memref<10xf32>, %{{.*}}: memref<10xf32>) {
-// CHECK-NEXT:     acc.parallel num_gangs([[C10]]) num_workers([[C10]]) private([[ARG2]]: memref<10xf32>) {
+// CHECK-NEXT:     acc.parallel num_gangs([[NUMGANG]]: i64) num_workers([[NUMWORKERS]]: i64) private([[ARG2]]: memref<10xf32>) {
 // CHECK-NEXT:       acc.loop gang {
 // CHECK-NEXT:         scf.for %{{.*}} = [[C0]] to [[C10]] step [[C1]] {
 // CHECK-NEXT:           acc.loop worker {
 // CHECK-NEXT:             scf.for %{{.*}} = [[C0]] to [[C10]] step [[C1]] {
 // CHECK-NEXT:               %{{.*}} = load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 // CHECK-NEXT:               %{{.*}} = load %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 // CHECK-NEXT:               %{{.*}} = addf %{{.*}}, %{{.*}} : f32
 // CHECK-NEXT:               store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>

 // CHECK-NEXT: }
 // CHECK-NEXT: acc.loop tile([[TILESIZE]]: i64, [[TILESIZE]]: i64) {
 // CHECK-NEXT:   "some.op"() : () -> ()
 // CHECK-NEXT:   acc.yield
 // CHECK-NEXT: }
 
 
 func @testparallelop() -> () {
-  %vectorLength = constant 128 : index
-  acc.parallel vector_length(%vectorLength) {
+  %i64value = constant 1 : i64
+  %i32value = constant 1 : i32
+  %idxValue = constant 1 : index
+  acc.parallel async(%i64value: i64) {
+  }
+  acc.parallel async(%i32value: i32) {
+  }
+  acc.parallel async(%idxValue: index) {
+  }
+  acc.parallel wait(%i64value: i64) {
+  }
+  acc.parallel wait(%i32value: i32) {
+  }
+  acc.parallel wait(%idxValue: index) {
+  }
+  acc.parallel wait(%i64value: i64, %i32value: i32, %idxValue: index) {
+  }
+  acc.parallel num_gangs(%i64value: i64) {
+  }
+  acc.parallel num_gangs(%i32value: i32) {
+  }
+  acc.parallel num_gangs(%idxValue: 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) {
+  }
+  acc.parallel vector_length(%i32value: i32) {
+  }
+  acc.parallel vector_length(%idxValue: index) {
   }
   return
 }
 
-// CHECK:      [[VECTORLENGTH:%.*]] = constant 128 : index
-// CHECK-NEXT: acc.parallel vector_length([[VECTORLENGTH]]) {
+// CHECK-LABEL: func @testparallelop
+// CHECK:      [[I64VALUE:%.*]] = constant 1 : i64
+// CHECK:      [[I32VALUE:%.*]] = constant 1 : i32
+// CHECK:      [[IDXVALUE:%.*]] = constant 1 : index
+// CHECK:      acc.parallel async([[I64VALUE]]: i64) {
+// CHECK-NEXT: }
+// CHECK:      acc.parallel async([[I32VALUE]]: i32) {
+// CHECK-NEXT: }
+// CHECK:      acc.parallel async([[IDXVALUE]]: index) {
+// CHECK-NEXT: }
+// CHECK:      acc.parallel wait([[I64VALUE]]: i64) {
+// CHECK-NEXT: }
+// CHECK:      acc.parallel wait([[I32VALUE]]: i32) {
+// CHECK-NEXT: }
+// CHECK:      acc.parallel wait([[IDXVALUE]]: index) {
+// CHECK-NEXT: }
+// CHECK:      acc.parallel wait([[I64VALUE]]: i64, [[I32VALUE]]: i32, [[IDXVALUE]]: 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_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: }
+// CHECK:      acc.parallel vector_length([[I32VALUE]]: i32) {
+// CHECK-NEXT: }
+// CHECK:      acc.parallel vector_length([[IDXVALUE]]: index) {
 // CHECK-NEXT: }