[Mlir-commits] [mlir] Added a check for unsigned integer before accessing (PR #204276)

[Balaji V. Iyer.]

https://github.com/bviyer updated https://github.com/llvm/llvm-project/pull/204276

>From aa944c5373e026cae0f8036aa22eb19d3d30a640 Mon Sep 17 00:00:00 2001
From: "Balaji V. Iyer" <bviyer at gmail.com>
Date: Tue, 16 Jun 2026 22:17:37 -0500
Subject: [PATCH 1/2] Added a check for unsigned integer before accessing

---
 mlir/test/Analysis/test-dataflow-framework.mlir  | 14 ++++++++++++++
 mlir/test/lib/Analysis/TestDataFlowFramework.cpp |  4 +++-
 2 files changed, 17 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Analysis/test-dataflow-framework.mlir

diff --git a/mlir/test/Analysis/test-dataflow-framework.mlir b/mlir/test/Analysis/test-dataflow-framework.mlir
new file mode 100644
index 0000000000000..ccb9f6a0721f8
--- /dev/null
+++ b/mlir/test/Analysis/test-dataflow-framework.mlir
@@ -0,0 +1,14 @@
+// RUN: mlir-opt -test-staged-analyses -split-input-file %s | FileCheck %s
+
+// Test that FooAnalysis does not crash when a func.func carries integer
+// attributes with signless i64 type (e.g. llvm.link_call_count).
+
+// CHECK-LABEL: func.func private @callee
+module {
+  func.func private @callee(%arg0: i32) -> i32 attributes {llvm.link_call_count = 10 : i64, sym.link_call_count = 11 : i64} {
+    %0:2 = "test.foo"() {foo = 5 : i32} : () -> (i32, i32)
+    %1 = "test.foo"() {foo = 7 : i32} : () -> i32
+    %2 = "test.hello"(%0#0, %1) : (i32, i32) -> i32
+    return %2 : i32
+  }
+}
diff --git a/mlir/test/lib/Analysis/TestDataFlowFramework.cpp b/mlir/test/lib/Analysis/TestDataFlowFramework.cpp
index 9af7e205aaee9..f29136853a9cd 100644
--- a/mlir/test/lib/Analysis/TestDataFlowFramework.cpp
+++ b/mlir/test/lib/Analysis/TestDataFlowFramework.cpp
@@ -239,7 +239,9 @@ void FooAnalysis::visitOperation(Operation *op) {
 
   // Modify the state with the attribute, if specified.
   if (auto attr = op->getAttrOfType<IntegerAttr>(kFooAttrName)) {
-    uint64_t value = attr.getUInt();
+    uint64_t value = attr.getType().isUnsignedInteger()
+                         ? attr.getUInt()
+                         : static_cast<uint64_t>(attr.getInt());
     result |= state->join(value);
   }
   propagateIfChanged(state, result);

>From 3119b23d9faa74d71632e6d2bc688853d1c61960 Mon Sep 17 00:00:00 2001
From: "Balaji V. Iyer" <bviyer at gmail.com>
Date: Wed, 17 Jun 2026 20:04:31 -0500
Subject: [PATCH 2/2] Moved the test to test-staged-analysis.mlir

---
 .../Analysis/DataFlow/test-staged-analyses.mlir   | 15 +++++++++++++++
 mlir/test/Analysis/test-dataflow-framework.mlir   | 14 --------------
 2 files changed, 15 insertions(+), 14 deletions(-)
 delete mode 100644 mlir/test/Analysis/test-dataflow-framework.mlir

diff --git a/mlir/test/Analysis/DataFlow/test-staged-analyses.mlir b/mlir/test/Analysis/DataFlow/test-staged-analyses.mlir
index 5da7ad8bf1274..740b648078f75 100644
--- a/mlir/test/Analysis/DataFlow/test-staged-analyses.mlir
+++ b/mlir/test/Analysis/DataFlow/test-staged-analyses.mlir
@@ -48,3 +48,18 @@ func.func @requires_staged_bar() {
   // CHECK: "test.branch"()[^bb{{[0-9]+}}] {bar_state = true, foo = 2 : ui64, foo_state = 5 : i64, tag = "annotate"} : () -> ()
   "test.branch"() [^bb1] {tag = "annotate", foo = 2 : ui64} : () -> ()
 }
+
+// -----
+
+// Test that FooAnalysis does not crash when a func.func carries integer
+// attributes with signless i64 type (e.g. llvm.link_call_count).
+
+// CHECK-LABEL: func.func private @callee
+module {
+  func.func private @callee(%arg0: i32) -> i32 attributes {llvm.link_call_count = 10 : i64, sym.link_call_count = 11 : i64} {
+    %0:2 = "test.foo"() {foo = 5 : i32} : () -> (i32, i32)
+    %1 = "test.foo"() {foo = 7 : i32} : () -> i32
+    %2 = "test.hello"(%0#0, %1) : (i32, i32) -> i32
+    return %2 : i32
+  }
+}
diff --git a/mlir/test/Analysis/test-dataflow-framework.mlir b/mlir/test/Analysis/test-dataflow-framework.mlir
deleted file mode 100644
index ccb9f6a0721f8..0000000000000
--- a/mlir/test/Analysis/test-dataflow-framework.mlir
+++ /dev/null
@@ -1,14 +0,0 @@
-// RUN: mlir-opt -test-staged-analyses -split-input-file %s | FileCheck %s
-
-// Test that FooAnalysis does not crash when a func.func carries integer
-// attributes with signless i64 type (e.g. llvm.link_call_count).
-
-// CHECK-LABEL: func.func private @callee
-module {
-  func.func private @callee(%arg0: i32) -> i32 attributes {llvm.link_call_count = 10 : i64, sym.link_call_count = 11 : i64} {
-    %0:2 = "test.foo"() {foo = 5 : i32} : () -> (i32, i32)
-    %1 = "test.foo"() {foo = 7 : i32} : () -> i32
-    %2 = "test.hello"(%0#0, %1) : (i32, i32) -> i32
-    return %2 : i32
-  }
-}