LLVM::GlobalOp: take address space as builder argument

Accept the address space of the global as a builder argument when constructing
an LLVM::GlobalOp instance. This decreases the reliance of LLVM::GlobalOp users
on the internal name of the attribute used for this purpose. Update several
uses of the address space in GPU to NVVM conversion.

PiperOrigin-RevId: 284233254

3 files changed, 8 insertions, 9 deletions

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index 66d9ba3f750..8aa5397651c 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -553,7 +553,8 @@ def LLVM_GlobalOp
   let builders = [
     OpBuilder<"Builder *builder, OperationState &result, LLVMType type, "
               "bool isConstant, Linkage linkage, StringRef name, "
-              "Attribute value, ArrayRef<NamedAttribute> attrs = {}">
+              "Attribute value, unsigned addrSpace = 0, "
+              "ArrayRef<NamedAttribute> attrs = {}">
   ];
 
   let extraClassDeclaration = [{
diff --git a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
index 50a2e2efd2c..949294cf692 100644
--- a/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
+++ b/mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
@@ -385,12 +385,10 @@ private:
 
     auto arrayType = LLVM::LLVMType::getArrayTy(elementType, numElements);
     StringRef name = "reduce_buffer";
-    auto addrSpace =
-        builder.getNamedAttr("addr_space", builder.getI32IntegerAttr(3));
     auto globalOp = builder.create<LLVM::GlobalOp>(
         loc, arrayType.cast<LLVM::LLVMType>(),
         /*isConstant=*/false, LLVM::Linkage::Internal, name,
-        /*value=*/Attribute(), llvm::makeArrayRef(addrSpace));
+        /*value=*/Attribute(), gpu::GPUDialect::getWorkgroupAddressSpace());
 
     return rewriter.create<LLVM::AddressOfOp>(loc, globalOp);
   }
@@ -481,15 +479,12 @@ struct FuncOpLowering : LLVMOpLowering {
       auto elementType =
           lowering.convertType(type.getElementType()).cast<LLVM::LLVMType>();
       auto arrayType = LLVM::LLVMType::getArrayTy(elementType, numElements);
-      auto addSpaceAttr = rewriter.getNamedAttr(
-          "addr_space", rewriter.getI32IntegerAttr(
-                            gpu::GPUDialect::getWorkgroupAddressSpace()));
       std::string name =
           llvm::formatv("__wg_{0}_{1}", gpuFuncOp.getName(), en.index());
       auto globalOp = rewriter.create<LLVM::GlobalOp>(
           gpuFuncOp.getLoc(), arrayType, /*isConstant=*/false,
           LLVM::Linkage::Internal, name, /*value=*/Attribute(),
-          llvm::makeArrayRef(addSpaceAttr));
+          gpu::GPUDialect::getWorkgroupAddressSpace());
       workgroupBuffers.push_back(globalOp);
     }
 
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index fb4555674eb..78da9998c6d 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -869,7 +869,8 @@ static StringRef getLinkageAttrName() { return "linkage"; }
 
 void GlobalOp::build(Builder *builder, OperationState &result, LLVMType type,
                      bool isConstant, Linkage linkage, StringRef name,
-                     Attribute value, ArrayRef<NamedAttribute> attrs) {
+                     Attribute value, unsigned addrSpace,
+                     ArrayRef<NamedAttribute> attrs) {
   result.addAttribute(SymbolTable::getSymbolAttrName(),
                       builder->getStringAttr(name));
   result.addAttribute("type", TypeAttr::get(type));
@@ -879,6 +880,8 @@ void GlobalOp::build(Builder *builder, OperationState &result, LLVMType type,
     result.addAttribute("value", value);
   result.addAttribute(getLinkageAttrName(), builder->getI64IntegerAttr(
                                                 static_cast<int64_t>(linkage)));
+  if (addrSpace != 0)
+    result.addAttribute("addr_space", builder->getI32IntegerAttr(addrSpace));
   result.attributes.append(attrs.begin(), attrs.end());
   result.addRegion();
 }