[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

87–91

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

152–154

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

152–154

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

87–91

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

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

152–154

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

[clementval]

Address review comments

[clementval]

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

87–91

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

87–91

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

87–91

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

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

def OpenACC_ReductionOpAttr : StrEnumAttr<"ReductionOpAttr",
      OpenACC_ReductionOpMin, OpenACC_ReductionOpAnd, OpenACC_ReductionOpOr,
      OpenACC_ReductionOpXor, OpenACC_ReductionOpLogEqv,
      OpenACC_ReductionOpLogNeqv, OpenACC_ReductionOpLogAnd,
      OpenACC_ReductionOpLogOr
     ]> {
   let cppNamespace = "::mlir::acc";
 }
 
+// Type used in operation below
+def IntOrIndex : AnyTypeOf<[AnyInteger, Index]>;
+
 //===----------------------------------------------------------------------===//
 // 2.5.1 parallel Construct
 //===----------------------------------------------------------------------===//
 
 def OpenACC_ParallelOp : OpenACC_Op<"parallel",
     [AttrSizedOperandSegments]> {
   let summary = "parallel construct";
   let description = [{
     The "acc.parallel" operation represents a parallel construct block. It has
     one region to be executued in parallel on the current device.
 
     Example:
 
     ```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<IntOrIndex>:$async,
+                       Variadic<IntOrIndex>:$waitOperands,
+                       Optional<IntOrIndex>:$numGangs,
+                       Optional<IntOrIndex>:$numWorkers,
+                       Optional<IntOrIndex>:$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 OptionalParseResult parseOptionalOperandAndType(OpAsmParser &parser,
+                                                       StringRef keyword,
+                                                       OperationState &result) {
+  OpAsmParser::OperandType operand;
+  Type type;
+  if (succeeded(parser.parseOptionalKeyword(keyword))) {
+    if (parser.parseLParen() || parser.parseOperand(operand) ||
+        parser.parseColonType(type) ||
+        parser.resolveOperand(operand, type, result.operands) ||
+        parser.parseRParen())
+      return failure();
+    return success();
+  }
+  return llvm::None;
+}
+
 //===----------------------------------------------------------------------===//
 // 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;
-  bool hasAsync = false, hasNumGangs = false, hasNumWorkers = false;
-  bool hasVectorLength = false, hasIfCond = false, hasSelfCond = false;
-
-  Type indexType = builder.getIndexType();
+  OpAsmParser::OperandType ifCond, selfCond;
+  bool hasIfCond = false, hasSelfCond = false;
+  OptionalParseResult async, numGangs, numWorkers, vectorLength;
   Type i1Type = builder.getI1Type();
 
   // async()?
-  if (failed(parseOptionalOperand(parser, ParallelOp::getAsyncKeyword(), async,
-                                  indexType, hasAsync, result)))
+  async = parseOptionalOperandAndType(parser, ParallelOp::getAsyncKeyword(),
+                                      result);
+  if (async.hasValue() && failed(*async))
     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)))
+  numGangs = parseOptionalOperandAndType(
+      parser, ParallelOp::getNumGangsKeyword(), result);
+  if (numGangs.hasValue() && failed(*numGangs))
     return failure();
 
   // num_workers(value)?
-  if (failed(parseOptionalOperand(parser, ParallelOp::getNumWorkersKeyword(),
-                                  numWorkers, indexType, hasNumWorkers,
-                                  result)))
+  numWorkers = parseOptionalOperandAndType(
+      parser, ParallelOp::getNumWorkersKeyword(), result);
+  if (numWorkers.hasValue() && failed(*numWorkers))
     return failure();
 
   // vector_length(value)?
-  if (failed(parseOptionalOperand(parser, ParallelOp::getVectorLengthKeyword(),
-                                  vectorLength, indexType, hasVectorLength,
-                                  result)))
+  vectorLength = parseOptionalOperandAndType(
+      parser, ParallelOp::getVectorLengthKeyword(), result);
+  if (vectorLength.hasValue() && failed(*vectorLength))
     return failure();
 
   // if()?
   if (failed(parseOptionalOperand(parser, ParallelOp::getIfKeyword(), ifCond,
                                   i1Type, hasIfCond, result)))
     return failure();
 
   // self()?
static ParseResult parseParallelOp(OpAsmParser &parser,
   if (failed(parseOperandList(parser, ParallelOp::getFirstPrivateKeyword(),
                               firstprivateOperands, operandTypes, result)))
     return failure();
 
   // Parallel op region
   if (failed(parseRegions<ParallelOp>(parser, result)))
     return failure();
 
-  result.addAttribute(ParallelOp::getOperandSegmentSizeAttr(),
-                      builder.getI32VectorAttr(
-                          {static_cast<int32_t>(hasAsync ? 1 : 0),
-                           static_cast<int32_t>(waitOperands.size()),
-                           static_cast<int32_t>(hasNumGangs ? 1 : 0),
-                           static_cast<int32_t>(hasNumWorkers ? 1 : 0),
-                           static_cast<int32_t>(hasVectorLength ? 1 : 0),
-                           static_cast<int32_t>(hasIfCond ? 1 : 0),
-                           static_cast<int32_t>(hasSelfCond ? 1 : 0),
-                           static_cast<int32_t>(reductionOperands.size()),
-                           static_cast<int32_t>(copyOperands.size()),
-                           static_cast<int32_t>(copyinOperands.size()),
-                           static_cast<int32_t>(copyoutOperands.size()),
-                           static_cast<int32_t>(createOperands.size()),
-                           static_cast<int32_t>(noCreateOperands.size()),
-                           static_cast<int32_t>(presentOperands.size()),
-                           static_cast<int32_t>(devicePtrOperands.size()),
-                           static_cast<int32_t>(attachOperands.size()),
-                           static_cast<int32_t>(privateOperands.size()),
-                           static_cast<int32_t>(firstprivateOperands.size())}));
+  result.addAttribute(
+      ParallelOp::getOperandSegmentSizeAttr(),
+          {static_cast<int32_t>(async.hasValue() ? 1 : 0),
+           static_cast<int32_t>(numGangs.hasValue() ? 1 : 0),
+           static_cast<int32_t>(numWorkers.hasValue() ? 1 : 0),
+           static_cast<int32_t>(vectorLength.hasValue() ? 1 : 0),
 
   // Additional attributes
   if (failed(parser.parseOptionalAttrDictWithKeyword(result.attributes)))
     return failure();
 
   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: }