[CIR][CodeGen][LowerToLLVM] End-to-end implementation of offload_* cases for OpenCL with SPIR-V target (llvm#724)

seven-mile · smeenai · commit c0e8df696aca · 2024-10-09T13:52:15.000-07:00
This PR implements `offload_*` cases discussed in [this thread](https://discourse.llvm.org/t/rfc-clangir-unified-address-space-design-in-clangir/79728). * Integrate target-specific CIR-to-LLVM address space map into `TargetLoweringInfo` * CIRGen: Implement these cases in `getValueFromLangAS` * Lowering: Extend the state of type converter with `LowerModule` When frontend provides a new LangAS like `opencl_generic`, it would be processed by CIRGenTypes and `Builder.getPointerTo()` and encoded as `offload_generic`. When we lower CIR to LLVM, * For pointer types without address space attribute, it's mapped to `ptr addrspace(0)` directly * For target cases `target<x>`, it's mapped to `ptr addrspace(x)` * For other defined cases, query the target info with a new virtual method `getTargetAddrSpaceFromCIRAddrSpace`. General targets like X86 and ARM64 map all known cases to 0. For SPIR-V target here, it maps `offload_generic` to `addrspace(4)`.
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
@@ -566,10 +566,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/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,8 @@ class TargetLoweringInfo {
   virtual ~TargetLoweringInfo();
 
   const ABIInfo &getABIInfo() const { return *Info; }
+  virtual unsigned getTargetAddrSpaceFromCIRAddrSpace(
+      mlir::cir::AddressSpaceAttr addressSpaceAttr) const = 0;
 };
 
 } // namespace cir
diff --git a/clang/lib/CIR/Dialect/Transforms/TargetLowering/Targets/AArch64.cpp b/clang/lib/CIR/Dialect/Transforms/TargetLowering/Targets/AArch64.cpp
@@ -62,6 +62,21 @@ class AArch64TargetLoweringInfo : public TargetLoweringInfo {
       : TargetLoweringInfo(std::make_unique<AArch64ABIInfo>(LT, Kind)) {
     assert(!MissingFeature::swift());
   }
+
+  unsigned getTargetAddrSpaceFromCIRAddrSpace(
+      mlir::cir::AddressSpaceAttr addressSpaceAttr) const override {
+    using Kind = mlir::cir::AddressSpaceAttr::Kind;
+    switch (addressSpaceAttr.getValue()) {
+    case Kind::offload_private:
+    case Kind::offload_local:
+    case Kind::offload_global:
+    case Kind::offload_constant:
+    case Kind::offload_generic:
+      return 0;
+    default:
+      llvm_unreachable("Unknown CIR address space for this target");
+    }
+  }
 };
 
 } // namespace
diff --git a/clang/lib/CIR/Dialect/Transforms/TargetLowering/Targets/SPIR.cpp b/clang/lib/CIR/Dialect/Transforms/TargetLowering/Targets/SPIR.cpp
@@ -41,6 +41,25 @@ class SPIRVTargetLoweringInfo : public TargetLoweringInfo {
 public:
   SPIRVTargetLoweringInfo(LowerTypes &LT)
       : TargetLoweringInfo(std::make_unique<SPIRVABIInfo>(LT)) {}
+
+  unsigned getTargetAddrSpaceFromCIRAddrSpace(
+      mlir::cir::AddressSpaceAttr addressSpaceAttr) const override {
+    using Kind = mlir::cir::AddressSpaceAttr::Kind;
+    switch (addressSpaceAttr.getValue()) {
+    case Kind::offload_private:
+      return 0;
+    case Kind::offload_local:
+      return 3;
+    case Kind::offload_global:
+      return 1;
+    case Kind::offload_constant:
+      return 2;
+    case Kind::offload_generic:
+      return 4;
+    default:
+      llvm_unreachable("Unknown CIR address space for this target");
+    }
+  }
 };
 
 } // namespace
