[Mlir-commits] [mlir] [mlir][NVVM] Add managed attribute for global variables (PR #189751)

[Zhen Wang]

https://github.com/wangzpgi updated https://github.com/llvm/llvm-project/pull/189751

>From 03d66dd315e6f58ebd70b00937a1b4c182aa8b7b Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Tue, 31 Mar 2026 13:56:54 -0700
Subject: [PATCH 1/2] [mlir][NVVM] Add managed attribute for global variables

Add support for the `nvvm.managed` attribute on `llvm.mlir.global` ops.
When present, the LLVM IR translation emits `!nvvm.annotations` metadata
with `!"managed"` for the global variable, which the NVPTX backend uses
to generate `.attribute(.managed)` in PTX output.

This enables CUDA managed memory support for frontends that lower
through MLIR.
---
 mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td   |  3 +++
 .../Dialect/NVVM/NVVMToLLVMIRTranslation.cpp  | 19 ++++++++++++++++++-
 mlir/test/Target/LLVMIR/nvvmir.mlir           |  7 +++++++
 3 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
index 0c5dae265e2ca..87fd75f5a3e19 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td
@@ -96,6 +96,9 @@ def NVVM_Dialect : Dialect {
     /// nvvm.cluster_dim attributes.
     static StringRef getBlocksAreClustersAttrName() { return "nvvm.blocksareclusters"; }
 
+    /// Get the name of the attribute used to annotate managed global variables.
+    static StringRef getManagedAttrName() { return "nvvm.managed"; }
+
     /// Verify an attribute from this dialect on the argument at 'argIndex' for
     /// the region at 'regionIndex' on the given operation. Returns failure if
     /// the verification failed, success otherwise. This hook may optionally be
diff --git a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
index 254be8e63590f..08e44949de1c0 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp
@@ -705,11 +705,28 @@ class NVVMDialectLLVMIRTranslationInterface
     return failure();
   }
 
-  /// Attaches module-level metadata for functions marked as kernels.
+  /// Attaches module-level metadata for functions marked as kernels
+  /// and managed annotations for global variables.
   LogicalResult
   amendOperation(Operation *op, ArrayRef<llvm::Instruction *> instructions,
                  NamedAttribute attribute,
                  LLVM::ModuleTranslation &moduleTranslation) const final {
+    if (auto globalOp = dyn_cast<LLVM::GlobalOp>(op)) {
+      if (attribute.getName() == NVVM::NVVMDialect::getManagedAttrName()) {
+        auto *gv = cast<llvm::GlobalVariable>(
+            moduleTranslation.lookupGlobal(globalOp));
+        llvm::Module *m = gv->getParent();
+        llvm::LLVMContext &ctx = m->getContext();
+        llvm::NamedMDNode *md = m->getOrInsertNamedMetadata("nvvm.annotations");
+        md->addOperand(llvm::MDNode::get(
+            ctx, {llvm::ConstantAsMetadata::get(gv),
+                  llvm::MDString::get(ctx, "managed"),
+                  llvm::ConstantAsMetadata::get(llvm::ConstantInt::get(
+                      llvm::Type::getInt32Ty(ctx), 1))}));
+      }
+      return success();
+    }
+
     auto func = dyn_cast<LLVM::LLVMFuncOp>(op);
     if (!func)
       return failure();
diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index 8a7e9bae4ec2e..f6360be2dfcc5 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -915,3 +915,10 @@ llvm.func @nanosleep(%duration: i32) {
   nvvm.nanosleep %duration
   llvm.return
 }
+
+// -----
+
+// CHECK: @managed_g = addrspace(1) global i32 0
+// CHECK: !nvvm.annotations = !{![[MANAGED:[0-9]+]]}
+// CHECK: ![[MANAGED]] = !{ptr addrspace(1) @managed_g, !"managed", i32 1}
+llvm.mlir.global external @managed_g(0 : i32) {addr_space = 1 : i32, nvvm.managed} : i32

>From 80f5b7277996650a0335a999bec6b2d8db0e01e4 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Tue, 31 Mar 2026 14:05:18 -0700
Subject: [PATCH 2/2] add test without an initializer

---
 mlir/test/Target/LLVMIR/nvvmir.mlir | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/mlir/test/Target/LLVMIR/nvvmir.mlir b/mlir/test/Target/LLVMIR/nvvmir.mlir
index f6360be2dfcc5..64c44be3e8182 100644
--- a/mlir/test/Target/LLVMIR/nvvmir.mlir
+++ b/mlir/test/Target/LLVMIR/nvvmir.mlir
@@ -919,6 +919,9 @@ llvm.func @nanosleep(%duration: i32) {
 // -----
 
 // CHECK: @managed_g = addrspace(1) global i32 0
-// CHECK: !nvvm.annotations = !{![[MANAGED:[0-9]+]]}
-// CHECK: ![[MANAGED]] = !{ptr addrspace(1) @managed_g, !"managed", i32 1}
+// CHECK: @managed_decl = external addrspace(1) global i32
+// CHECK: !nvvm.annotations = !{![[MANAGED0:[0-9]+]], ![[MANAGED1:[0-9]+]]}
+// CHECK: ![[MANAGED0]] = !{ptr addrspace(1) @managed_g, !"managed", i32 1}
+// CHECK: ![[MANAGED1]] = !{ptr addrspace(1) @managed_decl, !"managed", i32 1}
 llvm.mlir.global external @managed_g(0 : i32) {addr_space = 1 : i32, nvvm.managed} : i32
+llvm.mlir.global external @managed_decl() {addr_space = 1 : i32, nvvm.managed} : i32