[flang-commits] [flang] [Flang] Set address space during FIR pointer-like types lowering (PR #69599)

[Sergio Afonso via flang-commits]

https://github.com/skatrak created https://github.com/llvm/llvm-project/pull/69599

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

>From 6e868a3ef8c81ffd54e9f99a43bf45eee8ffa04d Mon Sep 17 00:00:00 2001
From: Sergio Afonso <safonsof at amd.com>
Date: Thu, 19 Oct 2023 12:49:35 +0100
Subject: [PATCH] [Flang] Set address space during FIR pointer-like types
 lowering

This patch modifies FIR pointer-like types lowering to LLVM dialect to use the
address space stated in the module's data layout.
---
 .../include/flang/Optimizer/CodeGen/TypeConverter.h  |  3 ++-
 flang/lib/Optimizer/CodeGen/TypeConverter.cpp        |  8 +++++++-
 flang/test/Fir/alloca-addrspace-2.fir                | 12 ++++++++++++
 flang/test/Fir/alloca-addrspace.fir                  | 12 ++++++++++++
 4 files changed, 33 insertions(+), 2 deletions(-)
 create mode 100644 flang/test/Fir/alloca-addrspace-2.fir
 create mode 100644 flang/test/Fir/alloca-addrspace.fir

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
+  }
+}