[Mlir-commits] [mlir] [MLIR][LLVM] Lift alignstack attribute ptr type restriction (PR #133195)

[Bruno Cardoso Lopes]

https://github.com/bcardosolopes created https://github.com/llvm/llvm-project/pull/133195

Current usage of alignstack is restricted to LLVM pointer types, whereas when it's used in parameters it's possible to use it for other types, see examples like `{i8, i8}, [2 x float], etc` in `llvm/test/CodeGen`. This PR lifts the restriction and add testcases.

>From 875918a978c44dd65142f375489596fa8d240d60 Mon Sep 17 00:00:00 2001
From: Bruno Cardoso Lopes <bruno.cardoso at gmail.com>
Date: Wed, 26 Mar 2025 19:12:10 -0700
Subject: [PATCH] [MLIR][LLVM] Allow alignstack attribute with more types

Current usage of alignstack is restricted to LLVM pointer types, whereas when
it's used in parameters it's possible to use it for other types, see examples
like `{i8, i8}, [2 x float], etc` in `llvm/test/CodeGen`. This PR lifts the
restriction and add testcases.
---
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp    | 10 ++++++--
 mlir/test/Dialect/LLVMIR/call-param.mlir      | 23 ++++++++++++++++++
 .../LLVMIR/parameter-attrs-invalid.mlir       |  5 ----
 mlir/test/Target/LLVMIR/llvmir.mlir           | 24 +++++++++++++++++++
 4 files changed, 55 insertions(+), 7 deletions(-)
 create mode 100644 mlir/test/Dialect/LLVMIR/call-param.mlir

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 18a70cc64628f..32b5e8bdcebeb 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -4032,8 +4032,7 @@ LogicalResult LLVMDialect::verifyParameterAttribute(Operation *op,
   // Check an integer attribute that is attached to a pointer value.
   if (name == LLVMDialect::getAlignAttrName() ||
       name == LLVMDialect::getDereferenceableAttrName() ||
-      name == LLVMDialect::getDereferenceableOrNullAttrName() ||
-      name == LLVMDialect::getStackAlignmentAttrName()) {
+      name == LLVMDialect::getDereferenceableOrNullAttrName()) {
     if (failed(checkIntegerAttrType()))
       return failure();
     if (verifyValueType && failed(checkPointerType()))
@@ -4041,6 +4040,13 @@ LogicalResult LLVMDialect::verifyParameterAttribute(Operation *op,
     return success();
   }
 
+  // Check an integer attribute that is attached to a pointer value.
+  if (name == LLVMDialect::getStackAlignmentAttrName()) {
+    if (failed(checkIntegerAttrType()))
+      return failure();
+    return success();
+  }
+
   // Check a unit attribute that can be attached to arbitrary types.
   if (name == LLVMDialect::getNoUndefAttrName() ||
       name == LLVMDialect::getInRegAttrName() ||
diff --git a/mlir/test/Dialect/LLVMIR/call-param.mlir b/mlir/test/Dialect/LLVMIR/call-param.mlir
new file mode 100644
index 0000000000000..f64c233c47625
--- /dev/null
+++ b/mlir/test/Dialect/LLVMIR/call-param.mlir
@@ -0,0 +1,23 @@
+// RUN: mlir-opt %s -split-input-file | FileCheck %s
+
+llvm.func local_unnamed_addr @testfn(!llvm.array<2 x f32> {llvm.alignstack = 8 : i64})
+llvm.func internal @g(%arg0: !llvm.array<2 x f32>) attributes {dso_local} {
+  // CHECK-LABEL: @g
+  // CHECK: llvm.call @testfn(%arg0) : (!llvm.array<2 x f32> {llvm.alignstack = 8 : i64}) -> ()
+  llvm.call @testfn(%arg0) : (!llvm.array<2 x f32> {llvm.alignstack = 8 : i64}) -> ()
+  llvm.return
+}
+llvm.func local_unnamed_addr @testfn2(!llvm.struct<(i8, i8)> {llvm.alignstack = 8 : i64})
+llvm.func internal @h(%arg0: !llvm.struct<(i8, i8)>) attributes {dso_local} {
+  // CHECK-LABEL: @h
+  // CHECK: llvm.call @testfn2(%arg0) : (!llvm.struct<(i8, i8)> {llvm.alignstack = 8 : i64}) -> ()
+  llvm.call @testfn2(%arg0) : (!llvm.struct<(i8, i8)> {llvm.alignstack = 8 : i64}) -> ()
+  llvm.return
+}
+llvm.func local_unnamed_addr @testfn3(i32 {llvm.alignstack = 8 : i64})
+llvm.func internal @i(%arg0: i32) attributes {dso_local} {
+  // CHECK-LABEL: @i
+  // CHECK: llvm.call @testfn3(%arg0) : (i32 {llvm.alignstack = 8 : i64}) -> ()
+  llvm.call @testfn3(%arg0) : (i32 {llvm.alignstack = 8 : i64}) -> ()
+  llvm.return
+}
diff --git a/mlir/test/Dialect/LLVMIR/parameter-attrs-invalid.mlir b/mlir/test/Dialect/LLVMIR/parameter-attrs-invalid.mlir
index 55b1d4faf207f..1cf561ac6c6a5 100644
--- a/mlir/test/Dialect/LLVMIR/parameter-attrs-invalid.mlir
+++ b/mlir/test/Dialect/LLVMIR/parameter-attrs-invalid.mlir
@@ -157,11 +157,6 @@ llvm.func @invalid_returned_attr_type(%0 : i32 {llvm.returned = !llvm.ptr})
 
 // -----
 
-// expected-error at below {{"llvm.alignstack" attribute attached to non-pointer LLVM type}}
-llvm.func @invalid_alignstack_arg_type(%0 : i32 {llvm.alignstack = 10 : i32})
-
-// -----
-
 // expected-error at below {{"llvm.alignstack" should be an integer attribute}}
 llvm.func @invalid_alignstack_attr_type(%0 : i32 {llvm.alignstack = "foo"})
 
diff --git a/mlir/test/Target/LLVMIR/llvmir.mlir b/mlir/test/Target/LLVMIR/llvmir.mlir
index 2476b1b15faaa..89faa6abb6cf7 100644
--- a/mlir/test/Target/LLVMIR/llvmir.mlir
+++ b/mlir/test/Target/LLVMIR/llvmir.mlir
@@ -2842,3 +2842,27 @@ llvm.func @call_alias_func() {
 
 // CHECK-LABEL: @call_alias_func
 // CHECK: call void dso_local_equivalent @alias_func()
+
+// -----
+
+llvm.func local_unnamed_addr @testfn(!llvm.array<2 x f32> {llvm.alignstack = 8 : i64})
+llvm.func internal @g(%arg0: !llvm.array<2 x f32>) attributes {dso_local} {
+  // CHECK-LABEL: @g
+  // CHECK: call void @testfn([2 x float] alignstack(8) %0)
+  llvm.call @testfn(%arg0) : (!llvm.array<2 x f32> {llvm.alignstack = 8 : i64}) -> ()
+  llvm.return
+}
+llvm.func local_unnamed_addr @testfn2(!llvm.struct<(i8, i8)> {llvm.alignstack = 8 : i64})
+llvm.func internal @h(%arg0: !llvm.struct<(i8, i8)>) attributes {dso_local} {
+  // CHECK-LABEL: @h
+  // CHECK: call void @testfn2({ i8, i8 } alignstack(8) %0)
+  llvm.call @testfn2(%arg0) : (!llvm.struct<(i8, i8)> {llvm.alignstack = 8 : i64}) -> ()
+  llvm.return
+}
+llvm.func local_unnamed_addr @testfn3(i32 {llvm.alignstack = 8 : i64})
+llvm.func internal @i(%arg0: i32) attributes {dso_local} {
+  // CHECK-LABEL: @i
+  // CHECK: call void @testfn3(i32 alignstack(8) %0)
+  llvm.call @testfn3(%arg0) : (i32 {llvm.alignstack = 8 : i64}) -> ()
+  llvm.return
+}