From c9f37fca379035b6334b50380ef05b00026de0cc Mon Sep 17 00:00:00 2001
From: Nicolas Vasilache <ntv@google.com>
Date: Sun, 18 Aug 2019 18:54:50 -0700
Subject: [PATCH] Add alignment support for llvm.alloca

Extend the LLVM dialect AllocaOp with an alignment attribute.

PiperOrigin-RevId: 264068306
---
 mlir/include/mlir/LLVMIR/LLVMOps.td | 33 +++++++++++++++++++++++++++++----
 mlir/lib/LLVMIR/IR/LLVMDialect.cpp  |  5 ++++-
 mlir/test/LLVMIR/invalid.mlir       |  7 +++++++
 mlir/test/LLVMIR/roundtrip.mlir     |  9 +++++++++
 mlir/test/Target/llvmir.mlir        | 10 ++++++++++
 5 files changed, 59 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/LLVMIR/LLVMOps.td b/mlir/include/mlir/LLVMIR/LLVMOps.td
index 3783ad9..d988f36 100644
--- 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 199d401..d051cc9 100644
--- 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 111d90a..a5e2248 100644
--- 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 4348b7b..36fe537 100644
--- 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 36fa128..e358c9e 100644
--- 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
+}
-- 
2.7.4