Add alignment support for llvm.alloca

Extend the LLVM dialect AllocaOp with an alignment attribute.

PiperOrigin-RevId: 264068306

diff --git a/mlir/include/mlir/LLVMIR/LLVMOps.td b/mlir/include/mlir/LLVMIR/LLVMOps.td
index 3783ad907f87830e4134761d85c0d2a033c58a37..d988f36d12c58649f4a2c2ba2bac5f44b9aaf0d7 100644 (file)
--- a/mlir/include/mlir/LLVMIR/LLVMOps.td
+++ b/mlir/include/mlir/LLVMIR/LLVMOps.td
@@ -220,14 +220,39 @@ def LLVM_FDivOp : LLVM_ArithmeticOp<"fdiv", "CreateFDiv">;
 def LLVM_FRemOp : LLVM_ArithmeticOp<"frem", "CreateFRem">;
 
 // Memory-related operations.
-def LLVM_AllocaOp : LLVM_OneResultOp<"alloca">,
-                    Arguments<(ins LLVM_Type:$arraySize)> {
+def LLVM_AllocaOp :
+    LLVM_OneResultOp<"alloca">,
+    Arguments<(ins LLVM_Type:$arraySize, OptionalAttr<I64Attr>:$alignment)> {
   string llvmBuilder = [{
-    $res = builder.CreateAlloca($_resultType->getPointerElementType(),
-                                $arraySize);
+    auto *alloca = builder.CreateAlloca(
+      $_resultType->getPointerElementType(), $arraySize);
+    if ($alignment.hasValue()) {
+      auto align = $alignment.getValue().getZExtValue();
+      if (align != 0)
+        alloca->setAlignment(align);
+    }
+    $res = alloca;
   }];
+  let builders = [OpBuilder<
+    "Builder *b, OperationState *result, Type resultType, Value *arraySize,"
+    "unsigned alignment = 0",
+    [{
+      if (!alignment)
+        return build(b, result, resultType, arraySize, IntegerAttr());
+      auto *ctx = resultType.getContext();
+      auto align = IntegerAttr::get(IntegerType::get(64, ctx), alignment);
+      build(b, result, resultType, arraySize, align);
+  }]>];
   let parser = [{ return parseAllocaOp(parser, result); }];
   let printer = [{ printAllocaOp(p, *this); }];
+  let verifier = [{
+    if (alignment().hasValue()) {
+      auto align = alignment().getValue().getSExtValue();
+      if (align < 0)
+        return emitOpError("expected positive alignment");
+    }
+    return success();
+  }];
 }
 def LLVM_GEPOp : LLVM_OneResultOp<"getelementptr", [NoSideEffect]>,
                  Arguments<(ins LLVM_Type:$base, Variadic<LLVM_Type>:$indices)>,

diff --git a/mlir/lib/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/LLVMIR/IR/LLVMDialect.cpp
index 199d40150dc7f4066377c32edd69fbe930779654..d051cc94d5b1c0aaf16dc57faa52e9a36eec26ba 100644 (file)
--- a/mlir/lib/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/LLVMIR/IR/LLVMDialect.cpp
@@ -131,7 +131,10 @@ static void printAllocaOp(OpAsmPrinter *p, AllocaOp &op) {
                                   op.getContext());
 
   *p << op.getOperationName() << ' ' << *op.arraySize() << " x " << elemTy;
-  p->printOptionalAttrDict(op.getAttrs());
+  if (op.alignment().hasValue() && op.alignment()->getSExtValue() != 0)
+    p->printOptionalAttrDict(op.getAttrs());
+  else
+    p->printOptionalAttrDict(op.getAttrs(), {"alignment"});
   *p << " : " << funcTy;
 }
 

diff --git a/mlir/test/LLVMIR/invalid.mlir b/mlir/test/LLVMIR/invalid.mlir
index 111d90a90a8ca9b1b73b67671a9de9354b2503e5..a5e2248487cf1740de02efe69cce630f89ccfc33 100644 (file)
--- a/mlir/test/LLVMIR/invalid.mlir
+++ b/mlir/test/LLVMIR/invalid.mlir
@@ -55,6 +55,13 @@ func @alloca_non_function_type() {
 
 // -----
 
+func @alloca_nonpositive_alignment(%size : !llvm.i64) {
+  // expected-error@+1 {{expected positive alignment}}
+  llvm.alloca %size x !llvm.i32 {alignment = -1} : (!llvm.i64) -> (!llvm<"i32*">)
+}
+
+// -----
+
 func @gep_missing_input_result_type(%pos : !llvm.i64, %base : !llvm<"float*">) {
   // expected-error@+1 {{expected trailing function type with at least one argument and one result}}
   llvm.getelementptr %base[%pos] : () -> ()

diff --git a/mlir/test/LLVMIR/roundtrip.mlir b/mlir/test/LLVMIR/roundtrip.mlir
index 4348b7b6a8daf4ed3987b26255217efe4354e7cf..36fe537b60590225b2121ba29a9dec90bf5b4de6 100644 (file)
--- a/mlir/test/LLVMIR/roundtrip.mlir
+++ b/mlir/test/LLVMIR/roundtrip.mlir
@@ -178,3 +178,12 @@ func @vect(%arg0: !llvm<"<4 x float>">, %arg1: !llvm.i32, %arg2: !llvm.float) {
   %2 = llvm.shufflevector %arg0, %arg0 [0 : i32, 0 : i32, 0 : i32, 0 : i32, 7 : i32] : !llvm<"<4 x float>">, !llvm<"<4 x float>">
   return
 }
+
+// CHECK-LABEL: @alloca
+func @alloca(%size : !llvm.i64) {
+  //      CHECK: llvm.alloca %{{.*}} x !llvm.i32 : (!llvm.i64) -> !llvm<"i32*">
+  llvm.alloca %size x !llvm.i32 {alignment = 0} : (!llvm.i64) -> (!llvm<"i32*">)
+  // CHECK-NEXT: llvm.alloca %{{.*}} x !llvm.i32 {alignment = 8 : i64} : (!llvm.i64) -> !llvm<"i32*">
+  llvm.alloca %size x !llvm.i32 {alignment = 8} : (!llvm.i64) -> (!llvm<"i32*">)
+  llvm.return
+}

diff --git a/mlir/test/Target/llvmir.mlir b/mlir/test/Target/llvmir.mlir
index 36fa12834fead646ba66c0262d28a36c75df3414..e358c9edb9ebba0096b7635c9030093c860c2d5d 100644 (file)
--- a/mlir/test/Target/llvmir.mlir
+++ b/mlir/test/Target/llvmir.mlir
@@ -902,3 +902,13 @@ func @vect(%arg0: !llvm<"<4 x float>">, %arg1: !llvm.i32, %arg2: !llvm.float) {
   %2 = llvm.shufflevector %arg0, %arg0 [0 : i32, 0 : i32, 0 : i32, 0 : i32, 7 : i32] : !llvm<"<4 x float>">, !llvm<"<4 x float>">
   llvm.return
 }
+
+// CHECK-LABEL: @alloca
+func @alloca(%size : !llvm.i64) {
+  //      CHECK: alloca
+  //  CHECK-NOT: align
+  llvm.alloca %size x !llvm.i32 {alignment = 0} : (!llvm.i64) -> (!llvm<"i32*">)
+  // CHECK-NEXT: alloca {{.*}} align 8
+  llvm.alloca %size x !llvm.i32 {alignment = 8} : (!llvm.i64) -> (!llvm<"i32*">)
+  llvm.return
+}