[mlir][openacc] Add missing attributes and operands for acc.loop

[clementval]

Summary

This patch add the missing operands to the acc.loop operation. Only the device_type
information is not part of the operation for now.

[clementval]

Change I1Attr to UnitAttr

[rriddle]

mlir/include/mlir/Dialect/OpenACC/OpenACC.h

38

This seems unrelated?

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

231

nit: UnitAttr is already an Optional construct, so this is redundant.

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

483–484

nit: Remove this extra newline.

612

nit: Nest this within the existing if for gangNum

621–635

nit: You don't need the == ... here.

630

Same here.

[kiranchandramohan]

LGTM. Please wait for @rriddle's approval.

mlir/test/Dialect/OpenACC/ops.mlir

79

Is it required to have the name attributes?

[clementval]

Address review comments

[clementval]

@rriddle and @kiranchandramohan Thanks for the reviews. I just updated the patche to address the comments.

mlir/include/mlir/Dialect/OpenACC/OpenACC.h

38

Seq is not really part of exec mapping but in group of attributes that includes auto, independent and seq and as they are mutually exclusive it doesn't really belong here anymore.

mlir/test/Dialect/OpenACC/ops.mlir

79

Those are UnitAttr so if they are present it means that are "active".

[clementval]

@rriddle Are you fine with the changes?

mlir/test/Dialect/OpenACC/ops.mlir

79

Sorry, I guess you meant the attributes keyword before the attributes-dict? Yes it is required. I used ParseResult parseOptionalAttrDictWithKeyword(NamedAttrList &result) in the parser for those ops so you need the keyword.

[rriddle]

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

518

nit: Drop trivial braces here.

[clementval]

Address review comments

mlir/include/mlir/Dialect/OpenACC/OpenACC.h

 
 #include "mlir/Dialect/OpenACC/OpenACCOpsDialect.h.inc"
 
 /// Enumeration used to encode the execution mapping on a loop construct.
 /// They refer directly to the OpenACC 3.0 standard:
 /// 2.9.2. gang
 /// 2.9.3. worker
 /// 2.9.4. vector
-/// 2.9.5. seq
 ///
 /// Value can be combined bitwise to reflect the mapping applied to the
 /// construct. e.g. `acc.loop gang vector`, the `gang` and `vector` could be
 /// combined and the final mapping value would be 5 (4 | 1).
-enum OpenACCExecMapping { NONE = 0, VECTOR = 1, WORKER = 2, GANG = 4, SEQ = 8 };
+enum OpenACCExecMapping { NONE = 0, VECTOR = 1, WORKER = 2, GANG = 4 };
 
 } // end namespace acc
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_OPENACC_OPENACC_H_

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,
+                       UnitAttr:$loopSeq,
+                       UnitAttr:$loopIndependent,
+                       UnitAttr:$loopAuto,
+                       Variadic<AnyInteger>:$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 getCollapseAttrName() { return "collapse"; }
-    static StringRef getExecutionMappingAttrName() { return "exec_mapping"; }
-    static StringRef getGangAttrName() { return "gang"; }
     static StringRef getSeqAttrName() { return "seq"; }
-    static StringRef getVectorAttrName() { return "vector"; }
-    static StringRef getWorkerAttrName() { return "worker"; }
+    static StringRef getIndependentAttrName() { return "independent"; }
+    static StringRef getAutoAttrName() { return "auto"; }
+    static StringRef getExecutionMappingAttrName() { return "exec_mapping"; }
+    static StringRef getGangKeyword() { return "gang"; }
+    static StringRef getGangNumKeyword() { return "num"; }
+    static StringRef getGangStaticKeyword() { return "static"; }
+    static StringRef getVectorKeyword() { return "vector"; }
+    static StringRef getWorkerKeyword() { return "worker"; }
+    static StringRef getTileKeyword() { return "tile"; }
     static StringRef getPrivateKeyword() { return "private"; }
     static StringRef getReductionKeyword() { return "reduction"; }
   }];
 
   let verifier = ?;
 }
 
 // Yield operation for the acc.loop and acc.parallel operations.

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

