[CIR][CodeGen][Lowering] Better handling of alloca address space with unified AS (llvm#682)

`TargetCodeGenInfo::getASTAllocaAddressSpace` is a bad design because it
requires the targets return `LangAS`, which enforce the targets to
consider which languages could be their upstream and unnecessarily
complicate the target info. Unified AS in CIR is a better abstraction
level for this kind of target info.

Apart from that, the languages also has some requirements on the result
address space of alloca.

```cpp
void func() {
  int x;
  // Here, the AS of pointer `&x` is the alloca AS defined by the language.
}
```

When we have inconsistency between the alloca AS defined by the language
and the one from target info, we have to perform `addrspacecast` from
the target one to the language one.

This PR includes
* New method `CGM.getLangTempAllocaAddressSpace` which explicitly
specifies the alloca address space defined by languages. It replaces the
vague `LangAS::Default` in the AS comparisons from OG CodeGen.
* Replace `getASTAllocaAddressSpace` with `getCIRAllocaAddressSpace`,
which returns CIR unified AS.
* Also use `getCIRAllocaAddressSpace` as the result address space of
`buildAlloca`.
* Fix the lowering of `cir.alloca` operation to be address-space-aware.

We don't perform any `addrspacecast` in this PR, i.e. all the related
code paths still remain NYI and it's fine. That's because we don't even
have a `(language, target)` pair holding the inconsistency.

After these changes, in the previous OpenCL testcases we will see all
the alloca pointers turning into private AS, without any
`addrspacecast`.

Author: seven-mile

clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp

@@ -1100,7 +1100,7 @@ RValue CIRGenFunction::buildBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID,
     // the AST level this is handled within CreateTempAlloca et al., but for the
     // builtin / dynamic alloca we have to handle it here.
     assert(!MissingFeatures::addressSpace());
-    auto AAS = builder.getAddrSpaceAttr(getASTAllocaAddressSpace());
+    auto AAS = getCIRAllocaAddressSpace();
     auto EAS = builder.getAddrSpaceAttr(
         E->getType()->getPointeeType().getAddressSpace());
     if (EAS != AAS) {

clang/lib/CIR/CodeGen/CIRGenExpr.cpp

@@ -2752,7 +2752,10 @@ mlir::Value CIRGenFunction::buildAlloca(StringRef name, mlir::Type ty,
                                         mlir::Location loc, CharUnits alignment,
                                         mlir::OpBuilder::InsertPoint ip,
                                         mlir::Value arraySize) {
-  auto localVarPtrTy = mlir::cir::PointerType::get(builder.getContext(), ty);
+  // CIR uses its own alloca AS rather than follow the target data layout like
+  // original CodeGen. The data layout awareness should be done in the lowering
+  // pass instead.
+  auto localVarPtrTy = builder.getPointerTo(ty, getCIRAllocaAddressSpace());
   auto alignIntAttr = CGM.getSize(alignment);
 
   mlir::Value addr;
@@ -2977,7 +2980,9 @@ Address CIRGenFunction::CreateTempAlloca(mlir::Type Ty, CharUnits Align,
   // be different from the type defined by the language. For example,
   // in C++ the auto variables are in the default address space. Therefore
   // cast alloca to the default address space when necessary.
-  if (getASTAllocaAddressSpace() != LangAS::Default) {
+  if (auto ASTAS =
+          builder.getAddrSpaceAttr(CGM.getLangTempAllocaAddressSpace());
+      getCIRAllocaAddressSpace() != ASTAS) {
     llvm_unreachable("Requires address space cast which is NYI");
   }
   return Address(V, Ty, Align);

clang/lib/CIR/CodeGen/CIRGenModule.cpp

@@ -159,7 +159,7 @@ CIRGenModule::CIRGenModule(mlir::MLIRContext &context,
   AllocaInt8PtrTy = UInt8PtrTy;
   // TODO: GlobalsInt8PtrTy
   // TODO: ConstGlobalsPtrTy
-  ASTAllocaAddressSpace = getTargetCIRGenInfo().getASTAllocaAddressSpace();
+  CIRAllocaAddressSpace = getTargetCIRGenInfo().getCIRAllocaAddressSpace();
 
   PtrDiffTy = ::mlir::cir::IntType::get(
       builder.getContext(), astCtx.getTargetInfo().getMaxPointerWidth(),
@@ -1404,6 +1404,16 @@ LangAS CIRGenModule::getGlobalConstantAddressSpace() const {
   return LangAS::Default;
 }
 
+// TODO(cir): this could be a common AST helper for both CIR and LLVM codegen.
+LangAS CIRGenModule::getLangTempAllocaAddressSpace() const {
+  if (getLangOpts().OpenCL)
+    return LangAS::opencl_private;
+  if (getLangOpts().SYCLIsDevice || getLangOpts().CUDAIsDevice ||
+    (getLangOpts().OpenMP && getLangOpts().OpenMPIsTargetDevice))
+    llvm_unreachable("NYI");
+  return LangAS::Default;
+}
+
 static mlir::cir::GlobalOp
 generateStringLiteral(mlir::Location loc, mlir::TypedAttr C,
                       mlir::cir::GlobalLinkageKind LT, CIRGenModule &CGM,

clang/lib/CIR/CodeGen/CIRGenModule.h

@@ -335,6 +335,12 @@ class CIRGenModule : public CIRGenTypeCache {
   /// in AST is always in default address space.
   LangAS getGlobalConstantAddressSpace() const;
 
+  /// Returns the address space for temporary allocations in the language. This
+  /// ensures that the allocated variable's address space matches the
+  /// expectations of the AST, rather than using the target's allocation address
+  /// space, which may lead to type mismatches in other parts of the IR.
+  LangAS getLangTempAllocaAddressSpace() const;
+
   /// Set attributes which are common to any form of a global definition (alias,
   /// Objective-C method, function, global variable).
   ///

clang/lib/CIR/CodeGen/CIRGenTypeCache.h

@@ -18,6 +18,7 @@
 #include "clang/AST/CharUnits.h"
 #include "clang/Basic/AddressSpaces.h"
 #include "clang/CIR/Dialect/IR/CIRTypes.h"
+#include "clang/CIR/Dialect/IR/CIRAttrs.h"
 #include "clang/CIR/MissingFeatures.h"
 
 namespace cir {
@@ -108,7 +109,7 @@ struct CIRGenTypeCache {
   //     unsigned char SizeAlignInBytes;
   //   };
 
-  clang::LangAS ASTAllocaAddressSpace;
+  mlir::cir::AddressSpaceAttr CIRAllocaAddressSpace;
 
   //   clang::CharUnits getSizeSize() const {
   //     return clang::CharUnits::fromQuantity(SizeSizeInBytes);
@@ -123,13 +124,8 @@ struct CIRGenTypeCache {
     return clang::CharUnits::fromQuantity(PointerAlignInBytes);
   }
 
-  clang::LangAS getASTAllocaAddressSpace() const {
-    // Address spaces are not yet fully supported, but the usage of the default
-    // alloca address space can be used for now only for comparison with the
-    // default address space.
-    assert(!MissingFeatures::addressSpace());
-    assert(ASTAllocaAddressSpace == clang::LangAS::Default);
-    return ASTAllocaAddressSpace;
+  mlir::cir::AddressSpaceAttr getCIRAllocaAddressSpace() const {
+    return CIRAllocaAddressSpace;
   }
 };
 

clang/lib/CIR/CodeGen/TargetInfo.cpp

@@ -263,6 +263,12 @@ class CommonSPIRTargetCIRGenInfo : public TargetCIRGenInfo {
   CommonSPIRTargetCIRGenInfo(std::unique_ptr<ABIInfo> ABIInfo)
       : TargetCIRGenInfo(std::move(ABIInfo)) {}
 
+  mlir::cir::AddressSpaceAttr getCIRAllocaAddressSpace() const override {
+    return mlir::cir::AddressSpaceAttr::get(
+        &getABIInfo().CGT.getMLIRContext(),
+        mlir::cir::AddressSpaceAttr::Kind::offload_private);
+  }
+
   unsigned getOpenCLKernelCallingConv() const override {
     return llvm::CallingConv::SPIR_KERNEL;
   }

clang/lib/CIR/CodeGen/TargetInfo.h

@@ -62,9 +62,11 @@ class TargetCIRGenInfo {
                            std::vector<LValue> &ResultRegDests,
                            std::string &AsmString, unsigned NumOutputs) const {}
 
-  /// Get the AST address space for alloca.
-  virtual clang::LangAS getASTAllocaAddressSpace() const {
-    return clang::LangAS::Default;
+  /// Get the CIR address space for alloca.
+  virtual mlir::cir::AddressSpaceAttr getCIRAllocaAddressSpace() const {
+    // Return the null attribute, which means the target does not care about the
+    // alloca address space.
+    return {};
   }
 
   /// Perform address space cast of an expression of pointer type.

clang/lib/CIR/Lowering/DirectToLLVM/LowerToLLVM.cpp

@@ -999,7 +999,8 @@ class CIRAllocaLowering
                   typeConverter->convertType(rewriter.getIndexType()),
                   rewriter.getIntegerAttr(rewriter.getIndexType(), 1));
     auto elementTy = getTypeConverter()->convertType(op.getAllocaType());
-    auto resultTy = mlir::LLVM::LLVMPointerType::get(getContext());
+    auto resultTy = getTypeConverter()->convertType(op.getResult().getType());
+    // TODO: Verification between the CIR alloca AS and the one from data layout
     rewriter.replaceOpWithNewOp<mlir::LLVM::AllocaOp>(
         op, resultTy, elementTy, size, op.getAlignmentAttr().getInt());
     return mlir::success();

clang/test/CIR/CodeGen/OpenCL/addrspace-alloca.cl

@@ -7,20 +7,28 @@
 // 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(offload_local)>, !cir.ptr<!cir.ptr<!s32i, addrspace(offload_local)>>, ["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)>, addrspace(offload_private)>, ["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}
+  // CIR-NEXT: %[[#ALLOCA_X:]] = cir.alloca !s32i, !cir.ptr<!s32i, addrspace(offload_private)>, ["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(offload_global)>, !cir.ptr<!cir.ptr<!s8i, addrspace(offload_global)>>, ["b"] {alignment = 8 : i64}
+  // CIR-NEXT: %[[#ALLOCA_B:]] = cir.alloca !cir.ptr<!s8i, addrspace(offload_global)>, !cir.ptr<!cir.ptr<!s8i, addrspace(offload_global)>, addrspace(offload_private)>, ["b"] {alignment = 8 : i64}
   // LLVM-NEXT: %[[#ALLOCA_B:]] = alloca ptr addrspace(1), i64 1, align 8
 
+  private int *ptr;
+  // CIR-NEXT: %[[#ALLOCA_PTR:]] = cir.alloca !cir.ptr<!s32i, addrspace(offload_private)>, !cir.ptr<!cir.ptr<!s32i, addrspace(offload_private)>, addrspace(offload_private)>, ["ptr"] {alignment = 8 : i64}
+  // LLVM-NEXT: %[[#ALLOCA_PTR:]] = alloca ptr, i64 1, align 8
+
   // Store of the argument `p`
-  // CIR-NEXT: cir.store %arg0, %[[#ALLOCA_P]] : !cir.ptr<!s32i, addrspace(offload_local)>, !cir.ptr<!cir.ptr<!s32i, addrspace(offload_local)>>
+  // CIR-NEXT: cir.store %arg0, %[[#ALLOCA_P]] : !cir.ptr<!s32i, addrspace(offload_local)>, !cir.ptr<!cir.ptr<!s32i, addrspace(offload_local)>, addrspace(offload_private)>
   // LLVM-NEXT: store ptr addrspace(3) %{{[0-9]+}}, ptr %[[#ALLOCA_P]], align 8
 
+  ptr = &x;
+  // CIR-NEXT: cir.store %[[#ALLOCA_X]], %[[#ALLOCA_PTR]] : !cir.ptr<!s32i, addrspace(offload_private)>, !cir.ptr<!cir.ptr<!s32i, addrspace(offload_private)>, addrspace(offload_private)>
+  // LLVM-NEXT: store ptr %[[#ALLOCA_X]], ptr %[[#ALLOCA_PTR]]
+
   return;
 }