[Mlir-commits] [mlir] [mlir][LLVMIR] Verify that `IntegerAttr` type matches result type in `llvm.mlir.constant` (PR #205304)

[Anutosh Bhat]

https://github.com/anutosh491 updated https://github.com/llvm/llvm-project/pull/205304

>From 7c92de44f3f34fbe96757ca6d3d6e26ce355079a Mon Sep 17 00:00:00 2001
From: anutosh491 <andersonbhat491 at gmail.com>
Date: Tue, 23 Jun 2026 15:04:03 +0530
Subject: [PATCH] [mlir][LLVMIR] Guard ShlOp/OrOp folds against APInt bitwidth
 mismatch

---
 mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp    | 12 +++++++++++
 .../test/Dialect/LLVMIR/constant-folding.mlir | 20 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 58f569abff8ea..2610139b2e8be 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -4001,6 +4001,12 @@ OpFoldResult LLVM::ShlOp::fold(FoldAdaptor adaptor) {
   if (!lhs)
     return {};
 
+  // Bail out if APInt bitwidths don't match the result type.
+  unsigned width = getType().getIntOrFloatBitWidth();
+  if (lhs.getValue().getBitWidth() != width ||
+      rhs.getValue().getBitWidth() != width)
+    return {};
+
   return IntegerAttr::get(getType(), lhs.getValue().shl(rhs.getValue()));
 }
 
@@ -4017,6 +4023,12 @@ OpFoldResult LLVM::OrOp::fold(FoldAdaptor adaptor) {
   if (!rhs)
     return {};
 
+  // Same bitwidth guard as ShlOp::fold.
+  unsigned width = getType().getIntOrFloatBitWidth();
+  if (lhs.getValue().getBitWidth() != width ||
+      rhs.getValue().getBitWidth() != width)
+    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 0616f19b8fddb..7ae017e1ba896 100644
--- a/mlir/test/Dialect/LLVMIR/constant-folding.mlir
+++ b/mlir/test/Dialect/LLVMIR/constant-folding.mlir
@@ -211,3 +211,23 @@ llvm.func @blockaddress_select(%arg: i1) -> !llvm.ptr {
   llvm.blocktag <id = 1>
   llvm.return %1 : !llvm.ptr
 }
+
+// -----
+
+// CHECK-LABEL: llvm.func @shl_bitwidth_mismatch
+llvm.func @shl_bitwidth_mismatch() -> i64 {
+  %0 = llvm.mlir.constant(1 : i8) : i64
+  // CHECK: llvm.shl
+  %1 = llvm.shl %0, %0 : i64
+  llvm.return %1 : i64
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @or_bitwidth_mismatch
+llvm.func @or_bitwidth_mismatch() -> i64 {
+  %0 = llvm.mlir.constant(1 : i8) : i64
+  // CHECK: llvm.or
+  %1 = llvm.or %0, %0 : i64
+  llvm.return %1 : i64
+}