[Mlir-commits] [mlir] [mlir][arith] Fix bug in `arith.bitcast` canonicalizer (PR #148795)

[Matthias Springer]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/148795

`bitcast(bitcast(x))` was incorrectly folded to `x`.


>From 7bf9cec682f7dd94245550910a7b5d0c661c1de5 Mon Sep 17 00:00:00 2001
From: Matthias Springer <me at m-sp.org>
Date: Tue, 15 Jul 2025 07:24:59 +0000
Subject: [PATCH] [mlir][arith] Fix bug in `arith.bitcast` canonicalizer

---
 .../Dialect/Arith/IR/ArithCanonicalization.td |  2 +-
 mlir/test/Dialect/Arith/canonicalize.mlir     | 23 +++++++++++++++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
index 64eccc76a6642..cb710cff89fb4 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
+++ b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
@@ -315,7 +315,7 @@ def IndexCastUIOfExtUI :
 
 // bitcast(bitcast(x)) -> x
 def BitcastOfBitcast :
-    Pat<(Arith_BitcastOp (Arith_BitcastOp $x)), (replaceWithValue $x)>;
+    Pat<(Arith_BitcastOp (Arith_BitcastOp $x)), (Arith_BitcastOp $x)>;
 
 //===----------------------------------------------------------------------===//
 // ExtSIOp
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index 076f3a99cd393..51caa261375ba 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -1940,6 +1940,29 @@ func.func @bitcastPoisonFPtoI() -> i32 {
 
 // -----
 
+// CHECK-LABEL: func @bitcast_chain_1(
+//  CHECK-SAME:     %[[arg:.*]]: i16)
+//       CHECK:   %[[cast:.*]] = arith.bitcast %[[arg]] : i16 to f16
+//       CHECK:   return %[[cast]]
+func.func @bitcast_chain_1(%arg: i16) -> f16 {
+  %0 = arith.bitcast %arg : i16 to bf16
+  %1 = arith.bitcast %0 : bf16 to f16
+  return %1 : f16
+}
+
+// -----
+
+// CHECK-LABEL: func @bitcast_chain_2(
+//  CHECK-SAME:     %[[arg:.*]]: f16)
+//       CHECK:   return %[[arg]]
+func.func @bitcast_chain_2(%arg: f16) -> f16 {
+  %0 = arith.bitcast %arg : f16 to bf16
+  %1 = arith.bitcast %0 : bf16 to f16
+  return %1 : f16
+}
+
+// -----
+
 // CHECK-LABEL: test_maxsi
 // CHECK-DAG: %[[C0:.+]] = arith.constant 42
 // CHECK-DAG: %[[MAX_INT_CST:.+]] = arith.constant 127