https://github.com/RiverDave updated 
https://github.com/llvm/llvm-project/pull/206576

>From 9d0f66f6fe743d75f3cae14b5e7412e646b6116a Mon Sep 17 00:00:00 2001
From: David Rivera <[email protected]>
Date: Mon, 29 Jun 2026 15:46:11 -0400
Subject: [PATCH 1/3] [CIR] Add offload container operation

Introduce cir.offload.container, a CIR dialect operation that groups one host 
CIR module with one or more device CIR modules in a single IR unit.

Nested modules are tagged with the new cir.offload.kind enum attribute using 
#cir.offload_kind<host> and #cir.offload_kind<device>. The verifier enforces 
the structural contract expected by follow-up offload merge/split pipeline 
patches: host module first, device modules after it, only nested builtin.module 
ops, and at least one device module.

This patch only adds the IR representation and verifier tests; the passes that 
create, consume, or split the container are left to later patches.
---
 .../clang/CIR/Dialect/IR/CIRDialect.td        |  1 +
 clang/include/clang/CIR/Dialect/IR/CIROps.td  | 59 ++++++++++++++++++
 clang/lib/CIR/Dialect/IR/CIRDialect.cpp       | 61 +++++++++++++++++++
 .../test/CIR/IR/invalid-offload-container.cir | 54 ++++++++++++++++
 clang/test/CIR/IR/offload-container.cir       | 32 ++++++++++
 5 files changed, 207 insertions(+)
 create mode 100644 clang/test/CIR/IR/invalid-offload-container.cir
 create mode 100644 clang/test/CIR/IR/offload-container.cir

