[NFC][OpenACC] addDeviceType to init/shutdown ops (#137372)

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.

Author: erichkeane

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);

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 `)`

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
 //===----------------------------------------------------------------------===//