[CIR][Dialect] Add offload_* cases in address space attribute

seven-mile · seven-mile · commit 6f1be9518ea3 · 2024-07-18T09:22:34.000+08:00
fix dirty include
diff --git a/clang/include/clang/CIR/Dialect/IR/CIRAttrs.td b/clang/include/clang/CIR/Dialect/IR/CIRAttrs.td
@@ -707,6 +707,7 @@ def AddressSpaceAttr : CIR_Attr<"AddressSpace", "addrspace"> {
   let extraClassDeclaration = [{
     static constexpr char kTargetKeyword[] = "}]#targetASCase.symbol#[{";
     static constexpr int32_t kFirstTargetASValue = }]#targetASCase.value#[{;
+    using MapTy = unsigned[kFirstTargetASValue];
   
     bool isLang() const;
     bool isTarget() const;
diff --git a/clang/include/clang/CIR/MissingFeatures.h b/clang/include/clang/CIR/MissingFeatures.h
@@ -161,7 +161,6 @@ struct MissingFeatures {
   static bool constantFoldsToSimpleInteger() { return false; }
   static bool checkFunctionCallABI() { return false; }
   static bool zeroInitializer() { return false; }
-  static bool targetLoweringInfoAddressSpaceMap() { return false; }
   static bool targetCodeGenInfoIsProtoCallVariadic() { return false; }
   static bool targetCodeGenInfoGetNullPointer() { return false; }
   static bool operandBundles() { return false; }
@@ -277,6 +276,10 @@ struct MissingFeatures {
   static bool returnValueDominatingStoreOptmiization() { return false; }
   // Globals (vars and functions) may have attributes that are target depedent.
   static bool setTargetAttributes() { return false; }
+
+  // CIR modules parsed from text form may not carry the triple or data layout
+  // specs. We should make it always present.
+  static bool makeTripleAlwaysPresent() { return false; }
 };
 
 } // namespace cir
diff --git a/clang/lib/CIR/Dialect/IR/CIRAttrs.cpp b/clang/lib/CIR/Dialect/IR/CIRAttrs.cpp
@@ -565,10 +565,16 @@ AddressSpaceAttr::getValueFromLangAS(clang::LangAS langAS) {
     // Default address space should be encoded as a null attribute.
     return std::nullopt;
   case LangAS::opencl_global:
+    return Kind::offload_global;
   case LangAS::opencl_local:
+    return Kind::offload_local;
   case LangAS::opencl_constant:
+    return Kind::offload_constant;
   case LangAS::opencl_private:
+    return Kind::offload_private;
   case LangAS::opencl_generic:
+    return Kind::offload_generic;
+
   case LangAS::opencl_global_device:
   case LangAS::opencl_global_host:
   case LangAS::cuda_device:
diff --git a/clang/lib/CIR/Dialect/Transforms/TargetLowering/TargetInfo.cpp b/clang/lib/CIR/Dialect/Transforms/TargetLowering/TargetInfo.cpp
@@ -8,5 +8,10 @@ TargetLoweringInfo::TargetLoweringInfo(std::unique_ptr<ABIInfo> Info)
 
 TargetLoweringInfo::~TargetLoweringInfo() = default;
 
+AddressSpaceAttr::MapTy const &TargetLoweringInfo::getCIRAddrSpaceMap() const {
+  static AddressSpaceAttr::MapTy defaultAddrSpaceMap = {0};
+  return defaultAddrSpaceMap;
+}
+
 } // namespace cir
 } // namespace mlir
diff --git a/clang/lib/CIR/Dialect/Transforms/TargetLowering/TargetLoweringInfo.h b/clang/lib/CIR/Dialect/Transforms/TargetLowering/TargetLoweringInfo.h
@@ -18,6 +18,8 @@
 #include "ABIInfo.h"
 #include <memory>
 
+#include "clang/CIR/Dialect/IR/CIRAttrs.h"
+
 namespace mlir {
 namespace cir {
 
@@ -30,6 +32,7 @@ class TargetLoweringInfo {
   virtual ~TargetLoweringInfo();
 
   const ABIInfo &getABIInfo() const { return *Info; }
+  virtual AddressSpaceAttr::MapTy const &getCIRAddrSpaceMap() const;
 };
 
 } // namespace cir
diff --git a/clang/lib/CIR/Dialect/Transforms/TargetLowering/Targets/SPIR.cpp b/clang/lib/CIR/Dialect/Transforms/TargetLowering/Targets/SPIR.cpp
@@ -12,6 +12,7 @@
 #include "TargetInfo.h"
 #include "TargetLoweringInfo.h"
 #include "clang/CIR/ABIArgInfo.h"
+#include "clang/CIR/Dialect/IR/CIRTypes.h"
 #include "clang/CIR/MissingFeatures.h"
 #include "llvm/Support/ErrorHandling.h"
 
@@ -37,10 +38,27 @@ class SPIRVABIInfo : public ABIInfo {
   }
 };
 
