[NFC][OpenACC] addDeviceType to init/shutdown ops

[erichkeane]

As a first step of attempting to make for a 'better' interface to lowering to the OpenACC dialect, this patch adds a helper function to InitOp and ShutdownOp to make adding a device-type clause easier.

[llvmbot]

@llvm/pr-subscribers-clang
@llvm/pr-subscribers-openacc
@llvm/pr-subscribers-mlir-openacc

@llvm/pr-subscribers-mlir

Author: Erich Keane (erichkeane)

Changes

As a first step of attempting to make for a 'better' interface to lowering to the OpenACC dialect, this patch adds a helper function to InitOp and ShutdownOp to make adding a device-type clause easier.

---

Full diff: https://github.com/llvm/llvm-project/pull/137372.diff

3 Files Affected:

• (modified) clang/lib/CIR/CodeGen/CIRGenStmtOpenACC.cpp (+5-19)
• (modified) mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td (+10)
• (modified) mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp (+20)

diff --git a/clang/lib/CIR/CodeGen/CIRGenStmtOpenACC.cpp b/clang/lib/CIR/CodeGen/CIRGenStmtOpenACC.cpp
index 688fca1bf2751..5720fec556ff2 100644
--- a/clang/lib/CIR/CodeGen/CIRGenStmtOpenACC.cpp
+++ b/clang/lib/CIR/CodeGen/CIRGenStmtOpenACC.cpp
@@ -238,25 +238,11 @@ class OpenACCClauseCIREmitter final
   void VisitDeviceTypeClause(const OpenACCDeviceTypeClause &clause) {
     lastDeviceTypeClause = &clause;
     if constexpr (isOneOfTypes<OpTy, InitOp, ShutdownOp>) {
-      llvm::SmallVector<mlir::Attribute> deviceTypes;
-      std::optional<mlir::ArrayAttr> existingDeviceTypes =
-          operation.getDeviceTypes();
-
-      // Ensure we keep the existing ones, and in the correct 'new' order.
-      if (existingDeviceTypes) {
-        for (mlir::Attribute attr : *existingDeviceTypes)
-          deviceTypes.push_back(mlir::acc::DeviceTypeAttr::get(
-              builder.getContext(),
-              cast<mlir::acc::DeviceTypeAttr>(attr).getValue()));
-      }
-
-      for (const DeviceTypeArgument &arg : clause.getArchitectures()) {
-        deviceTypes.push_back(mlir::acc::DeviceTypeAttr::get(
-            builder.getContext(), decodeDeviceType(arg.getIdentifierInfo())));
-      }
-      operation.removeDeviceTypesAttr();
-      operation.setDeviceTypesAttr(
-          mlir::ArrayAttr::get(builder.getContext(), deviceTypes));
+      llvm::for_each(
+          clause.getArchitectures(), [this](const DeviceTypeArgument &arg) {
+            operation.addDeviceType(builder.getContext(),
+                                    decodeDeviceType(arg.getIdentifierInfo()));
+          });
     } else if constexpr (isOneOfTypes<OpTy, SetOp>) {
       assert(!operation.getDeviceTypeAttr() && "already have device-type?");
       assert(clause.getArchitectures().size() <= 1);
diff --git a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
index 5e249e639d837..2167129e9e1c7 100644
--- a/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
+++ b/mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td
@@ -2614,6 +2614,11 @@ def OpenACC_InitOp : OpenACC_Op<"init", [AttrSizedOperandSegments]> {
                        Optional<IntOrIndex>:$deviceNum,
                        Optional<I1>:$ifCond);
 
+  let extraClassDeclaration = [{
+    /// Adds a device type to the list of device types for this directive.
+    void addDeviceType(MLIRContext *, mlir::acc::DeviceType);
+  }];
+
   let assemblyFormat = [{
     oilist(`device_num` `(` $deviceNum `:` type($deviceNum) `)`
       | `if` `(` $ifCond `)`
@@ -2645,6 +2650,11 @@ def OpenACC_ShutdownOp : OpenACC_Op<"shutdown", [AttrSizedOperandSegments]> {
                        Optional<IntOrIndex>:$deviceNum,
                        Optional<I1>:$ifCond);
 
+  let extraClassDeclaration = [{
+    /// Adds a device type to the list of device types for this directive.
+    void addDeviceType(MLIRContext *, mlir::acc::DeviceType);
+  }];
+
   let assemblyFormat = [{
     oilist(`device_num` `(` $deviceNum `:` type($deviceNum) `)`
     |`if` `(` $ifCond `)`
diff --git a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
index 06a1f9ef0c8cb..04cbe200eafe9 100644
--- a/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
+++ b/mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp
@@ -2954,6 +2954,16 @@ LogicalResult acc::InitOp::verify() {
   return success();
 }
 
+void acc::InitOp::addDeviceType(MLIRContext *context,
+                                mlir::acc::DeviceType deviceType) {
+  llvm::SmallVector<mlir::Attribute> deviceTypes;
+  if (getDeviceTypesAttr())
+    llvm::copy(getDeviceTypesAttr(), std::back_inserter(deviceTypes));
+
+  deviceTypes.push_back(acc::DeviceTypeAttr::get(context, deviceType));
+  setDeviceTypesAttr(mlir::ArrayAttr::get(context, deviceTypes));
+}
+
 //===----------------------------------------------------------------------===//
 // ShutdownOp
 //===----------------------------------------------------------------------===//
@@ -2966,6 +2976,16 @@ LogicalResult acc::ShutdownOp::verify() {
   return success();
 }
 
+void acc::ShutdownOp::addDeviceType(MLIRContext *context,
+                                    mlir::acc::DeviceType deviceType) {
+  llvm::SmallVector<mlir::Attribute> deviceTypes;
+  if (getDeviceTypesAttr())
+    llvm::copy(getDeviceTypesAttr(), std::back_inserter(deviceTypes));
+
+  deviceTypes.push_back(acc::DeviceTypeAttr::get(context, deviceType));
+  setDeviceTypesAttr(mlir::ArrayAttr::get(context, deviceTypes));
+}
+
 //===----------------------------------------------------------------------===//
 // SetOp
 //===----------------------------------------------------------------------===//

[clementval]

LGTM

[razvanlupusoru]

Thank you! I am happy to see this improvement!