diff --git a/clang/include/clang/CIR/Dialect/IR/CIRDialect.td 
b/clang/include/clang/CIR/Dialect/IR/CIRDialect.td
index c20af04f97a1a..7d189361eee48 100644
--- a/clang/include/clang/CIR/Dialect/IR/CIRDialect.td
+++ b/clang/include/clang/CIR/Dialect/IR/CIRDialect.td
@@ -50,6 +50,7 @@ def CIR_Dialect : Dialect {
     static llvm::StringRef getModuleLevelAsmAttrName() { return 
"cir.module_asm"; }
     static llvm::StringRef getGlobalCtorsAttrName() { return 
"cir.global_ctors"; }
     static llvm::StringRef getGlobalDtorsAttrName() { return 
"cir.global_dtors"; }
+    static llvm::StringRef getOffloadKindAttrName() { return 
"cir.offload.kind"; }
     static llvm::StringRef getOperandSegmentSizesAttrName() { return 
"operandSegmentSizes"; }
     static llvm::StringRef getNoCallerSavedRegsAttrName() { return 
"no_caller_saved_registers"; }
     static llvm::StringRef getNoCallbackAttrName() { return "nocallback"; }
diff --git a/clang/include/clang/CIR/Dialect/IR/CIROps.td 
b/clang/include/clang/CIR/Dialect/IR/CIROps.td
index 7d1c48b994b27..0c18aac7f4eb4 100644
--- a/clang/include/clang/CIR/Dialect/IR/CIROps.td
+++ b/clang/include/clang/CIR/Dialect/IR/CIROps.td
@@ -5672,6 +5672,65 @@ def CIR_VecSplatOp : CIR_Op<"vec.splat", [
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// OffloadKind
+//===----------------------------------------------------------------------===//
+
+def CIR_OffloadKind : CIR_I32EnumAttr<"OffloadKind", "offload kind", [
+  I32EnumAttrCase<"Host", 0, "host">,
+  I32EnumAttrCase<"Device", 1, "device">
+]> {
+  let genSpecializedAttr = 0;
+}
+
+def CIR_OffloadKindAttr : CIR_EnumAttr<CIR_OffloadKind, "offload_kind"> {
+  let summary = "Offload kind (host or device)";
+}
+
+//===----------------------------------------------------------------------===//
+// OffloadContainerOp
+//===----------------------------------------------------------------------===//
+
+def CIR_OffloadContainerOp : CIR_Op<"offload.container", [
+    NoRegionArguments, NoTerminator, SingleBlock, SymbolTable]> {
+  let summary = "Container for host and device CIR modules";
+  let description = [{
+    `cir.offload.container` groups one host CIR module with one or more device
+    CIR modules for offload-aware analysis and transformation.
+
+    The body holds nested `builtin.module` operations. The first nested
+    module is the host module and must carry
+    `cir.offload.kind = #cir.offload_kind<host>`. All remaining nested
+    modules are device modules and must carry
+    `cir.offload.kind = #cir.offload_kind<device>`. There must be at least
+    one device module.
+
+    Example:
+
+    ```mlir
+    cir.offload.container {
+      builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
+      }
+      builtin.module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
+      }
+    }
+    ```
+  }];
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let assemblyFormat = "$body attr-dict";
+
+  let hasVerifier = 1;
+  let hasLLVMLowering = false;
+
+  let extraClassDeclaration = [{
+    mlir::ModuleOp getHostModule();
+    llvm::iterator_range<mlir::Block::op_iterator<mlir::ModuleOp>>
+    getDeviceModules();
+  }];
+}
+
 
//===----------------------------------------------------------------------===//
 // BaseClassAddrOp
 
//===----------------------------------------------------------------------===//
diff --git a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp 
b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
index 67cc5e09f26d0..96501d5808168 100644
--- a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
+++ b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
@@ -18,6 +18,7 @@
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Value.h"
@@ -2334,6 +2335,66 @@ LogicalResult cir::VTTAddrPointOp::verify() {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// OffloadContainerOp
+//===----------------------------------------------------------------------===//
+
+mlir::ModuleOp cir::OffloadContainerOp::getHostModule() {
+  return mlir::cast<mlir::ModuleOp>(getBody().front().front());
+}
+
+llvm::iterator_range<mlir::Block::op_iterator<mlir::ModuleOp>>
+cir::OffloadContainerOp::getDeviceModules() {
+  mlir::Block &body = getBody().front();
+  auto begin = body.op_begin<mlir::ModuleOp>();
+  auto end = body.op_end<mlir::ModuleOp>();
+  if (begin != end)
+    ++begin;
+  return {begin, end};
+}
+
+static LogicalResult checkOffloadKind(mlir::ModuleOp module,
+                                      cir::OffloadKind expected) {
+  auto attr = module->getAttrOfType<cir::OffloadKindAttr>(
+      cir::CIRDialect::getOffloadKindAttrName());
+  if (!attr)
+    return module.emitOpError()
+           << "expects '" << cir::CIRDialect::getOffloadKindAttrName()
+           << "' offload kind attribute";
+  if (attr.getValue() != expected)
+    return module.emitOpError()
+           << "expects '" << cir::CIRDialect::getOffloadKindAttrName()
+           << "' value '" << cir::stringifyOffloadKind(expected) << "'";
+  return success();
+}
+
+LogicalResult cir::OffloadContainerOp::verify() {
+  mlir::Block &body = getBody().front();
+  if (body.empty())
+    return emitOpError() << "expects host module as the first nested op";
+
+  auto host = mlir::dyn_cast<mlir::ModuleOp>(body.front());
+  if (!host)
+    return emitOpError() << "expects host module as the first nested op";
+  if (failed(checkOffloadKind(host, cir::OffloadKind::Host)))
+    return failure();
+
+  unsigned numDevices = 0;
+  auto it = body.begin();
+  for (++it; it != body.end(); ++it) {
+    auto module = mlir::dyn_cast<mlir::ModuleOp>(*it);
+    if (!module)
+      return emitOpError() << "expects only nested builtin.module ops";
+    if (failed(checkOffloadKind(module, cir::OffloadKind::Device)))
+      return failure();
+    ++numDevices;
+  }
+
+  if (numDevices == 0)
+    return emitOpError() << "expects at least one device module";
+  return success();
+}
+
 
//===----------------------------------------------------------------------===//
 // FuncOp
 
//===----------------------------------------------------------------------===//
diff --git a/clang/test/CIR/IR/invalid-offload-container.cir 
b/clang/test/CIR/IR/invalid-offload-container.cir
new file mode 100644
index 0000000000000..219442f39674a
--- /dev/null
+++ b/clang/test/CIR/IR/invalid-offload-container.cir
@@ -0,0 +1,54 @@
+// RUN: cir-opt %s -verify-diagnostics -split-input-file
+
+module {
+  cir.offload.container {
+    builtin.module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} { // expected-error {{expects 'cir.offload.kind' 
value 'host'}}
+    }
+  }
+}
+
+// -----
+
+module {
+  cir.offload.container {
+    builtin.module @host_0 attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
+    }
+    builtin.module @host_1 attributes {cir.offload.kind = 
#cir.offload_kind<host>} { // expected-error {{expects 'cir.offload.kind' value 
'device'}}
+    }
+  }
+}
+
+// -----
+
+module {
+  cir.offload.container {
+    builtin.module @host { // expected-error {{expects 'cir.offload.kind' 
offload kind attribute}}
+    }
+  }
+}
+
+// -----
+
+module {
+  cir.offload.container { // expected-error {{expects only nested 
builtin.module ops}}
+    builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
+    }
+    cir.const #cir.int<0> : !cir.int<s, 32>
+  }
+}
+
+// -----
+
+module {
+  cir.offload.container { // expected-error {{expects at least one device 
module}}
+    builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
+    }
+  }
+}
+
+// -----
+
+module {
+  cir.offload.container { // expected-error {{expects host module as the first 
nested op}}
+  }
+}
diff --git a/clang/test/CIR/IR/offload-container.cir 
b/clang/test/CIR/IR/offload-container.cir
new file mode 100644
index 0000000000000..b7b53aa973ce5
--- /dev/null
+++ b/clang/test/CIR/IR/offload-container.cir
@@ -0,0 +1,32 @@
+// RUN: cir-opt %s -split-input-file --verify-roundtrip | FileCheck %s
+
+module {
+  cir.offload.container {
+    builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
+    }
+    builtin.module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
+    }
+  }
+}
+
+// CHECK: cir.offload.container {
+// CHECK:   builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
+// CHECK:   builtin.module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
+
+// -----
+
+module {
+  cir.offload.container {
+    builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
+    }
+    builtin.module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
+    }
+    builtin.module @device_1 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
+    }
+  }
+}
+
+// CHECK: cir.offload.container {
+// CHECK:   builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
+// CHECK:   builtin.module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
+// CHECK:   builtin.module @device_1 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {

