[Mlir-commits] [mlir] [mlir][tosa] Allow i64 accumulators (PR #228427)

[Michael Platings]

https://github.com/mplatings updated https://github.com/llvm/llvm-project/pull/228427

>From 87d8abf9149f7bd636d1e682d84e8a240ffcd376 Mon Sep 17 00:00:00 2001
From: Michael Platings <michael.platings at arm.com>
Date: Fri, 2 Oct 2026 13:43:15 +0100
Subject: [PATCH 1/2] [mlir][tosa] Allow i64 accumulators

This can be helpful for incremental lowerings to/from TOSA.
---
 mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td | 2 +-
 mlir/test/Dialect/Tosa/accumulator-i64.mlir  | 9 +++++++++
 2 files changed, 10 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Dialect/Tosa/accumulator-i64.mlir

diff --git a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
index 8b5803c3074f0..339061b2aee83 100644
--- a/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
+++ b/mlir/include/mlir/Dialect/Tosa/IR/TosaOps.td
@@ -103,7 +103,7 @@ def Tosa_ArgMinOp : Tosa_InferShapedTypeOp<"argmin", [NoMemoryEffect]> {
 // Accumulator types.
 //===----------------------------------------------------------------------===//
 
-def Tosa_AccType : AnyTypeOf<[I<32>, I<48>, F16, F32, BF16]>;
+def Tosa_AccType : AnyTypeOf<[I<32>, I<48>, I<64>, F16, F32, BF16]>;
 
 //===----------------------------------------------------------------------===//
 // Operator: avg_pool2d
diff --git a/mlir/test/Dialect/Tosa/accumulator-i64.mlir b/mlir/test/Dialect/Tosa/accumulator-i64.mlir
new file mode 100644
index 0000000000000..6f1dec31084a5
--- /dev/null
+++ b/mlir/test/Dialect/Tosa/accumulator-i64.mlir
@@ -0,0 +1,9 @@
+// RUN: mlir-opt %s --verify-each | FileCheck %s
+
+// Allow i64 accumulators to facilitate incremental lowerings to/from TOSA.
+// CHECK-LABEL: func.func @conv2d_i64
+// CHECK: acc_type(i64)
+func.func @conv2d_i64(%input: tensor<1x1x1x1xi16>, %weight: tensor<1x1x1x1xi8>, %bias: tensor<1xi64>, %izp: tensor<1xi16>, %wzp: tensor<1xi8>) -> tensor<1x1x1x1xi64> {
+  %0 = tosa.conv2d %input, %weight, %bias, %izp, %wzp pad([0, 0, 0, 0]) stride([1, 1]) dilation([1, 1]) acc_type(i64) : (tensor<1x1x1x1xi16>, tensor<1x1x1x1xi8>, tensor<1xi64>, tensor<1xi16>, tensor<1xi8>) -> tensor<1x1x1x1xi64>
+  return %0 : tensor<1x1x1x1xi64>
+}

>From b00f4d1d914645239162bf18b7e3e24c42893cca Mon Sep 17 00:00:00 2001
From: Michael Platings <michael.platings at arm.com>
Date: Fri, 2 Oct 2026 17:56:09 +0100
Subject: [PATCH 2/2] fixup! [mlir][tosa] Allow i64 accumulators

---
 mlir/test/Dialect/Tosa/accumulator-i64.mlir | 9 ---------
 mlir/test/Dialect/Tosa/ops.mlir             | 9 +++++++++
 2 files changed, 9 insertions(+), 9 deletions(-)
 delete mode 100644 mlir/test/Dialect/Tosa/accumulator-i64.mlir

diff --git a/mlir/test/Dialect/Tosa/accumulator-i64.mlir b/mlir/test/Dialect/Tosa/accumulator-i64.mlir
deleted file mode 100644
index 6f1dec31084a5..0000000000000
--- a/mlir/test/Dialect/Tosa/accumulator-i64.mlir
+++ /dev/null
@@ -1,9 +0,0 @@
-// RUN: mlir-opt %s --verify-each | FileCheck %s
-
-// Allow i64 accumulators to facilitate incremental lowerings to/from TOSA.
-// CHECK-LABEL: func.func @conv2d_i64
-// CHECK: acc_type(i64)
-func.func @conv2d_i64(%input: tensor<1x1x1x1xi16>, %weight: tensor<1x1x1x1xi8>, %bias: tensor<1xi64>, %izp: tensor<1xi16>, %wzp: tensor<1xi8>) -> tensor<1x1x1x1xi64> {
-  %0 = tosa.conv2d %input, %weight, %bias, %izp, %wzp pad([0, 0, 0, 0]) stride([1, 1]) dilation([1, 1]) acc_type(i64) : (tensor<1x1x1x1xi16>, tensor<1x1x1x1xi8>, tensor<1xi64>, tensor<1xi16>, tensor<1xi8>) -> tensor<1x1x1x1xi64>
-  return %0 : tensor<1x1x1x1xi64>
-}
diff --git a/mlir/test/Dialect/Tosa/ops.mlir b/mlir/test/Dialect/Tosa/ops.mlir
index 2592b08b35d39..9f2b7f1e5ee1e 100644
--- a/mlir/test/Dialect/Tosa/ops.mlir
+++ b/mlir/test/Dialect/Tosa/ops.mlir
@@ -149,6 +149,15 @@ func.func @test_conv2d(%arg0: tensor<1x4x4x4xf32>, %arg1: tensor<8x1x1x4xf32>, %
   return %0 : tensor<1x4x4x8xf32>
 }
 
+// -----
+// Allow i64 accumulators to facilitate incremental lowerings to/from TOSA.
+// CHECK-LABEL: conv2d_i64
+// CHECK: acc_type(i64)
+func.func @test_conv2d_i64(%input: tensor<1x1x1x1xi16>, %weight: tensor<1x1x1x1xi8>, %bias: tensor<1xi64>, %izp: tensor<1xi16>, %wzp: tensor<1xi8>) -> tensor<1x1x1x1xi64> {
+  %0 = tosa.conv2d %input, %weight, %bias, %izp, %wzp pad([0, 0, 0, 0]) stride([1, 1]) dilation([1, 1]) acc_type(i64) : (tensor<1x1x1x1xi16>, tensor<1x1x1x1xi8>, tensor<1xi64>, tensor<1xi16>, tensor<1xi8>) -> tensor<1x1x1x1xi64>
+  return %0 : tensor<1x1x1x1xi64>
+}
+
 // -----
 // CHECK-LABEL: conv2d_mxfp
 func.func @test_conv2d_mxfp(%arg0: tensor<1x4x4x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, %arg1: tensor<8x1x1x32x!tosa.block_scaled<BLOCK_SHAPE_32:f8E8M0FNU:f4E2M1FN>>, %arg2: tensor<8xf16>) -> tensor<1x4x4x8xf16> {