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

[krzysz00]

Summary

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

[mehdi_amini]

Can this be reflected in a test?

[ftynse]

Please add a test for this, looks good otherwise.

[krzysz00]

• Add a test for nonzero address space on llvm.alloca

[krzysz00]

Rebase so there's a working build

mlir/include/mlir/Dialect/LLVMIR/LLVMOps.td

 }
 
 // Memory-related operations.
 def LLVM_AllocaOp : LLVM_Op<"alloca">, MemoryOpWithAlignmentBase {
   let arguments = (ins AnyInteger:$arraySize,
                    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;
   }];
   let builders = [
     OpBuilder<(ins "Type":$resultType, "Value":$arraySize,
       "unsigned":$alignment),
     [{
       if (alignment == 0)

mlir/test/Target/LLVMIR/llvmir.mlir

 // CHECK-LABEL: @alloca
 llvm.func @alloca(%size : i64) {
   // Alignment automatically set by the LLVM IR builder when alignment attribute
   // is 0.
   //  CHECK: alloca {{.*}} align 4
   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
 }
 
 // CHECK-LABEL: @constants
 llvm.func @constants() -> vector<4xf32> {
   // CHECK: ret <4 x float> <float 4.2{{0*}}e+01, float 0.{{0*}}e+00, float 0.{{0*}}e+00, float 0.{{0*}}e+00>
   %0 = llvm.mlir.constant(sparse<[0], [4.2e+01]> : vector<4xf32>) : vector<4xf32>
   llvm.return %0 : vector<4xf32>