>From 999b3bb1a5f14302ee224dc073863035bec829be Mon Sep 17 00:00:00 2001
From: David Rivera <[email protected]>
Date: Mon, 29 Jun 2026 16:00:27 -0400
Subject: [PATCH 2/3] fix fmt

---
 clang/lib/CIR/Dialect/IR/CIRDialect.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp 
b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
index 96501d5808168..1c1924099db7d 100644
--- a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
+++ b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
@@ -17,8 +17,8 @@
 #include "clang/CIR/Dialect/IR/CIRTypes.h"
 
 #include "mlir/IR/Attributes.h"
-#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Value.h"

>From 3faf82003e71ddca8b6bdfcf3b99edd172d954cf Mon Sep 17 00:00:00 2001
From: David Rivera <[email protected]>
Date: Fri, 14 Aug 2026 08:11:59 -0400
Subject: [PATCH 3/3] [CIR] Drop `cir.offload.container` and represent offload
 modules through a unit market on MLIR modules.

---
 .../include/clang/CIR/Dialect/IR/CIRAttrs.td  |  42 +++++
 .../include/clang/CIR/Dialect/IR/CIRDialect.h |  12 ++
 .../clang/CIR/Dialect/IR/CIRDialect.td        |   3 +
 clang/include/clang/CIR/Dialect/IR/CIROps.td  |  59 -------
 clang/lib/CIR/Dialect/IR/CIRDialect.cpp       | 166 +++++++++++-------
 .../test/CIR/IR/invalid-offload-container.cir |  70 +++++---
 clang/test/CIR/IR/offload-container.cir       |  38 ++--
 7 files changed, 228 insertions(+), 162 deletions(-)

