[Mlir-commits] [mlir] [mlir][LLVMIR] Skip or folding for mismatched constant types (PR #226026)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-llvm

Author: Akun (AkunCube)

<details>
<summary>Changes</summary>

A constant-like operation can fold to an integer attribute whose type differs from its SSA result type. After arith.ori is converted to llvm.or, LLVM::OrOp::fold tries to construct an i32 IntegerAttr from an i64 APInt and aborts in an assertions-enabled build.

Check both folded operand attribute types against the llvm.or result type before folding. Add regression cases for mismatches on either operand; the existing matching-type case continues to fold.

Fixes: #<!-- -->225659

---
Full diff: https://github.com/llvm/llvm-project/pull/226026.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp (+2-2) 
- (modified) mlir/test/Dialect/LLVMIR/constant-folding.mlir (+30) 


``````````diff
diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index b931631369e6f..74d6934307559 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -4148,11 +4148,11 @@ OpFoldResult LLVM::ShlOp::fold(FoldAdaptor adaptor) {
 
 OpFoldResult LLVM::OrOp::fold(FoldAdaptor adaptor) {
   auto lhs = dyn_cast_or_null<IntegerAttr>(adaptor.getLhs());
-  if (!lhs)
+  if (!lhs || lhs.getType() != getType())
     return {};
 
   auto rhs = dyn_cast_or_null<IntegerAttr>(adaptor.getRhs());
-  if (!rhs)
+  if (!rhs || rhs.getType() != getType())
     return {};
 
   return IntegerAttr::get(getType(), lhs.getValue() | rhs.getValue());
diff --git a/mlir/test/Dialect/LLVMIR/constant-folding.mlir b/mlir/test/Dialect/LLVMIR/constant-folding.mlir
index 0a6227bd2607b..cf5f062a6fffa 100644
--- a/mlir/test/Dialect/LLVMIR/constant-folding.mlir
+++ b/mlir/test/Dialect/LLVMIR/constant-folding.mlir
@@ -88,6 +88,36 @@ llvm.func @or_basic() -> i32 {
 
 // -----
 
+// A constant-like op can fold to an attribute whose type differs from its
+// result type. Do not fold llvm.or using such an attribute.
+// CHECK-LABEL: llvm.func @or_mismatched_lhs
+llvm.func @or_mismatched_lhs() -> i32 {
+  // CHECK: %[[LHS:.*]] = "test.constant"() <{value = 1 : i64}> : () -> i32
+  %lhs = "test.constant"() {value = 1 : i64} : () -> i32
+  // CHECK: %[[RHS:.*]] = "test.constant"() <{value = 2 : i32}> : () -> i32
+  %rhs = "test.constant"() {value = 2 : i32} : () -> i32
+  // CHECK: %[[RESULT:.*]] = llvm.or %[[LHS]], %[[RHS]] : i32
+  %result = llvm.or %lhs, %rhs : i32
+  // CHECK: llvm.return %[[RESULT]] : i32
+  llvm.return %result : i32
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @or_mismatched_rhs
+llvm.func @or_mismatched_rhs() -> i32 {
+  // CHECK: %[[LHS:.*]] = "test.constant"() <{value = 1 : i32}> : () -> i32
+  %lhs = "test.constant"() {value = 1 : i32} : () -> i32
+  // CHECK: %[[RHS:.*]] = "test.constant"() <{value = 2 : i64}> : () -> i32
+  %rhs = "test.constant"() {value = 2 : i64} : () -> i32
+  // CHECK: %[[RESULT:.*]] = llvm.or %[[LHS]], %[[RHS]] : i32
+  %result = llvm.or %lhs, %rhs : i32
+  // CHECK: llvm.return %[[RESULT]] : i32
+  llvm.return %result : i32
+}
+
+// -----
+
 // CHECK-LABEL: llvm.func @addressof
 llvm.func @addressof() {
   // CHECK-NEXT: %[[ADDRESSOF:.+]] = llvm.mlir.addressof @foo

``````````

</details>


https://github.com/llvm/llvm-project/pull/226026