printer.printOptionalAttrDictWithKeyword(
       op.getAttrs(), ParallelOp::getOperandSegmentSizeAttr());
 }
 
 //===----------------------------------------------------------------------===//
 // LoopOp
 //===----------------------------------------------------------------------===//
 
 /// Parse acc.loop operation
-/// operation := `acc.loop` `gang`? `vector`? `worker`? `seq`?
+/// operation := `acc.loop` `gang`? `vector`? `worker`?
 ///                         `private` `(` value-list `)`?
 ///                         `reduction` `(` value-list `)`?
 ///                         region attr-dict?
 static ParseResult parseLoopOp(OpAsmParser &parser, OperationState &result) {
-
   Builder &builder = parser.getBuilder();
   unsigned executionMapping = 0;
   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();
 
   // gang?
-  if (succeeded(parser.parseOptionalKeyword(LoopOp::getGangAttrName())))
+  if (succeeded(parser.parseOptionalKeyword(LoopOp::getGangKeyword())))
     executionMapping |= OpenACCExecMapping::GANG;
 
-  // vector?
-  if (succeeded(parser.parseOptionalKeyword(LoopOp::getVectorAttrName())))
-    executionMapping |= OpenACCExecMapping::VECTOR;
+  // 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)) {
+        return failure();
+      }
+    }
+    parser.parseOptionalComma();
+    if (succeeded(
+            parser.parseOptionalKeyword(LoopOp::getGangStaticKeyword()))) {
+      hasGangStatic = true;
+      parser.parseColon();
+      if (parser.parseOperand(gangStatic) ||
+          parser.resolveOperand(gangStatic, intType, result.operands)) {
+        return failure();
+      }
+    }
+    if (failed(parser.parseRParen()))
+      return failure();
+  }
 
   // worker?
-  if (succeeded(parser.parseOptionalKeyword(LoopOp::getWorkerAttrName())))
+  if (succeeded(parser.parseOptionalKeyword(LoopOp::getWorkerKeyword())))
     executionMapping |= OpenACCExecMapping::WORKER;
 