diff --git a/clang/include/clang/CIR/Dialect/IR/CIRAttrs.td 
b/clang/include/clang/CIR/Dialect/IR/CIRAttrs.td
index f5f4f28f8993c..4695871176285 100644
--- a/clang/include/clang/CIR/Dialect/IR/CIRAttrs.td
+++ b/clang/include/clang/CIR/Dialect/IR/CIRAttrs.td
@@ -110,6 +110,48 @@ def CIR_SourceLanguageAttr : 
CIR_EnumAttr<CIR_SourceLanguage, "lang"> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// OffloadKindAttr
+//===----------------------------------------------------------------------===//
+
+def CIR_OffloadKind : CIR_I32EnumAttr<"OffloadKind", "offload kind", [
+  I32EnumAttrCase<"Host", 0, "host">,
+  I32EnumAttrCase<"Device", 1, "device">
+]> {
+  let genSpecializedAttr = 0;
+}
+
+def CIR_OffloadKindAttr : CIR_EnumAttr<CIR_OffloadKind, "offload_kind"> {
+  let summary = "Offload kind of a module in an offload container";
+  let description = [{
+    Marks a module as either the host module or one of the device modules of
+    an offload container, keeping a host CIR module and its associated device
+    CIR modules in one IR unit while offload merge/split pipeline passes need
+    visibility into both sides.
+
+    A module carrying the `cir.offload.container` unit attribute is such a
+    container. Its body holds only nested modules: the first one is the host
+    module and must carry `cir.offload.kind = #cir.offload_kind<host>`, and
+    all remaining ones are device modules and must carry
+    `cir.offload.kind = #cir.offload_kind<device>`. There must be at least one
+    device module. Keeping the host module first gives later passes a simple
+    convention for finding the host side while iterating the remaining device
+    modules.
+
+    Example:
+    ```mlir
+    module attributes {cir.offload.container} {
+      module @host attributes {cir.offload.kind = #cir.offload_kind<host>} {}
+      module @device attributes {cir.offload.kind = #cir.offload_kind<device>} 
{}
+    }
+    ```
+
+    The attribute names `cir.offload.container` and `cir.offload.kind` are
+    defined by the `getOffloadContainerAttrName` and `getOffloadKindAttrName`
+    methods in the CIRDialect class.
+  }];
+}
+
 
//===----------------------------------------------------------------------===//
 // ArgPassingKind + RecordLayoutAttr
 
//===----------------------------------------------------------------------===//
diff --git a/clang/include/clang/CIR/Dialect/IR/CIRDialect.h 
b/clang/include/clang/CIR/Dialect/IR/CIRDialect.h
index 970a6984a5b05..77201fe08274b 100644
--- a/clang/include/clang/CIR/Dialect/IR/CIRDialect.h
+++ b/clang/include/clang/CIR/Dialect/IR/CIRDialect.h
@@ -45,6 +45,18 @@ void buildTerminatedBody(mlir::OpBuilder &builder, 
mlir::Location loc);
 /// Look up the RecordLayoutAttr for a named record in the module's
 /// cir.record_layouts dictionary.  Asserts if the entry is missing.
 RecordLayoutAttr getRecordLayout(mlir::ModuleOp module, mlir::StringAttr name);
+
+/// Returns whether the module is an offload container, i.e. whether it carries
+/// the cir.offload.container unit attribute.  See CIR_OffloadKindAttr for the
+/// structure such a module is required to have.
+bool isOffloadContainer(mlir::ModuleOp module);
+
+/// Returns the host module of an offload container.
+mlir::ModuleOp getOffloadHostModule(mlir::ModuleOp container);
+
+/// Returns the device modules of an offload container, in container order.
+llvm::iterator_range<mlir::Block::op_iterator<mlir::ModuleOp>>
+getOffloadDeviceModules(mlir::ModuleOp container);
 } // namespace cir
 
 // TableGen'erated files for MLIR dialects require that a macro be defined when
