[openacc] Add acc routine support to acc dialect

[razvanlupusoru]

Summary

Adds representation for acc routine under new operation named
acc.routine. This operation is associated with a function symbol.
It also gets its own compiler generated synthetic symbol name so
that it can be referenced from the associated function. The clauses
associated with the acc routine directive are captured in the
acc.routine op.

The linking between the func.func and its acc.routine declaration
is done through the acc.routine_info attribute. In practice, a
single acc routine is associated with a function. But the spec does
not specifically restrict this - thus the 1:N relationship between
func.func and acc.routine allowed in the dialect. Additionally, it
makes sense that multiple acc routines could be used for a single
function depending on loop context - to allow flexible parallelization.

Most acc routine clauses are supported including gang, gang(dim:),
vector, worker, seq, nohost, and bind. The only one not
supported is device_type. This is because most other clauses also
miss this and the effort to add support for it needs to be coordinated
and consistent.

[clementval]

LGTM. Just couple of formatting nits. Also the pre-commit ci is failing for a clang-format issue.

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

1051–1054

Remove brace

1074–1076

Remove brace. I know that in some other place we keep them but in lowering we follow the llvm guidelines for this.

1077–1079

Remove brace

1092–1095

Remove brace

[vzakhari]

Looks good modulo Valentin's remarks.

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

1363

nit: it may be worth verifying somewhere that the symbols are actually defined by acc.routine operations.

[razvanlupusoru]

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

1363

I agree. I will need to figure out how to do this effectively - I possibly need a verifier pass. The attribute itself cannot do the verification because a lookup in symbol table needs an operation - and the attribute verifier does not have that.

[razvanlupusoru]

Add description to operation and attribute. Address reviewer comments.

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

def OpenACC_GlobalDestructorOp : OpenACC_Op<"global_dtor",
   let assemblyFormat = [{
     $sym_name $region attr-dict-with-keyword
   }];
 
   let hasVerifier = 0;
 }
 
 //===----------------------------------------------------------------------===//