-  // seq?
-  if (succeeded(parser.parseOptionalKeyword(LoopOp::getSeqAttrName())))
-    executionMapping |= OpenACCExecMapping::SEQ;
+  // optional worker operand
+  if (succeeded(parser.parseOptionalLParen())) {
+    hasWorkerNum = true;
+    if (parser.parseOperand(workerNum) ||
+        parser.resolveOperand(workerNum, intType, 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.parseRParen()) {
+      return failure();
+    }
+  }
+
+  // tile()?
+  if (failed(parseOperandList(parser, LoopOp::getTileKeyword(), tileOperands,
+                              operandTypes, result)))
+    return failure();
 
   // private()?
   if (failed(parseOperandList(parser, LoopOp::getPrivateKeyword(),
                               privateOperands, operandTypes, result)))
     return failure();
 
   // reduction()?
   if (failed(parseOperandList(parser, LoopOp::getReductionKeyword(),
                               reductionOperands, operandTypes, result)))
     return failure();
 
   if (executionMapping != 0)
     result.addAttribute(LoopOp::getExecutionMappingAttrName(),
                         builder.getI64IntegerAttr(executionMapping));
 
   // Parse optional results in case there is a reduce.
   if (parser.parseOptionalArrowTypeList(result.types))
     return failure();
 
   if (failed(parseRegions<LoopOp>(parser, result)))
     return failure();
 
   result.addAttribute(LoopOp::getOperandSegmentSizeAttr(),
                       builder.getI32VectorAttr(
-                          {static_cast<int32_t>(privateOperands.size()),
+                          {static_cast<int32_t>(hasGangNum ? 1 : 0),
+                           static_cast<int32_t>(hasGangStatic ? 1 : 0),
+                           static_cast<int32_t>(hasWorkerNum ? 1 : 0),
+                           static_cast<int32_t>(hasVectorLength ? 1 : 0),
+                           static_cast<int32_t>(tileOperands.size()),
+                           static_cast<int32_t>(privateOperands.size()),
                            static_cast<int32_t>(reductionOperands.size())}));
 
   if (failed(parser.parseOptionalAttrDictWithKeyword(result.attributes)))
     return failure();
 
   return success();
 }
 
 static void print(OpAsmPrinter &printer, LoopOp &op) {
   printer << LoopOp::getOperationName();
 
   unsigned execMapping =
       (op.getAttrOfType<IntegerAttr>(LoopOp::getExecutionMappingAttrName()) !=
        nullptr)
           ? op.getAttrOfType<IntegerAttr>(LoopOp::getExecutionMappingAttrName())
                 .getInt()
           : 0;
-  if ((execMapping & OpenACCExecMapping::GANG) == OpenACCExecMapping::GANG)
-    printer << " " << LoopOp::getGangAttrName();
 
-  if ((execMapping & OpenACCExecMapping::WORKER) == OpenACCExecMapping::WORKER)
-    printer << " " << LoopOp::getWorkerAttrName();
+  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;
+        if (gangStatic)
+          printer << ", ";
+      }
+      if (gangStatic)
+        printer << LoopOp::getGangStaticKeyword() << ": " << gangStatic;
+      printer << ")";
+    }
+  }
+
+  if (execMapping & OpenACCExecMapping::WORKER) {
+    printer << " " << LoopOp::getWorkerKeyword();
 
-  if ((execMapping & OpenACCExecMapping::VECTOR) == OpenACCExecMapping::VECTOR)
-    printer << " " << LoopOp::getVectorAttrName();
+    // Print optional worker operand if present
+    if (Value workerNum = op.workerNum())
+      printer << "(" << workerNum << ")";
+  }
+
+  if (execMapping & OpenACCExecMapping::VECTOR) {
+    printer << " " << LoopOp::getVectorKeyword();
+
+    // Print optional vector operand if present
+    if (Value vectorLength = op.vectorLength())
+      printer << "(" << vectorLength << ")";
+  }
 
-  if ((execMapping & OpenACCExecMapping::SEQ) == OpenACCExecMapping::SEQ)
-    printer << " " << LoopOp::getSeqAttrName();
+  // tile()?
+  printOperandList(op.tileOperands(), LoopOp::getTileKeyword(), printer);
 
   // private()?
   printOperandList(op.privateOperands(), LoopOp::getPrivateKeyword(), printer);
 
   // reduction()?
   printOperandList(op.reductionOperands(), LoopOp::getReductionKeyword(),
                    printer);
 

mlir/test/Dialect/OpenACC/ops.mlir

 //  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 seq {
+    acc.loop {
       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
             %co = addf %cij, %p : f32
             store %co, %C[%arg3, %arg4] : memref<10x10xf32>
           }
         }
       }
       acc.yield
-    }
+    } attributes {seq}
     acc.yield
   }
 
   return %C : memref<10x10xf32>
 }
 
 // CHECK-LABEL: func @compute2(
 //  CHECK-NEXT:   %{{.*}} = constant 0 : index
 //  CHECK-NEXT:   %{{.*}} = constant 10 : index
 //  CHECK-NEXT:   %{{.*}} = constant 1 : index
 //  CHECK-NEXT:   acc.parallel {
-//  CHECK-NEXT:     acc.loop seq {
+//  CHECK-NEXT:     acc.loop {
 //  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:             %{{.*}} = addf %{{.*}}, %{{.*}} : f32
 //  CHECK-NEXT:             store %{{.*}}, %{{.*}}[%{{.*}}, %{{.*}}] : memref<10x10xf32>
 //  CHECK-NEXT:           }
 //  CHECK-NEXT:         }
 //  CHECK-NEXT:       }
 //  CHECK-NEXT:       acc.yield