diff --git a/clang/include/clang/CIR/Dialect/IR/CIRDialect.td 
b/clang/include/clang/CIR/Dialect/IR/CIRDialect.td
index 7d189361eee48..6ac63ab253e70 100644
--- a/clang/include/clang/CIR/Dialect/IR/CIRDialect.td
+++ b/clang/include/clang/CIR/Dialect/IR/CIRDialect.td
@@ -26,6 +26,8 @@ def CIR_Dialect : Dialect {
 
   let useDefaultAttributePrinterParser = 1;
 
+  let hasOperationAttrVerify = 1;
+
   // Enable constant materialization for the CIR dialect. This generates a
   // declaration for the cir::CIRDialect::materializeConstant function. This
   // hook is necessary for canonicalization to properly handle attributes
@@ -50,6 +52,7 @@ def CIR_Dialect : Dialect {
     static llvm::StringRef getModuleLevelAsmAttrName() { return 
"cir.module_asm"; }
     static llvm::StringRef getGlobalCtorsAttrName() { return 
"cir.global_ctors"; }
     static llvm::StringRef getGlobalDtorsAttrName() { return 
"cir.global_dtors"; }
+    static llvm::StringRef getOffloadContainerAttrName() { return 
"cir.offload.container"; }
     static llvm::StringRef getOffloadKindAttrName() { return 
"cir.offload.kind"; }
     static llvm::StringRef getOperandSegmentSizesAttrName() { return 
"operandSegmentSizes"; }
     static llvm::StringRef getNoCallerSavedRegsAttrName() { return 
"no_caller_saved_registers"; }
diff --git a/clang/include/clang/CIR/Dialect/IR/CIROps.td 
b/clang/include/clang/CIR/Dialect/IR/CIROps.td
index 0c18aac7f4eb4..7d1c48b994b27 100644
--- a/clang/include/clang/CIR/Dialect/IR/CIROps.td
+++ b/clang/include/clang/CIR/Dialect/IR/CIROps.td
@@ -5672,65 +5672,6 @@ def CIR_VecSplatOp : CIR_Op<"vec.splat", [
   }];
 }
 
-//===----------------------------------------------------------------------===//
-// OffloadKind
-//===----------------------------------------------------------------------===//
-
-def CIR_OffloadKind : CIR_I32EnumAttr<"OffloadKind", "offload kind", [
-  I32EnumAttrCase<"Host", 0, "host">,
-  I32EnumAttrCase<"Device", 1, "device">
-]> {
-  let genSpecializedAttr = 0;
-}
-
-def CIR_OffloadKindAttr : CIR_EnumAttr<CIR_OffloadKind, "offload_kind"> {
-  let summary = "Offload kind (host or device)";
-}
-
-//===----------------------------------------------------------------------===//
-// OffloadContainerOp
-//===----------------------------------------------------------------------===//
-
-def CIR_OffloadContainerOp : CIR_Op<"offload.container", [
-    NoRegionArguments, NoTerminator, SingleBlock, SymbolTable]> {
-  let summary = "Container for host and device CIR modules";
-  let description = [{
-    `cir.offload.container` groups one host CIR module with one or more device
-    CIR modules for offload-aware analysis and transformation.
-
-    The body holds nested `builtin.module` operations. The first nested
-    module is the host module and must carry
-    `cir.offload.kind = #cir.offload_kind<host>`. All remaining nested
-    modules are device modules and must carry
-    `cir.offload.kind = #cir.offload_kind<device>`. There must be at least
-    one device module.
-
-    Example:
-
-    ```mlir
-    cir.offload.container {
-      builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
-      }
-      builtin.module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
-      }
-    }
-    ```
-  }];
-
-  let regions = (region SizedRegion<1>:$body);
-
-  let assemblyFormat = "$body attr-dict";
-
-  let hasVerifier = 1;
-  let hasLLVMLowering = false;
-
-  let extraClassDeclaration = [{
-    mlir::ModuleOp getHostModule();
-    llvm::iterator_range<mlir::Block::op_iterator<mlir::ModuleOp>>
-    getDeviceModules();
-  }];
-}
-
 
//===----------------------------------------------------------------------===//
 // BaseClassAddrOp
 
//===----------------------------------------------------------------------===//
diff --git a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp 
b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
index 1c1924099db7d..58a8f064a0611 100644
--- a/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
+++ b/clang/lib/CIR/Dialect/IR/CIRDialect.cpp
@@ -111,6 +111,112 @@ Operation 
*cir::CIRDialect::materializeConstant(mlir::OpBuilder &builder,
                                  mlir::cast<mlir::TypedAttr>(value));
 }
 
