Add alignment support for llvm.alloca

Extend the LLVM dialect AllocaOp with an alignment attribute.

PiperOrigin-RevId: 264068306

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)>,

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;
 }
 

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] : () -> ()

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
+}

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
+}