+//===----------------------------------------------------------------------===//
+// SPIR-V other target-specific information
+//===----------------------------------------------------------------------===//
+
+static AddressSpaceAttr::MapTy SPIRVAddrSpaceMap = {
+    0, // None
+    0, // offload_private
+    3, // offload_local
+    1, // offload_global
+    2, // offload_constant
+    4, // offload_generic
+};
+
 class SPIRVTargetLoweringInfo : public TargetLoweringInfo {
 public:
   SPIRVTargetLoweringInfo(LowerTypes &LT)
       : TargetLoweringInfo(std::make_unique<SPIRVABIInfo>(LT)) {}
+
+  AddressSpaceAttr::MapTy const &getCIRAddrSpaceMap() const override {
+    return SPIRVAddrSpaceMap;
+  }
 };
 
 } // namespace
diff --git a/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp b/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp
@@ -68,6 +68,8 @@
 #include <optional>
 #include <set>
 
+#include "LowerModule.h"
+
 using namespace cir;
 using namespace llvm;
 
@@ -3583,24 +3585,44 @@ void populateCIRToLLVMConversionPatterns(mlir::RewritePatternSet &patterns,
 }
 
 namespace {
+
+using IRASMap = mlir::cir::AddressSpaceAttr::MapTy;
+
+/// Returns the CIR address space map when available. Returns nullptr if the
+/// triple is not present in the module.
+IRASMap const *getIRAddrSpaceMap(mlir::ModuleOp module) {
+  // If the triple is not present, e.g. CIR modules parsed from text, we
+  // cannot init LowerModule properly.
+  assert(!::cir::MissingFeatures::makeTripleAlwaysPresent());
+  // Here we will use a default AS map to ignore all AS stuff.
+  if (!module->hasAttr("cir.triple"))
+    return nullptr;
+  mlir::PatternRewriter rewriter{module->getContext()};
+  auto lowerModule = mlir::cir::createLowerModule(module, rewriter);
+  // All CIR AS maps have static lifetime, it's safe to ref it
+  return &lowerModule.getTargetLoweringInfo().getCIRAddrSpaceMap();
+};
+
 void prepareTypeConverter(mlir::LLVMTypeConverter &converter,
-                          mlir::DataLayout &dataLayout) {
-  converter.addConversion([&](mlir::cir::PointerType type) -> mlir::Type {
+                          mlir::DataLayout &dataLayout,
+                          IRASMap const *cirASMap) {
+  converter.addConversion([&, cirASMap](
+                              mlir::cir::PointerType type) -> mlir::Type {
     // Drop pointee type since LLVM dialect only allows opaque pointers.
 
     auto addrSpace =
         mlir::cast_if_present<mlir::cir::AddressSpaceAttr>(type.getAddrSpace());
-    // null addrspace attribute indicates the default addrspace
+    // Null addrspace attribute indicates the default addrspace
     if (!addrSpace)
       return mlir::LLVM::LLVMPointerType::get(type.getContext());
 
-    // TODO(cir): Query the target-specific address space map to lower other ASs
-    // like `opencl_private`.
-    assert(!MissingFeatures::targetLoweringInfoAddressSpaceMap());
-    assert(addrSpace.isTarget() && "NYI");
+    assert(cirASMap && "CIR AS map is not available");
+    // Pass through target addrspace and map CIR addrspace to LLVM addrspace.
+    unsigned targetAS = addrSpace.isTarget()
+                            ? addrSpace.getTargetValue()
+                            : (*cirASMap)[addrSpace.getValue()];
 
-    return mlir::LLVM::LLVMPointerType::get(type.getContext(),
-                                            addrSpace.getTargetValue());
+    return mlir::LLVM::LLVMPointerType::get(type.getContext(), targetAS);
   });
   converter.addConversion([&](mlir::cir::DataMemberType type) -> mlir::Type {
     return mlir::IntegerType::get(type.getContext(),
@@ -3825,7 +3847,7 @@ void ConvertCIRToLLVMPass::runOnOperation() {
   auto module = getOperation();
   mlir::DataLayout dataLayout(module);
   mlir::LLVMTypeConverter converter(&getContext());
-  prepareTypeConverter(converter, dataLayout);
+  prepareTypeConverter(converter, dataLayout, getIRAddrSpaceMap(module));
 
   mlir::RewritePatternSet patterns(&getContext());
 
diff --git a/clang/test/CIR/CodeGen/OpenCL/addrspace-alloca.cl b/clang/test/CIR/CodeGen/OpenCL/addrspace-alloca.cl
@@ -3,25 +3,23 @@
 // RUN: %clang_cc1 -cl-std=CL3.0 -O0 -fclangir -emit-llvm -triple spirv64-unknown-unknown %s -o %t.ll
 // RUN: FileCheck --input-file=%t.ll %s --check-prefix=LLVM
 
-// Lowering of language-specific AS not supported
-// XFAIL: *
 
-// CIR: cir.func @func(%arg0: !cir.ptr<!s32i, addrspace(target<3>)>
+// CIR: cir.func @func(%arg0: !cir.ptr<!s32i, addrspace(offload_local)>
 // LLVM: @func(ptr addrspace(3)
 kernel void func(local int *p) {
-  // CIR-NEXT: %[[#ALLOCA_P:]] = cir.alloca !cir.ptr<!s32i, addrspace(target<3>)>, !cir.ptr<!cir.ptr<!s32i, addrspace(target<3>)>>, ["p", init] {alignment = 8 : i64}
+  // CIR-NEXT: %[[#ALLOCA_P:]] = cir.alloca !cir.ptr<!s32i, addrspace(offload_local)>, !cir.ptr<!cir.ptr<!s32i, addrspace(offload_local)>>, ["p", init] {alignment = 8 : i64}
   // LLVM-NEXT: %[[#ALLOCA_P:]] = alloca ptr addrspace(3), i64 1, align 8
 
   int x;
   // CIR-NEXT: %[[#ALLOCA_X:]] = cir.alloca !s32i, !cir.ptr<!s32i>, ["x"] {alignment = 4 : i64}
   // LLVM-NEXT: %[[#ALLOCA_X:]] = alloca i32, i64 1, align 4
 
   global char *b;
-  // CIR-NEXT: %[[#ALLOCA_B:]] = cir.alloca !cir.ptr<!s8i, addrspace(target<1>)>, !cir.ptr<!cir.ptr<!s8i, addrspace(target<1>)>>, ["b"] {alignment = 8 : i64}
+  // CIR-NEXT: %[[#ALLOCA_B:]] = cir.alloca !cir.ptr<!s8i, addrspace(offload_global)>, !cir.ptr<!cir.ptr<!s8i, addrspace(offload_global)>>, ["b"] {alignment = 8 : i64}
   // LLVM-NEXT: %[[#ALLOCA_B:]] = alloca ptr addrspace(1), i64 1, align 8
 
   // Store of the argument `p`
-  // CIR-NEXT: cir.store %arg0, %[[#ALLOCA_P]] : !cir.ptr<!s32i, addrspace(target<3>)>, !cir.ptr<!cir.ptr<!s32i, addrspace(target<3>)>>
+  // CIR-NEXT: cir.store %arg0, %[[#ALLOCA_P]] : !cir.ptr<!s32i, addrspace(offload_local)>, !cir.ptr<!cir.ptr<!s32i, addrspace(offload_local)>>
   // LLVM-NEXT: store ptr addrspace(3) %{{[0-9]+}}, ptr %[[#ALLOCA_P]], align 8
 
   return;
diff --git a/clang/test/CIR/CodeGen/OpenCL/spirv-target.cl b/clang/test/CIR/CodeGen/OpenCL/spirv-target.cl
@@ -4,8 +4,6 @@
 // RUN: %clang_cc1 -cl-std=CL3.0 -fclangir -emit-llvm -triple spirv64-unknown-unknown %s -o %t_64.ll
 // RUN: FileCheck --input-file=%t_64.ll %s --check-prefix=LLVM-SPIRV64
 
-// Lowering of language-specific AS not supported
-// XFAIL: *
 
 // CIR-SPIRV64: cir.triple = "spirv64-unknown-unknown"
 // LLVM-SPIRV64: target triple = "spirv64-unknown-unknown"
diff --git a/clang/test/CIR/Lowering/address-space.cir b/clang/test/CIR/Lowering/address-space.cir
@@ -3,7 +3,10 @@
 
 !s32i = !cir.int<s, 32>
 
-module  {
+module attributes {
+  cir.triple = "spirv64-unknown-unknown",
+  llvm.data_layout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-G1"
+} {
   // LLVM: define void @foo(ptr %0)
   cir.func @foo(%arg0: !cir.ptr<!s32i>) {
     // LLVM-NEXT: alloca ptr,
@@ -24,4 +27,11 @@ module  {
     %0 = cir.alloca !cir.ptr<!s32i, addrspace(target<0>)>, !cir.ptr<!cir.ptr<!s32i, addrspace(target<0>)>>, ["arg", init] {alignment = 8 : i64}
     cir.return
   }
+
+  // LLVM: define void @qux(ptr addrspace(4) %0)
+  cir.func @qux(%arg0: !cir.ptr<!s32i, addrspace(offload_generic)>) {
+    // LLVM-NEXT: alloca ptr addrspace(4),
+    %0 = cir.alloca !cir.ptr<!s32i, addrspace(offload_generic)>, !cir.ptr<!cir.ptr<!s32i, addrspace(offload_generic)>>, ["arg", init] {alignment = 8 : i64}
+    cir.return
+  }
 }