+//===----------------------------------------------------------------------===//
+// Offload container helpers
+//===----------------------------------------------------------------------===//
+
+bool cir::isOffloadContainer(mlir::ModuleOp module) {
+  return module->hasAttr(cir::CIRDialect::getOffloadContainerAttrName());
+}
+
+mlir::ModuleOp cir::getOffloadHostModule(mlir::ModuleOp container) {
+  assert(isOffloadContainer(container) && "expected an offload container");
+  return mlir::cast<mlir::ModuleOp>(container.getBody()->front());
+}
+
+llvm::iterator_range<mlir::Block::op_iterator<mlir::ModuleOp>>
+cir::getOffloadDeviceModules(mlir::ModuleOp container) {
+  assert(isOffloadContainer(container) && "expected an offload container");
+  mlir::Block &body = *container.getBody();
+  auto begin = body.op_begin<mlir::ModuleOp>();
+  auto end = body.op_end<mlir::ModuleOp>();
+  if (begin != end)
+    ++begin;
+  return {begin, end};
+}
+
+//===----------------------------------------------------------------------===//
+// Dialect attribute verification
+//===----------------------------------------------------------------------===//
+
+static LogicalResult verifyOffloadKind(mlir::ModuleOp module,
+                                       cir::OffloadKind expected) {
+  auto attr = module->getAttrOfType<cir::OffloadKindAttr>(
+      cir::CIRDialect::getOffloadKindAttrName());
+  if (!attr)
+    return module.emitOpError()
+           << "expects '" << cir::CIRDialect::getOffloadKindAttrName()
+           << "' offload kind attribute";
+  if (attr.getValue() != expected)
+    return module.emitOpError()
+           << "expects '" << cir::CIRDialect::getOffloadKindAttrName()
+           << "' value '" << cir::stringifyOffloadKind(expected) << "'";
+  return success();
+}
+
+// A module marked with `cir.offload.container` holds the host module followed
+// by one or more device modules, each tagged with `cir.offload.kind`. Keeping
+// the host module first gives later offload passes a simple convention for
+// finding the host side while iterating the remaining device modules.
+static LogicalResult verifyOffloadContainer(mlir::Operation *op) {
+  auto container = mlir::dyn_cast<mlir::ModuleOp>(op);
+  if (!container)
+    return op->emitError() << "expects '"
+                           << cir::CIRDialect::getOffloadContainerAttrName()
+                           << "' attribute to be attached to '"
+                           << mlir::ModuleOp::getOperationName() << "'";
+
+  mlir::Block &body = *container.getBody();
+  if (body.empty())
+    return container.emitOpError()
+           << "expects host module as the first nested op";
+
+  auto host = mlir::dyn_cast<mlir::ModuleOp>(body.front());
+  if (!host)
+    return container.emitOpError()
+           << "expects host module as the first nested op";
+  if (failed(verifyOffloadKind(host, cir::OffloadKind::Host)))
+    return failure();
+
+  unsigned numDevices = 0;
+  auto it = body.begin();
+  for (++it; it != body.end(); ++it) {
+    auto module = mlir::dyn_cast<mlir::ModuleOp>(*it);
+    if (!module)
+      return container.emitOpError()
+             << "expects only nested builtin.module ops";
+    if (failed(verifyOffloadKind(module, cir::OffloadKind::Device)))
+      return failure();
+    ++numDevices;
+  }
+
+  if (numDevices == 0)
+    return container.emitOpError() << "expects at least one device module";
+  return success();
+}
+
+LogicalResult
+cir::CIRDialect::verifyOperationAttribute(mlir::Operation *op,
+                                          mlir::NamedAttribute attr) {
+  if (attr.getName() == getOffloadContainerAttrName()) {
+    if (!mlir::isa<mlir::UnitAttr>(attr.getValue()))
+      return op->emitError() << "expects '" << getOffloadContainerAttrName()
+                             << "' to be a unit attribute";
+    return verifyOffloadContainer(op);
+  }
+
+  // The container verifier owns the structural contract between a container
+  // and the modules it holds. All this can add is that the kind attribute
+  // never lands on something that is not a module.
+  if (attr.getName() == getOffloadKindAttrName() &&
+      !mlir::isa<mlir::ModuleOp>(op))
+    return op->emitError() << "expects '" << getOffloadKindAttrName()
+                           << "' attribute to be attached to '"
+                           << mlir::ModuleOp::getOperationName() << "'";
+
+  return success();
+}
+
 
//===----------------------------------------------------------------------===//
 // Helpers
 
