[Mlir-commits] [mlir] 5f6ce97 - [MLIR] Add support for alloca() of non-zero address spaces

[Krzysztof Drewniak]

Author: Krzysztof Drewniak
Date: 2022-02-15T18:41:30Z
New Revision: 5f6ce97408e978945d44ba33066514fa11f9a416

URL: https://github.com/llvm/llvm-project/commit/5f6ce97408e978945d44ba33066514fa11f9a416
DIFF: https://github.com/llvm/llvm-project/commit/5f6ce97408e978945d44ba33066514fa11f9a416.diff

LOG: [MLIR] Add support for alloca() of non-zero address spaces

This change is needed when lowering alloca()-using code on targets
such as ROCDL that represent private scratch space as a separate
address space.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D119775

Added: 
    

Modified: 
    mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
    mlir/test/Target/LLVMIR/llvmir.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
index d9883f33086e8..73efeadbdbda2 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td
@@ -296,8 +296,9 @@ def LLVM_AllocaOp : LLVM_Op<"alloca">, MemoryOpWithAlignmentBase {
                    OptionalAttr<I64Attr>:$alignment);
   let results = (outs Res<LLVM_AnyPointer, "", [MemAlloc<AutomaticAllocationScopeResource>]>:$res);
   string llvmBuilder = [{
+    auto addrSpace = $_resultType->getPointerAddressSpace();
     auto *inst = builder.CreateAlloca(
-      $_resultType->getPointerElementType(), $arraySize);
+        $_resultType->getPointerElementType(), addrSpace, $arraySize);
     }] # setAlignmentCode # [{
     $res = inst;
   }];

diff  --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index ef369714d94f1..9be89b28df82e 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -1176,6 +1176,8 @@ llvm.func @alloca(%size : i64) {
   llvm.alloca %size x i32 {alignment = 0} : (i64) -> (!llvm.ptr<i32>)
   // CHECK-NEXT: alloca {{.*}} align 8
   llvm.alloca %size x i32 {alignment = 8} : (i64) -> (!llvm.ptr<i32>)
+  // CHECK-NEXT: alloca {{.*}} addrspace(3)
+  llvm.alloca %size x i32 {alignment = 0} : (i64) -> (!llvm.ptr<i32, 3>)
   llvm.return
 }