[CIR][CIRGen][Dialect] Use the correct alloca address space

Author: seven-mile

clang/include/clang/CIR/Dialect/IR/CIRDataLayout.h

@@ -100,6 +100,16 @@ class CIRDataLayout {
                                          getPointerTypeSizeInBits(Ty), false);
     return IntTy;
   }
+
+  unsigned getAllocaMemorySpace() const {
+    mlir::Attribute memorySpaceAttr = layout.getAllocaMemorySpace();
+    // The empty attribute represents the default address space (0)
+    if (!memorySpaceAttr)
+      return 0;
+    auto intAttr = mlir::dyn_cast<mlir::IntegerAttr>(memorySpaceAttr);
+    assert(intAttr && "Expects integer attributes for memory space");
+    return static_cast<unsigned>(intAttr.getUInt());
+  }
 };
 
 } // namespace cir

clang/include/clang/CIR/MissingFeatures.h

@@ -32,6 +32,7 @@ struct MissingFeatures {
   // Address space related
   static bool addressSpace() { return false; }
   static bool addressSpaceInGlobalVar() { return false; }
+  static bool addressSpaceInAlloca() { return false; }
 
   // Clang codegen options
   static bool strictVTablePointers() { return false; }

clang/lib/CIR/CodeGen/CIRGenBuilder.h

@@ -417,10 +417,12 @@ class CIRGenBuilderTy : public CIRBaseBuilderTy {
 
   // Fetch the type representing a pointer to unsigned int values.
   mlir::cir::PointerType getUInt8PtrTy(unsigned AddrSpace = 0) {
-    return typeCache.UInt8PtrTy;
+    if (AddrSpace == 0)
+      return typeCache.UInt8PtrTy;
+    return mlir::cir::PointerType::get(getContext(), typeCache.UInt8Ty, AddrSpace);
   }
   mlir::cir::PointerType getUInt32PtrTy(unsigned AddrSpace = 0) {
-    return mlir::cir::PointerType::get(getContext(), typeCache.UInt32Ty);
+    return mlir::cir::PointerType::get(getContext(), typeCache.UInt32Ty, AddrSpace);
   }
 
   /// Get a CIR anonymous struct type.

clang/lib/CIR/CodeGen/CIRGenBuiltin.cpp

@@ -986,8 +986,8 @@ RValue CIRGenFunction::buildBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID,
     // pointer of type `void *`. This will require a change to the allocaOp
     // verifier.
     auto AllocaAddr = builder.createAlloca(
-        getLoc(E->getSourceRange()), builder.getUInt8PtrTy(),
-        builder.getUInt8Ty(), "bi_alloca", SuitableAlignmentInBytes, Size);
+        getLoc(E->getSourceRange()), CGM.AllocaInt8PtrTy, builder.getUInt8Ty(),
+        "bi_alloca", SuitableAlignmentInBytes, Size);
 
     // Initialize the allocated buffer if required.
     if (BuiltinID != Builtin::BI__builtin_alloca_uninitialized)

clang/lib/CIR/CodeGen/CIRGenExpr.cpp

@@ -2718,7 +2718,8 @@ 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);
+  auto localVarPtrTy = mlir::cir::PointerType::get(
+      builder.getContext(), ty, CGM.getDataLayout().getAllocaMemorySpace());
   auto alignIntAttr = CGM.getSize(alignment);
 
   mlir::Value addr;

clang/lib/CIR/CodeGen/CIRGenModule.cpp

@@ -31,6 +31,7 @@
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/IR/Verifier.h"
+#include "mlir/Target/LLVMIR/Import.h"
 #include "clang/CIR/MissingFeatures.h"
 
 #include "clang/AST/ASTConsumer.h"
@@ -96,6 +97,14 @@ static CIRGenCXXABI *createCXXABI(CIRGenModule &CGM) {
   }
 }
 
