[Flang] Set address space during FIR pointer-like types lowering

[skatrak]

This patch modifies FIR pointer-like types lowering to LLVM dialect to use the address space stated in the module's data layout.

[llvmbot]

@llvm/pr-subscribers-flang-fir-hlfir

@llvm/pr-subscribers-flang-codegen

Author: Sergio Afonso (skatrak)

Changes

This patch modifies FIR pointer-like types lowering to LLVM dialect to use the address space stated in the module's data layout.

---

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

4 Files Affected:

• (modified) flang/include/flang/Optimizer/CodeGen/TypeConverter.h (+2-1)
• (modified) flang/lib/Optimizer/CodeGen/TypeConverter.cpp (+7-1)
• (added) flang/test/Fir/alloca-addrspace-2.fir (+12)
• (added) flang/test/Fir/alloca-addrspace.fir (+12)

diff --git a/flang/include/flang/Optimizer/CodeGen/TypeConverter.h b/flang/include/flang/Optimizer/CodeGen/TypeConverter.h
index 29d0a902f556269..99323723d83eed1 100644
--- a/flang/include/flang/Optimizer/CodeGen/TypeConverter.h
+++ b/flang/include/flang/Optimizer/CodeGen/TypeConverter.h
@@ -119,7 +119,7 @@ class LLVMTypeConverter : public mlir::LLVMTypeConverter {
     if (eleTy.isa<fir::BaseBoxType>())
       return convertType(eleTy);
 
-    return mlir::LLVM::LLVMPointerType::get(convertType(eleTy));
+    return mlir::LLVM::LLVMPointerType::get(convertType(eleTy), addressSpace);
   }
 
   // convert a front-end kind value to either a std or LLVM IR dialect type
@@ -145,6 +145,7 @@ class LLVMTypeConverter : public mlir::LLVMTypeConverter {
   KindMapping kindMapping;
   std::unique_ptr<CodeGenSpecifics> specifics;
   std::unique_ptr<TBAABuilder> tbaaBuilder;
+  unsigned addressSpace;
 };
 
 } // namespace fir
diff --git a/flang/lib/Optimizer/CodeGen/TypeConverter.cpp b/flang/lib/Optimizer/CodeGen/TypeConverter.cpp
index 104018030bffd5c..99980e03a3f58fd 100644
--- a/flang/lib/Optimizer/CodeGen/TypeConverter.cpp
+++ b/flang/lib/Optimizer/CodeGen/TypeConverter.cpp
@@ -40,7 +40,13 @@ LLVMTypeConverter::LLVMTypeConverter(mlir::ModuleOp module, bool applyTBAA,
                                       getTargetTriple(module),
                                       getKindMapping(module))),
       tbaaBuilder(std::make_unique<TBAABuilder>(module->getContext(), applyTBAA,
-                                                forceUnifiedTBAATree)) {
+                                                forceUnifiedTBAATree)),
+      addressSpace(0) {
+  // Get default alloca address space for the current target
+  if (mlir::Attribute addrSpace =
+          mlir::DataLayout(module).getAllocaMemorySpace())
+    addressSpace = addrSpace.cast<mlir::IntegerAttr>().getUInt();
+
   LLVM_DEBUG(llvm::dbgs() << "FIR type converter\n");
 
   // Each conversion should return a value of type mlir::Type.
diff --git a/flang/test/Fir/alloca-addrspace-2.fir b/flang/test/Fir/alloca-addrspace-2.fir
new file mode 100644
index 000000000000000..8551cf8083635a4
--- /dev/null
+++ b/flang/test/Fir/alloca-addrspace-2.fir
@@ -0,0 +1,12 @@
+// RUN: fir-opt --fir-to-llvm-ir %s | FileCheck %s
+// RUN: tco --fir-to-llvm-ir %s | FileCheck %s
+
+module attributes { dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memory_space", 5 : ui32>> } {
+  // CHECK-LABEL: llvm.func @set_addrspace
+  func.func @set_addrspace() {
+    // CHECK: llvm.alloca {{.*}} x i32
+    // CHECK-SAME: -> !llvm.ptr<i32, 5>
+    %0 = fir.alloca i32
+    return
+  }
+}
diff --git a/flang/test/Fir/alloca-addrspace.fir b/flang/test/Fir/alloca-addrspace.fir
new file mode 100644
index 000000000000000..20bf59b7a568d5f
--- /dev/null
+++ b/flang/test/Fir/alloca-addrspace.fir
@@ -0,0 +1,12 @@
+// RUN: fir-opt --fir-to-llvm-ir %s | FileCheck %s
+// RUN: tco --fir-to-llvm-ir %s | FileCheck %s
+
+module {
+  // CHECK-LABEL: llvm.func @default_addrspace
+  func.func @default_addrspace() {
+    // CHECK: llvm.alloca {{.*}} x i32
+    // CHECK-SAME: -> !llvm.ptr<i32>
+    %0 = fir.alloca i32
+    return
+  }
+}

[jeanPerier]

Not an address space expert, but looks reasonable to me. However, this is far from the only place where LLVMPointerType::get is called when lowering FIR. I just made a patch to update our codegen to use opaque pointers #69692, and there are quite a few other places? Don't you need to do it in every places?

Although, wouldn't it be more generic to do this during the LLVM Dialect to LLVM IR translation to cover all the passes/translations/dialects that may generate llvm.ptr type inside a module with a non default addressing space?

[skatrak]

Thank you @jeanPerier for looking into this and for your suggestions. I focused on the particular issue of fir.alloca lowering, so I didn't realize about all these other spots in FIR lowering where LLVMPointerTypes can be instantiated.

I like the idea of making the LLVM dialect responsible for this, but I think it currently can't be done without making changes to the LLVMPointerType. The issue is that its addressSpace parameter is default-valued to 0, so after creation there would be no way of telling whether a pointer was meant to be specifically for address space 0 or if it was supposed to take the data layout's default.

I'm not that familiar with the specifics of address spaces either, but if my concern is right that would mean that LLVM dialect's users must specify the address space of any LLVMPointerType instances that are not supposed to go to address space 0, and handle the data layout themselves. I think that the alternative solution of moving the responsibility of checking for the data layout to the LLVM dialect would mean allowing the address space to be optional, so that during translation to LLVM IR we could check for the data layout before falling back to 0 as a last option.

I'm not sure what the right approach here would be, so I'd be interested to hear some opinions.

[skatrak]

Just updated the patch with the addition of the data layout-based default address space to all places where LLVMPointerTypes are created in Flang. I'm avoiding making any changes to the LLVM dialect at least for now, so that there is at least a complete implementation we can discuss.

[agozillon]

Perhaps a good idea to rename this to allocaAddressSpace, rather than addresSspace, as there's a few different types of address space that can be stored in the data layout!

The getAddressSpace function could also be renamed to getAllocaAddressSpace, but you will probably get the variations I have in the other address space related PR we've made soon on a rebase (hopefully it'll land without any issue, if not can always just take them directly from the PR).