//===----------------------------------------------------------------------===//
@@ -2335,66 +2441,6 @@ LogicalResult cir::VTTAddrPointOp::verify() {
   return success();
 }
 
-//===----------------------------------------------------------------------===//
-// OffloadContainerOp
-//===----------------------------------------------------------------------===//
-
-mlir::ModuleOp cir::OffloadContainerOp::getHostModule() {
-  return mlir::cast<mlir::ModuleOp>(getBody().front().front());
-}
-
-llvm::iterator_range<mlir::Block::op_iterator<mlir::ModuleOp>>
-cir::OffloadContainerOp::getDeviceModules() {
-  mlir::Block &body = getBody().front();
-  auto begin = body.op_begin<mlir::ModuleOp>();
-  auto end = body.op_end<mlir::ModuleOp>();
-  if (begin != end)
-    ++begin;
-  return {begin, end};
-}
-
-static LogicalResult checkOffloadKind(mlir::ModuleOp module,
-                                      cir::OffloadKind expected) {
-  auto attr = module->getAttrOfType<cir::OffloadKindAttr>(
-      cir::CIRDialect::getOffloadKindAttrName());
-  if (!attr)
-    return module.emitOpError()
-           << "expects '" << cir::CIRDialect::getOffloadKindAttrName()
-           << "' offload kind attribute";
-  if (attr.getValue() != expected)
-    return module.emitOpError()
-           << "expects '" << cir::CIRDialect::getOffloadKindAttrName()
-           << "' value '" << cir::stringifyOffloadKind(expected) << "'";
-  return success();
-}
-
-LogicalResult cir::OffloadContainerOp::verify() {
-  mlir::Block &body = getBody().front();
-  if (body.empty())
-    return emitOpError() << "expects host module as the first nested op";
-
-  auto host = mlir::dyn_cast<mlir::ModuleOp>(body.front());
-  if (!host)
-    return emitOpError() << "expects host module as the first nested op";
-  if (failed(checkOffloadKind(host, cir::OffloadKind::Host)))
-    return failure();
-
-  unsigned numDevices = 0;
-  auto it = body.begin();
-  for (++it; it != body.end(); ++it) {
-    auto module = mlir::dyn_cast<mlir::ModuleOp>(*it);
-    if (!module)
-      return emitOpError() << "expects only nested builtin.module ops";
-    if (failed(checkOffloadKind(module, cir::OffloadKind::Device)))
-      return failure();
-    ++numDevices;
-  }
-
-  if (numDevices == 0)
-    return emitOpError() << "expects at least one device module";
-  return success();
-}
-
 
//===----------------------------------------------------------------------===//
 // FuncOp
 
//===----------------------------------------------------------------------===//
diff --git a/clang/test/CIR/IR/invalid-offload-container.cir 
b/clang/test/CIR/IR/invalid-offload-container.cir
index 219442f39674a..92d4213765489 100644
--- a/clang/test/CIR/IR/invalid-offload-container.cir
+++ b/clang/test/CIR/IR/invalid-offload-container.cir
@@ -1,54 +1,80 @@
 // RUN: cir-opt %s -verify-diagnostics -split-input-file
 
-module {
-  cir.offload.container {
-    builtin.module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} { // expected-error {{expects 'cir.offload.kind' 
value 'host'}}
-    }
+module attributes {cir.offload.container} {
+  module @device_0 attributes {cir.offload.kind = #cir.offload_kind<device>} { 
// expected-error {{expects 'cir.offload.kind' value 'host'}}
   }
 }
 
 // -----
 
-module {
-  cir.offload.container {
-    builtin.module @host_0 attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
-    }
-    builtin.module @host_1 attributes {cir.offload.kind = 
#cir.offload_kind<host>} { // expected-error {{expects 'cir.offload.kind' value 
'device'}}
-    }
+module attributes {cir.offload.container} {
+  module @host_0 attributes {cir.offload.kind = #cir.offload_kind<host>} {
+  }
+  module @host_1 attributes {cir.offload.kind = #cir.offload_kind<host>} { // 
expected-error {{expects 'cir.offload.kind' value 'device'}}
   }
 }
 
 // -----
 