-//  CHECK-NEXT:     }
+//  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
acc.parallel num_gangs(%c10) num_workers(%c10) private(%c : memref<10xf32>) {
               %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>
             }
             acc.yield
           }
 
-          acc.loop seq {
+          acc.loop {
             // for i = 0 to 10 step 1
             //   d[x] += c[i]
             scf.for %i = %lb to %c10 step %st {
               %ci = load %c[%i] : memref<10xf32>
               %dx = load %d[%x] : memref<10xf32>
               %z = addf %ci, %dx : f32
               store %z, %d[%x] : memref<10xf32>
             }
             acc.yield
-          }
+          } attributes {seq}
         }
         acc.yield
       }
       acc.yield
     }
     acc.terminator
   }
 

 // 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.yield
 // CHECK-NEXT:           }
-// CHECK-NEXT:           acc.loop seq {
+// CHECK-NEXT:           acc.loop {
 // CHECK-NEXT:             scf.for %{{.*}} = [[C0]] to [[C10]] step [[C1]] {
 // CHECK-NEXT:               %{{.*}} = load %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:               %{{.*}} = load %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:               %{{.*}} = addf %{{.*}}, %{{.*}} : f32
 // CHECK-NEXT:               store %{{.*}}, %{{.*}}[%{{.*}}] : memref<10xf32>
 // CHECK-NEXT:             }
 // CHECK-NEXT:             acc.yield
-// CHECK-NEXT:           }
+// CHECK-NEXT:           } attributes {seq}
 // CHECK-NEXT:         }
 // CHECK-NEXT:         acc.yield
 // CHECK-NEXT:       }
 // CHECK-NEXT:       acc.yield
 // CHECK-NEXT:     }
 // CHECK-NEXT:     acc.terminator
 // CHECK-NEXT:   }
 // CHECK-NEXT:   return %{{.*}} : memref<10xf32>
-// CHECK-NEXT: }
+// CHECK-NEXT: }
+
+func @testop() -> () {
+  %workerNum = constant 1 : i64
+  %vectorLength = constant 128 : i64
+  %gangNum = constant 8 : i64
+  %gangStatic = constant 2 : i64
+  %tileSize = constant 2 : i64
+  acc.loop gang worker vector {
+  }
+  acc.loop gang(num: %gangNum) {
+  }
+  acc.loop gang(static: %gangStatic) {
+  }
+  acc.loop worker(%workerNum) {
+  }
+  acc.loop vector(%vectorLength) {
+  }
+  acc.loop gang(num: %gangNum) worker vector {
+  }
+  acc.loop gang(num: %gangNum, static: %gangStatic) worker(%workerNum) vector(%vectorLength) {
+  }
+  acc.loop tile(%tileSize : i64, %tileSize : i64) {
+  }
+  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-NEXT: }
+// CHECK-NEXT: acc.loop gang(num: [[GANGNUM]]) {
+// CHECK-NEXT: }
+// CHECK-NEXT: acc.loop gang(static: [[GANGSTATIC]]) {
+// CHECK-NEXT: }
+// CHECK-NEXT: acc.loop worker([[WORKERNUM]]) {
+// CHECK-NEXT: }
+// CHECK-NEXT: acc.loop vector([[VECTORLENGTH]]) {
+// CHECK-NEXT: }
+// CHECK-NEXT: acc.loop gang(num: [[GANGNUM]]) worker vector {
+// CHECK-NEXT: }
+// CHECK-NEXT: acc.loop gang(num: [[GANGNUM]], static: [[GANGSTATIC]]) worker([[WORKERNUM]]) vector([[VECTORLENGTH]]) {
+// CHECK-NEXT: }
+// CHECK-NEXT: acc.loop tile([[TILESIZE]]: i64, [[TILESIZE]]: i64) {
+// CHECK-NEXT: }