diff --git a/clang/lib/CIR/Dialect/Transforms/TargetLowering/Targets/X86.cpp b/clang/lib/CIR/Dialect/Transforms/TargetLowering/Targets/X86.cpp
@@ -105,6 +105,21 @@ class X86_64TargetLoweringInfo : public TargetLoweringInfo {
       : TargetLoweringInfo(std::make_unique<X86_64ABIInfo>(LM, AVXLevel)) {
     assert(!::cir::MissingFeatures::swift());
   }
+
+  unsigned getTargetAddrSpaceFromCIRAddrSpace(
+      mlir::cir::AddressSpaceAttr addressSpaceAttr) const override {
+    using Kind = mlir::cir::AddressSpaceAttr::Kind;
+    switch (addressSpaceAttr.getValue()) {
+    case Kind::offload_private:
+    case Kind::offload_local:
+    case Kind::offload_global:
+    case Kind::offload_constant:
+    case Kind::offload_generic:
+      return 0;
+    default:
+      llvm_unreachable("Unknown CIR address space for this target");
+    }
+  }
 };
 
 void X86_64ABIInfo::classify(Type Ty, uint64_t OffsetBase, Class &Lo, Class &Hi,
diff --git a/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp b/clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp
@@ -69,6 +69,8 @@
 #include <optional>
 #include <set>
 
+#include "LowerModule.h"
+
 using namespace cir;
 using namespace llvm;
 
@@ -3482,24 +3484,43 @@ void populateCIRToLLVMConversionPatterns(mlir::RewritePatternSet &patterns,
 }
 
 namespace {
+
+std::unique_ptr<mlir::cir::LowerModule>
+prepareLowerModule(mlir::ModuleOp module) {
+  mlir::PatternRewriter rewriter{module->getContext()};
+  // If the triple is not present, e.g. CIR modules parsed from text, we
+  // cannot init LowerModule properly.
+  assert(!::cir::MissingFeatures::makeTripleAlwaysPresent());
+  if (!module->hasAttr("cir.triple"))
+    return {};
+  return mlir::cir::createLowerModule(module, rewriter);
+}
+
+// FIXME: change the type of lowerModule to `LowerModule &` to have better
+// lambda capturing experience. Also blocked by makeTripleAlwaysPresent.
 void prepareTypeConverter(mlir::LLVMTypeConverter &converter,
-                          mlir::DataLayout &dataLayout) {
-  converter.addConversion([&](mlir::cir::PointerType type) -> mlir::Type {
+                          mlir::DataLayout &dataLayout,
+                          mlir::cir::LowerModule *lowerModule) {
+  converter.addConversion([&, lowerModule](
+                              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(lowerModule && "CIR AS map is not available");
+    // Pass through target addrspace and map CIR addrspace to LLVM addrspace by
+    // querying the target info.
+    unsigned targetAS =
+        addrSpace.isTarget()
+            ? addrSpace.getTargetValue()
+            : lowerModule->getTargetLoweringInfo()
+                  .getTargetAddrSpaceFromCIRAddrSpace(addrSpace);
 
-    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(),
@@ -3724,7 +3745,9 @@ void ConvertCIRToLLVMPass::runOnOperation() {
   auto module = getOperation();
   mlir::DataLayout dataLayout(module);
   mlir::LLVMTypeConverter converter(&getContext());
-  prepareTypeConverter(converter, dataLayout);
+  std::unique_ptr<mlir::cir::LowerModule> lowerModule =
+      prepareLowerModule(module);
+  prepareTypeConverter(converter, dataLayout, lowerModule.get());
 
   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,19 @@ 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 @test_lower_offload_as()
+  cir.func @test_lower_offload_as() {
+    %0 = cir.alloca !cir.ptr<!s32i, addrspace(offload_private)>, !cir.ptr<!cir.ptr<!s32i, addrspace(offload_private)>>, ["arg0", init] {alignment = 8 : i64}
+    // LLVM-NEXT: alloca ptr,
+    %1 = cir.alloca !cir.ptr<!s32i, addrspace(offload_global)>, !cir.ptr<!cir.ptr<!s32i, addrspace(offload_global)>>, ["arg1", init] {alignment = 8 : i64}
+    // LLVM-NEXT: alloca ptr addrspace(1),
+    %2 = cir.alloca !cir.ptr<!s32i, addrspace(offload_constant)>, !cir.ptr<!cir.ptr<!s32i, addrspace(offload_constant)>>, ["arg2", init] {alignment = 8 : i64}
+    // LLVM-NEXT: alloca ptr addrspace(2),
+    %3 = cir.alloca !cir.ptr<!s32i, addrspace(offload_local)>, !cir.ptr<!cir.ptr<!s32i, addrspace(offload_local)>>, ["arg3", init] {alignment = 8 : i64}
+    // LLVM-NEXT: alloca ptr addrspace(3),
+    %4 = cir.alloca !cir.ptr<!s32i, addrspace(offload_generic)>, !cir.ptr<!cir.ptr<!s32i, addrspace(offload_generic)>>, ["arg4", init] {alignment = 8 : i64}
+    // LLVM-NEXT: alloca ptr addrspace(4),
+    cir.return
+  }
 }