-module {
-  cir.offload.container {
-    builtin.module @host { // expected-error {{expects 'cir.offload.kind' 
offload kind attribute}}
-    }
+module attributes {cir.offload.container} {
+  module @host { // expected-error {{expects 'cir.offload.kind' offload kind 
attribute}}
+  }
+}
+
+// -----
+
+// expected-error@+1 {{expects only nested builtin.module ops}}
+module attributes {cir.offload.container} {
+  module @host attributes {cir.offload.kind = #cir.offload_kind<host>} {
+  }
+  cir.func @f() {
+    cir.return
+  }
+}
+
+// -----
+
+// expected-error@+1 {{expects at least one device module}}
+module attributes {cir.offload.container} {
+  module @host attributes {cir.offload.kind = #cir.offload_kind<host>} {
   }
 }
 
 // -----
 
 module {
-  cir.offload.container { // expected-error {{expects only nested 
builtin.module ops}}
-    builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
-    }
-    cir.const #cir.int<0> : !cir.int<s, 32>
+  // The op below keeps the CIR dialect loaded, which is what makes the
+  // dialect attribute verifier run on the empty container module.
+  cir.func @f() {
+    cir.return
+  }
+  // expected-error@+1 {{expects host module as the first nested op}}
+  module @container attributes {cir.offload.container} {
   }
 }
 
 // -----
 
 module {
-  cir.offload.container { // expected-error {{expects at least one device 
module}}
-    builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
-    }
+  cir.func @f() {
+    cir.return {cir.offload.container} // expected-error {{expects 
'cir.offload.container' attribute to be attached to 'builtin.module'}}
   }
 }
 
 // -----
 
 module {
-  cir.offload.container { // expected-error {{expects host module as the first 
nested op}}
+  cir.func @f() {
+    cir.return {cir.offload.kind = #cir.offload_kind<host>} // expected-error 
{{expects 'cir.offload.kind' attribute to be attached to 'builtin.module'}}
+  }
+}
+
+// -----
+
+// expected-error@+1 {{expects 'cir.offload.container' to be a unit attribute}}
+module attributes {cir.offload.container = 0 : i32} {
+  module @host attributes {cir.offload.kind = #cir.offload_kind<host>} {
+  }
+  module @device_0 attributes {cir.offload.kind = #cir.offload_kind<device>} {
   }
 }
diff --git a/clang/test/CIR/IR/offload-container.cir 
b/clang/test/CIR/IR/offload-container.cir
index b7b53aa973ce5..f1e00bc63fed7 100644
--- a/clang/test/CIR/IR/offload-container.cir
+++ b/clang/test/CIR/IR/offload-container.cir
@@ -1,32 +1,28 @@
 // RUN: cir-opt %s -split-input-file --verify-roundtrip | FileCheck %s
 
-module {
-  cir.offload.container {
-    builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
-    }
-    builtin.module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
-    }
+module attributes {cir.offload.container} {
+  module @host attributes {cir.offload.kind = #cir.offload_kind<host>} {
+  }
+  module @device_0 attributes {cir.offload.kind = #cir.offload_kind<device>} {
   }
 }
 
-// CHECK: cir.offload.container {
-// CHECK:   builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
-// CHECK:   builtin.module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
+// CHECK: module attributes {cir.offload.container} {
+// CHECK:   module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
+// CHECK:   module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
 
 // -----
 
-module {
-  cir.offload.container {
-    builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
-    }
-    builtin.module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
-    }
-    builtin.module @device_1 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
-    }
+module attributes {cir.offload.container} {
+  module @host attributes {cir.offload.kind = #cir.offload_kind<host>} {
+  }
+  module @device_0 attributes {cir.offload.kind = #cir.offload_kind<device>} {
+  }
+  module @device_1 attributes {cir.offload.kind = #cir.offload_kind<device>} {
   }
 }
 
-// CHECK: cir.offload.container {
-// CHECK:   builtin.module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
-// CHECK:   builtin.module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
-// CHECK:   builtin.module @device_1 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
+// CHECK: module attributes {cir.offload.container} {
+// CHECK:   module @host attributes {cir.offload.kind = 
#cir.offload_kind<host>} {
+// CHECK:   module @device_0 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {
+// CHECK:   module @device_1 attributes {cir.offload.kind = 
#cir.offload_kind<device>} {

_______________________________________________
cfe-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to