+static void setMLIRDataLayout(mlir::ModuleOp &mod, const llvm::DataLayout &dl) {
+  auto *context = mod.getContext();
+  mod->setAttr(mlir::LLVM::LLVMDialect::getDataLayoutAttrName(),
+               mlir::StringAttr::get(context, dl.getStringRepresentation()));
+  mlir::DataLayoutSpecInterface dlSpec = mlir::translateDataLayout(dl, context);
+  mod->setAttr(mlir::DLTIDialect::kDataLayoutAttrName, dlSpec);
+}
+
 CIRGenModule::CIRGenModule(mlir::MLIRContext &context,
                            clang::ASTContext &astctx,
                            const clang::CodeGenOptions &CGO,
@@ -106,6 +115,10 @@ CIRGenModule::CIRGenModule(mlir::MLIRContext &context,
       target(astCtx.getTargetInfo()), ABI(createCXXABI(*this)), genTypes{*this},
       VTables{*this}, openMPRuntime(new CIRGenOpenMPRuntime(*this)) {
 
+  // Initialize DataLayout in the module op.
+  auto layout = llvm::DataLayout(astCtx.getTargetInfo().getDataLayoutString());
+  setMLIRDataLayout(theModule, layout);
+
   // Initialize CIR signed integer types cache.
   SInt8Ty =
       ::mlir::cir::IntType::get(builder.getContext(), 8, /*isSigned=*/true);
@@ -156,7 +169,8 @@ CIRGenModule::CIRGenModule(mlir::MLIRContext &context,
       /*isSigned=*/false);
   UInt8PtrTy = builder.getPointerTo(UInt8Ty);
   UInt8PtrPtrTy = builder.getPointerTo(UInt8PtrTy);
-  AllocaInt8PtrTy = UInt8PtrTy;
+  AllocaInt8PtrTy =
+      builder.getPointerTo(UInt8Ty, getDataLayout().getAllocaMemorySpace());
   // TODO: GlobalsInt8PtrTy
   // TODO: ConstGlobalsPtrTy
   ASTAllocaAddressSpace = getTargetCIRGenInfo().getASTAllocaAddressSpace();

clang/lib/CIR/CodeGen/CIRGenerator.cpp

@@ -40,14 +40,6 @@ CIRGenerator::~CIRGenerator() {
   assert(DeferredInlineMemberFuncDefs.empty() || Diags.hasErrorOccurred());
 }
 
-static void setMLIRDataLayout(mlir::ModuleOp &mod, const llvm::DataLayout &dl) {
-  auto *context = mod.getContext();
-  mod->setAttr(mlir::LLVM::LLVMDialect::getDataLayoutAttrName(),
-               mlir::StringAttr::get(context, dl.getStringRepresentation()));
-  mlir::DataLayoutSpecInterface dlSpec = mlir::translateDataLayout(dl, context);
-  mod->setAttr(mlir::DLTIDialect::kDataLayoutAttrName, dlSpec);
-}
-
 void CIRGenerator::Initialize(ASTContext &astCtx) {
   using namespace llvm;
 
@@ -62,9 +54,6 @@ void CIRGenerator::Initialize(ASTContext &astCtx) {
   mlirCtx->getOrLoadDialect<mlir::omp::OpenMPDialect>();
   CGM = std::make_unique<CIRGenModule>(*mlirCtx.get(), astCtx, codeGenOpts,
                                        Diags);
-  auto mod = CGM->getModule();
-  auto layout = llvm::DataLayout(astCtx.getTargetInfo().getDataLayoutString());
-  setMLIRDataLayout(mod, layout);
 }
 
 bool CIRGenerator::verifyModule() { return CGM->verifyModule(); }

clang/lib/CIR/Dialect/Transforms/LoweringPrepare.cpp

@@ -21,6 +21,7 @@
 #include "clang/CIR/Dialect/IR/CIRDialect.h"
 #include "clang/CIR/Dialect/Passes.h"
 #include "clang/CIR/Interfaces/ASTAttrInterfaces.h"
+#include "clang/CIR/MissingFeatures.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringRef.h"
@@ -526,6 +527,7 @@ static void lowerArrayDtorCtorIntoLoop(CIRBaseBuilderTy &builder,
   mlir::Value end = builder.create<mlir::cir::PtrStrideOp>(
       loc, eltTy, begin, numArrayElementsConst);
 
+  assert(!::cir::MissingFeatures::addressSpaceInAlloca());
   auto tmpAddr = builder.createAlloca(
       loc, /*addr type*/ builder.getPointerTo(eltTy),
       /*var type*/ eltTy, "__array_idx", clang::CharUnits::One());