[mlir][openacc] Support Index and AnyInteger in loop op

[clementval]

Summary

Following patch D87712, this patch switch AnyInteger for operands gangNum, gangStatic,
workerNum, vectoreLength and tileOperands to Index and AnyInteger.

[ftynse]

mlir/test/Dialect/OpenACC/ops.mlir

209–233

Syntax nit: we had a similar model with two colons in attribute dictionaries some time ago, but later switched to keyword=value:type. You may want to use such a syntax here as well, but I won't insist.

[clementval]

mlir/test/Dialect/OpenACC/ops.mlir

209–233

I think it is a good idea so I'll update the syntax before pushing.

[clementval]

Address review comment

[clementval]

Update comment to reflect operation syntax.

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

acc.loop gang vector {
       }
       acc.yield
     } attributes { collapse = 3 }
     ```
   }];
 
 
   let arguments = (ins OptionalAttr<I64Attr>:$collapse,
-                       Optional<AnyInteger>:$gangNum,
-                       Optional<AnyInteger>:$gangStatic,
-                       Optional<AnyInteger>:$workerNum,
-                       Optional<AnyInteger>:$vectorLength,
+                       Optional<IntOrIndex>:$gangNum,
+                       Optional<IntOrIndex>:$gangStatic,
+                       Optional<IntOrIndex>:$workerNum,
+                       Optional<IntOrIndex>:$vectorLength,
                        UnitAttr:$seq,
                        UnitAttr:$independent,
                        UnitAttr:$auto_,
-                       Variadic<AnyInteger>:$tileOperands,
+                       Variadic<IntOrIndex>:$tileOperands,
                        Variadic<AnyType>:$privateOperands,
                        OptionalAttr<OpenACC_ReductionOpAttr>:$reductionOp,
                        Variadic<AnyType>:$reductionOperands,
                        DefaultValuedAttr<I64Attr, "0">:$exec_mapping);
 
   let results = (outs Variadic<AnyType>:$results);
 
   let regions = (region AnyRegion:$region);

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

printer.printOptionalAttrDictWithKeyword(
       op.getAttrs(), ParallelOp::getOperandSegmentSizeAttr());
 }
 
 //===----------------------------------------------------------------------===//
 // LoopOp
 //===----------------------------------------------------------------------===//
 
 /// Parse acc.loop operation
-/// operation := `acc.loop` `gang`? `vector`? `worker`?
-///                         `private` `(` value-list `)`?
-///                         `reduction` `(` value-list `)`?
-///                         region attr-dict?
+/// operation := `acc.loop`
+///              (`gang` ( `(` (`num=` value)? (`,` `static=` value `)`)? )? )?
+///              (`vector` ( `(` value `)` )? )? (`worker` (`(` value `)`)? )?
+///              (`vector_length` `(` value `)`)?
+///              (`tile` `(` value-list `)`)?
+///              (`private` `(` value-list `)`)?
+///              (`reduction` `(` value-list `)`)?
+///              region attr-dict?
 static ParseResult parseLoopOp(OpAsmParser &parser, OperationState &result) {
   Builder &builder = parser.getBuilder();
   unsigned executionMapping = OpenACCExecMapping::NONE;
   SmallVector<Type, 8> operandTypes;
   SmallVector<OpAsmParser::OperandType, 8> privateOperands, reductionOperands;
   SmallVector<OpAsmParser::OperandType, 8> tileOperands;
   bool hasWorkerNum = false, hasVectorLength = false, hasGangNum = false;
   bool hasGangStatic = false;
   OpAsmParser::OperandType workerNum, vectorLength, gangNum, gangStatic;
-  Type intType = builder.getI64Type();
+  Type gangNumType, gangStaticType, workerType, vectorLengthType;
 
   // gang?
   if (succeeded(parser.parseOptionalKeyword(LoopOp::getGangKeyword())))
     executionMapping |= OpenACCExecMapping::GANG;
 
   // optional gang operand
   if (succeeded(parser.parseOptionalLParen())) {
     if (succeeded(parser.parseOptionalKeyword(LoopOp::getGangNumKeyword()))) {
       hasGangNum = true;
-      parser.parseColon();
-      if (parser.parseOperand(gangNum) ||
-          parser.resolveOperand(gangNum, intType, result.operands)) {
+      parser.parseEqual();
+      if (parser.parseOperand(gangNum) || parser.parseColonType(gangNumType) ||
+          parser.resolveOperand(gangNum, gangNumType, result.operands)) {
         return failure();
       }
     }
     parser.parseOptionalComma();
     if (succeeded(
             parser.parseOptionalKeyword(LoopOp::getGangStaticKeyword()))) {
       hasGangStatic = true;
-      parser.parseColon();
+      parser.parseEqual();
       if (parser.parseOperand(gangStatic) ||
-          parser.resolveOperand(gangStatic, intType, result.operands)) {
+          parser.parseColonType(gangStaticType) ||
+          parser.resolveOperand(gangStatic, gangStaticType, result.operands)) {
         return failure();
       }
     }
     if (failed(parser.parseRParen()))
       return failure();
   }
 
   // worker?
   if (succeeded(parser.parseOptionalKeyword(LoopOp::getWorkerKeyword())))
     executionMapping |= OpenACCExecMapping::WORKER;
 
   // optional worker operand
   if (succeeded(parser.parseOptionalLParen())) {
     hasWorkerNum = true;
-    if (parser.parseOperand(workerNum) ||
-        parser.resolveOperand(workerNum, intType, result.operands) ||
+    if (parser.parseOperand(workerNum) || parser.parseColonType(workerType) ||
+        parser.resolveOperand(workerNum, workerType, result.operands) ||
         parser.parseRParen()) {
       return failure();
     }
   }
 
   // vector?
   if (succeeded(parser.parseOptionalKeyword(LoopOp::getVectorKeyword())))
     executionMapping |= OpenACCExecMapping::VECTOR;
 
   // optional vector operand
   if (succeeded(parser.parseOptionalLParen())) {
     hasVectorLength = true;
     if (parser.parseOperand(vectorLength) ||
-        parser.resolveOperand(vectorLength, intType, result.operands) ||
+        parser.parseColonType(vectorLengthType) ||
+        parser.resolveOperand(vectorLength, vectorLengthType,
+                              result.operands) ||
         parser.parseRParen()) {
       return failure();
     }
   }
 
   // tile()?
   if (failed(parseOperandList(parser, LoopOp::getTileKeyword(), tileOperands,
                               operandTypes, result)))
if (execMapping & OpenACCExecMapping::GANG) {
     printer << " " << LoopOp::getGangKeyword();
     Value gangNum = op.gangNum();
     Value gangStatic = op.gangStatic();
 
     // Print optional gang operands
     if (gangNum || gangStatic) {
       printer << "(";
       if (gangNum) {
-        printer << LoopOp::getGangNumKeyword() << ": " << gangNum;
+        printer << LoopOp::getGangNumKeyword() << "=" << gangNum << ": "
+                << gangNum.getType();
         if (gangStatic)
           printer << ", ";
       }
       if (gangStatic)
-        printer << LoopOp::getGangStaticKeyword() << ": " << gangStatic;
+        printer << LoopOp::getGangStaticKeyword() << "=" << gangStatic << ": "
+                << gangStatic.getType();
       printer << ")";
     }
   }
 
   if (execMapping & OpenACCExecMapping::WORKER) {
     printer << " " << LoopOp::getWorkerKeyword();
 
     // Print optional worker operand if present
     if (Value workerNum = op.workerNum())
-      printer << "(" << workerNum << ")";
+      printer << "(" << workerNum << ": " << workerNum.getType() << ")";
   }
 
   if (execMapping & OpenACCExecMapping::VECTOR) {
     printer << " " << LoopOp::getVectorKeyword();
 
     // Print optional vector operand if present
     if (Value vectorLength = op.vectorLength())
-      printer << "(" << vectorLength << ")";
+      printer << "(" << vectorLength << ": " << vectorLength.getType() << ")";
   }
 
   // tile()?
   printOperandList(op.tileOperands(), LoopOp::getTileKeyword(), printer);
 
   // private()?
   printOperandList(op.privateOperands(), LoopOp::getPrivateKeyword(), printer);
 

mlir/test/Dialect/OpenACC/ops.mlir

-// RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s
+// RUN: mlir-opt -split-input-file -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
+// RUN: mlir-opt -split-input-file -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
+// RUN: mlir-opt -split-input-file -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(%async: i64) {

 //  CHECK-NEXT:       }
 //  CHECK-NEXT:       acc.yield
 //  CHECK-NEXT:     } attributes {collapse = 3 : i64}
 //  CHECK-NEXT:     acc.yield
 //  CHECK-NEXT:   }
 //  CHECK-NEXT:   return %{{.*}} : memref<10x10xf32>
 //  CHECK-NEXT: }
 
+// -----
+
 func @compute2(%A: memref<10x10xf32>, %B: memref<10x10xf32>, %C: memref<10x10xf32>) -> memref<10x10xf32> {
   %c0 = constant 0 : index
   %c10 = constant 10 : index
   %c1 = constant 1 : index
 
   acc.parallel {
     acc.loop {
       scf.for %arg3 = %c0 to %c10 step %c1 {

 //  CHECK-NEXT:       }
 //  CHECK-NEXT:       acc.yield
 //  CHECK-NEXT:     } attributes {seq}
 //  CHECK-NEXT:     acc.yield
 //  CHECK-NEXT:   }
 //  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
 

 // CHECK-NEXT:       }
 // CHECK-NEXT:       acc.yield
 // CHECK-NEXT:     }
 // CHECK-NEXT:     acc.terminator
 // CHECK-NEXT:   }
 // CHECK-NEXT:   return %{{.*}} : memref<10xf32>
 // CHECK-NEXT: }
 
-func @testop(%a: memref<10xf32>) -> () {
-  %workerNum = constant 1 : i64
-  %vectorLength = constant 128 : i64
-  %gangNum = constant 8 : i64
-  %gangStatic = constant 2 : i64
-  %tileSize = constant 2 : i64
+// -----
+
+func @testloopop() -> () {
+  %i64Value = constant 1 : i64
+  %i32Value = constant 128 : i32
+  %idxValue = constant 8 : index
+
   acc.loop gang worker vector {
     "some.op"() : () -> ()
     acc.yield
   }
-  acc.loop gang(num: %gangNum) {
+  acc.loop gang(num=%i64Value: i64) {
+    "some.op"() : () -> ()
+    acc.yield
+  }
+  acc.loop gang(static=%i64Value: i64) {
+    "some.op"() : () -> ()
+    acc.yield
+  }
+  acc.loop worker(%i64Value: i64) {
+    "some.op"() : () -> ()
+    acc.yield
+  }
+  acc.loop worker(%i32Value: i32) {
+    "some.op"() : () -> ()
+    acc.yield
+  }
+  acc.loop worker(%idxValue: index) {
+    "some.op"() : () -> ()
+    acc.yield
+  }
+  acc.loop vector(%i64Value: i64) {
+    "some.op"() : () -> ()
+    acc.yield
+  }
+  acc.loop vector(%i32Value: i32) {
     "some.op"() : () -> ()
     acc.yield
   }
-  acc.loop gang(static: %gangStatic) {
+  acc.loop vector(%idxValue: index) {
     "some.op"() : () -> ()
     acc.yield
   }
-  acc.loop worker(%workerNum) {
+  acc.loop gang(num=%i64Value: i64) worker vector {
     "some.op"() : () -> ()
     acc.yield
   }
-  acc.loop vector(%vectorLength) {
+  acc.loop gang(num=%i64Value: i64, static=%i64Value: i64) worker(%i64Value: i64) vector(%i64Value: i64) {
     "some.op"() : () -> ()
     acc.yield
   }
-  acc.loop gang(num: %gangNum) worker vector {
+  acc.loop gang(num=%i32Value: i32, static=%idxValue: index) {
     "some.op"() : () -> ()
     acc.yield
   }
-  acc.loop gang(num: %gangNum, static: %gangStatic) worker(%workerNum) vector(%vectorLength) {
+  acc.loop tile(%i64Value: i64, %i64Value: i64) {
     "some.op"() : () -> ()
     acc.yield
   }
-  acc.loop tile(%tileSize : i64, %tileSize : i64) {
+  acc.loop tile(%i32Value: i32, %i32Value: i32) {
     "some.op"() : () -> ()
     acc.yield
   }
   return
 }
 
-// CHECK:      [[WORKERNUM:%.*]] = constant 1 : i64
-// CHECK-NEXT: [[VECTORLENGTH:%.*]] = constant 128 : i64
-// CHECK-NEXT: [[GANGNUM:%.*]] = constant 8 : i64
-// CHECK-NEXT: [[GANGSTATIC:%.*]] = constant 2 : i64
-// CHECK-NEXT: [[TILESIZE:%.*]] = constant 2 : i64
-// CHECK-NEXT: acc.loop gang worker vector {
+// CHECK:      [[I64VALUE:%.*]] = constant 1 : i64
+// CHECK-NEXT: [[I32VALUE:%.*]] = constant 128 : i32
+// CHECK-NEXT: [[IDXVALUE:%.*]] = constant 8 : index
+// CHECK:      acc.loop gang worker vector {
+// CHECK-NEXT:   "some.op"() : () -> ()
+// CHECK-NEXT:   acc.yield
+// CHECK-NEXT: }
+// CHECK:      acc.loop gang(num=[[I64VALUE]]: i64) {
+// CHECK-NEXT:   "some.op"() : () -> ()
+// CHECK-NEXT:   acc.yield
+// CHECK-NEXT: }
+// CHECK:      acc.loop gang(static=[[I64VALUE]]: i64) {
+// CHECK-NEXT:   "some.op"() : () -> ()
+// CHECK-NEXT:   acc.yield
+// CHECK-NEXT: }
+// CHECK:      acc.loop worker([[I64VALUE]]: i64) {
+// CHECK-NEXT:   "some.op"() : () -> ()
+// CHECK-NEXT:   acc.yield
+// CHECK-NEXT: }
+// CHECK:      acc.loop worker([[I32VALUE]]: i32) {
+// CHECK-NEXT:   "some.op"() : () -> ()
+// CHECK-NEXT:   acc.yield
+// CHECK-NEXT: }
+// CHECK:      acc.loop worker([[IDXVALUE]]: index) {
+// CHECK-NEXT:   "some.op"() : () -> ()
+// CHECK-NEXT:   acc.yield
+// CHECK-NEXT: }
+// CHECK:      acc.loop vector([[I64VALUE]]: i64) {
 // CHECK-NEXT:   "some.op"() : () -> ()
 // CHECK-NEXT:   acc.yield
 // CHECK-NEXT: }
-// CHECK-NEXT: acc.loop gang(num: [[GANGNUM]]) {
+// CHECK:      acc.loop vector([[I32VALUE]]: i32) {
 // CHECK-NEXT:   "some.op"() : () -> ()
 // CHECK-NEXT:   acc.yield
 // CHECK-NEXT: }
-// CHECK-NEXT: acc.loop gang(static: [[GANGSTATIC]]) {
+// CHECK:      acc.loop vector([[IDXVALUE]]: index) {
 // CHECK-NEXT:   "some.op"() : () -> ()
 // CHECK-NEXT:   acc.yield
 // CHECK-NEXT: }
-// CHECK-NEXT: acc.loop worker([[WORKERNUM]]) {
+// CHECK:      acc.loop gang(num=[[I64VALUE]]: i64) worker vector {
 // CHECK-NEXT:   "some.op"() : () -> ()
 // CHECK-NEXT:   acc.yield
 // CHECK-NEXT: }
-// CHECK-NEXT: acc.loop vector([[VECTORLENGTH]]) {
+// CHECK:      acc.loop gang(num=[[I64VALUE]]: i64, static=[[I64VALUE]]: i64) worker([[I64VALUE]]: i64) vector([[I64VALUE]]: i64) {
 // CHECK-NEXT:   "some.op"() : () -> ()
 // CHECK-NEXT:   acc.yield
 // CHECK-NEXT: }
-// CHECK-NEXT: acc.loop gang(num: [[GANGNUM]]) worker vector {
+// CHECK:      acc.loop gang(num=[[I32VALUE]]: i32, static=[[IDXVALUE]]: index) {
 // CHECK-NEXT:   "some.op"() : () -> ()
 // CHECK-NEXT:   acc.yield
 // CHECK-NEXT: }
-// CHECK-NEXT: acc.loop gang(num: [[GANGNUM]], static: [[GANGSTATIC]]) worker([[WORKERNUM]]) vector([[VECTORLENGTH]]) {
+// CHECK:      acc.loop tile([[I64VALUE]]: i64, [[I64VALUE]]: i64) {
 // CHECK-NEXT:   "some.op"() : () -> ()
 // CHECK-NEXT:   acc.yield
 // CHECK-NEXT: }
-// CHECK-NEXT: acc.loop tile([[TILESIZE]]: i64, [[TILESIZE]]: i64) {
+// CHECK:      acc.loop tile([[I32VALUE]]: i32, [[I32VALUE]]: i32) {
 // CHECK-NEXT:   "some.op"() : () -> ()
 // CHECK-NEXT:   acc.yield
 // CHECK-NEXT: }
 
+// -----
 
 func @testparallelop(%a: memref<10xf32>, %b: memref<10xf32>, %c: memref<10x10xf32>) -> () {
   %i64value = constant 1 : i64
   %i32value = constant 1 : i32
   %idxValue = constant 1 : index
   acc.parallel async(%i64value: i64) {
   }
   acc.parallel async(%i32value: i32) {