+// 2.15.1 Routine Directive
+//===----------------------------------------------------------------------===//
+
+def OpenACC_RoutineOp : OpenACC_Op<"routine", [IsolatedFromAbove]> {
+  let summary = "acc routine operation";
+
+  let description = [{
+    The `acc.routine` operation is used to capture the clauses of acc
+    routine directive, including the associated function name. The associated
+    function keeps track of its corresponding routine declaration through
+    the `RoutineInfoAttr`.
+
+    Example:
+
+    ```mlir
+    func.func @acc_func(%a : i64) -> () attributes {acc.routine_info = #acc.routine_info<[@acc_func_rout1]>} {
+      return
+    }
+    acc.routine @acc_func_rout1 func(@acc_func) gang
+    ```
+  }];
+
+  let arguments = (ins SymbolNameAttr:$sym_name,
+                       SymbolNameAttr:$func_name,
+                       OptionalAttr<StrAttr>:$bind_name,
+                       OptionalAttr<UnitAttr>:$gang,
+                       OptionalAttr<UnitAttr>:$worker,
+                       OptionalAttr<UnitAttr>:$vector,
+                       OptionalAttr<UnitAttr>:$seq,
+                       OptionalAttr<UnitAttr>:$nohost,
+                       OptionalAttr<UnitAttr>:$implicit,
+                       OptionalAttr<APIntAttr>:$gangDim);
+
+  let extraClassDeclaration = [{
+    static StringRef getGangDimKeyword() { return "dim"; }
+  }];
+
+  let assemblyFormat = [{
+    $sym_name `func` `(` $func_name `)`
+    oilist (
+        `bind` `(` $bind_name `)`
+      | `gang` `` custom<RoutineGangClause>($gang, $gangDim)
+      | `worker` $worker
+      | `vector` $vector
+      | `seq` $seq
+      | `nohost` $nohost
+      | `implicit` $implicit
+    ) attr-dict-with-keyword
+  }];
+
+  let hasVerifier = 1;
+}
+
+def RoutineInfoAttr : OpenACC_Attr<"RoutineInfo", "routine_info"> {
+  let summary = "Keeps track of associated acc routine information";
+
+  let description = [{
+    This attribute is used to create the association between a function and
+    its `acc.routine` operation. A `func.func` uses this if its name
+    was referenced in an `acc routine` directive.
+  }];
+
+  let parameters = (ins ArrayRefParameter<"SymbolRefAttr", "">:$accRoutines);
+  let assemblyFormat = "`<` `[` `` $accRoutines `]` `>`";
+}
+
+//===----------------------------------------------------------------------===//
 // 2.14.1. Init Directive
 //===----------------------------------------------------------------------===//
 
 def OpenACC_InitOp : OpenACC_Op<"init", [AttrSizedOperandSegments]> {
   let summary = "init operation";
 
   let description = [{
     The "acc.init" operation represents the OpenACC init executable

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

 // DeclareExitOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult acc::DeclareExitOp::verify() {
   return checkDeclareOperands(*this, this->getDataClauseOperands());
 }
 
 //===----------------------------------------------------------------------===//
+// RoutineOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult acc::RoutineOp::verify() {
+  int parallelism = 0;
+  parallelism += getGang() ? 1 : 0;
+  parallelism += getWorker() ? 1 : 0;
+  parallelism += getVector() ? 1 : 0;
+  parallelism += getSeq() ? 1 : 0;
+
+  if (parallelism > 1)
+    return emitError() << "only one of `gang`, `worker`, `vector`, `seq` can "
+                          "be present at the same time";
+
+  return success();
+}
+
+static ParseResult parseRoutineGangClause(OpAsmParser &parser, UnitAttr &gang,
+                                          IntegerAttr &gangDim) {
+  // Since gang clause exists, ensure that unit attribute is set.
+  gang = UnitAttr::get(parser.getBuilder().getContext());
+
+  // Next, look for dim on gang. Don't initialize `gangDim` yet since
+  // we leave it without attribute if there is no `dim` specifier.
+  if (succeeded(parser.parseOptionalLParen())) {
+    // Look for syntax that looks like `dim = 1 : i32`.
+    // Thus first look for `dim =`
+    if (failed(parser.parseKeyword(RoutineOp::getGangDimKeyword())) ||
+        failed(parser.parseEqual()))
+      return failure();
+
+    int64_t dimValue;
+    Type valueType;
+    // Now look for `1 : i32`
+    if (failed(parser.parseInteger(dimValue)) ||
+        failed(parser.parseColonType(valueType)))
+      return failure();
+
+    gangDim = IntegerAttr::get(valueType, dimValue);
+
+    if (failed(parser.parseRParen()))
+      return failure();
+  }
+
+  return success();
+}
+
+void printRoutineGangClause(OpAsmPrinter &p, Operation *op, UnitAttr gang,
+                            IntegerAttr gangDim) {
+  if (gangDim)
+    p << "(" << RoutineOp::getGangDimKeyword() << " = " << gangDim.getValue()
+      << " : " << gangDim.getType() << ")";
+}
+
+//===----------------------------------------------------------------------===//
 // InitOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult acc::InitOp::verify() {
   Operation *currOp = *this;
   while ((currOp = currOp->getParentOp()))
     if (isComputeOperation(currOp))
       return emitOpError("cannot be nested in a compute operation");

mlir/test/Dialect/OpenACC/ops.mlir

 // CHECK: %[[ADDR:.*]] = llvm.mlir.addressof @globalvar {acc.declare = #acc.declare<dataClause = acc_create>} : !llvm.ptr<i32>
 // CHECK-NEXT: %[[CREATE:.*]] = acc.create varPtr(%[[ADDR]] : !llvm.ptr<i32>) -> !llvm.ptr<i32>
 // CHECK-NEXT: acc.declare_enter dataOperands(%[[CREATE]] : !llvm.ptr<i32>)
 // CHECK: acc.global_dtor @acc_destructor
 // CHECK: %[[ADDR:.*]] = llvm.mlir.addressof @globalvar {acc.declare = #acc.declare<dataClause = acc_create>} : !llvm.ptr<i32>
 // CHECK-NEXT: %[[DELETE:.*]] = acc.getdeviceptr varPtr(%[[ADDR]] : !llvm.ptr<i32>) -> !llvm.ptr<i32> {dataClause = #acc<data_clause acc_create>}
 // CHECK-NEXT: acc.declare_exit dataOperands(%[[DELETE]] : !llvm.ptr<i32>)
 // CHECK-NEXT: acc.delete accPtr(%[[DELETE]] : !llvm.ptr<i32>)
+
+// -----
+
+func.func @acc_func(%a : i64) -> () attributes {acc.routine_info = #acc.routine_info<[@acc_func_rout1,@acc_func_rout2,@acc_func_rout3,
+    @acc_func_rout4,@acc_func_rout5,@acc_func_rout6,@acc_func_rout7,@acc_func_rout8,@acc_func_rout9]>} {
+  return
+}
+
+acc.routine @acc_func_rout1 func(@acc_func)
+acc.routine @acc_func_rout2 func(@acc_func) bind("acc_func_gpu")
+acc.routine @acc_func_rout3 func(@acc_func) bind("acc_func_gpu_gang") gang
+acc.routine @acc_func_rout4 func(@acc_func) bind("acc_func_gpu_vector") vector
+acc.routine @acc_func_rout5 func(@acc_func) bind("acc_func_gpu_worker") worker
+acc.routine @acc_func_rout6 func(@acc_func) bind("acc_func_gpu_seq") seq
+acc.routine @acc_func_rout7 func(@acc_func) bind("acc_func_gpu_imp_gang") implicit gang
+acc.routine @acc_func_rout8 func(@acc_func) bind("acc_func_gpu_vector_nohost") vector nohost
+acc.routine @acc_func_rout9 func(@acc_func) bind("acc_func_gpu_gang_dim1") gang(dim = 1 : i32)
+
+// CHECK-LABEL: func.func @acc_func(
+// CHECK: attributes {acc.routine_info = #acc.routine_info<[@acc_func_rout1, @acc_func_rout2, @acc_func_rout3,
+// CHECK: @acc_func_rout4, @acc_func_rout5, @acc_func_rout6, @acc_func_rout7, @acc_func_rout8, @acc_func_rout9]>}
+// CHECK: acc.routine @acc_func_rout1 func(@acc_func)
+// CHECK: acc.routine @acc_func_rout2 func(@acc_func) bind("acc_func_gpu")
+// CHECK: acc.routine @acc_func_rout3 func(@acc_func) bind("acc_func_gpu_gang") gang
+// CHECK: acc.routine @acc_func_rout4 func(@acc_func) bind("acc_func_gpu_vector") vector
+// CHECK: acc.routine @acc_func_rout5 func(@acc_func) bind("acc_func_gpu_worker") worker
+// CHECK: acc.routine @acc_func_rout6 func(@acc_func) bind("acc_func_gpu_seq") seq
+// CHECK: acc.routine @acc_func_rout7 func(@acc_func) bind("acc_func_gpu_imp_gang") gang implicit
+// CHECK: acc.routine @acc_func_rout8 func(@acc_func) bind("acc_func_gpu_vector_nohost") vector nohost
+// CHECK: acc.routine @acc_func_rout9 func(@acc_func) bind("acc_func_gpu_gang_dim1") gang(dim = 1 : i32)