[Mlir-commits] [mlir] [mlir][tosa] Add a pass to downgrade TOSA `1.1.draft` to `1.0` (PR #194971)
Luke Hutton
llvmlistbot at llvm.org
Wed May 6 12:53:53 PDT 2026
https://github.com/lhutton1 updated https://github.com/llvm/llvm-project/pull/194971
>From 79750b4a9f5d1111455581cf455072b6f586819e Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Tue, 28 Apr 2026 17:41:24 +0100
Subject: [PATCH 1/2] [mlir][tosa] Add a pass to downgrade TOSA `1.1.draft` to
`1.0`
This commit adds a pass that will allow 1.1.draft operations
to be rewritten to their 1.0 counterparts where possible.
The pass currently covers the following operations:
- avg/max_pool2d_adaptive -> avg/max_pool2d when kernel/stride/pad are
direct const_shape operands
- bool <-> fp32 casts via i8 bridge casts
- bool gather/scatter with i32 indices via i8 payload rewrites
Note that the downgrade is 'best-effort' and the pass does not perform
any validation itself. The validation pass should be run after
downgrading to check that the resulting IR was downgraded successfully.
Motivation: This decouples the target specification version in
legalizations and backends. Legalizations from higher level frameworks
may be updated to support producing TOSA 1.1.draft variants of operations,
while backends can still consume TOSA 1.0 IR after running the downgrade
pass.
Change-Id: I4c16c099b83827ffb264270e4d7afed9f22e6398
---
.../mlir/Dialect/Tosa/Transforms/Passes.td | 11 +
.../Dialect/Tosa/Transforms/CMakeLists.txt | 1 +
.../Tosa/Transforms/TosaDowngrade1p1To1p0.cpp | 217 ++++++++++++++++++
.../Tosa/tosa-downgrade-1-1-to-1-0.mlir | 161 +++++++++++++
4 files changed, 390 insertions(+)
create mode 100644 mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
create mode 100644 mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir
diff --git a/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td b/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
index 5979ce4962e55..005cbfab782df 100644
--- a/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
@@ -185,6 +185,17 @@ def TosaAttachTarget : Pass<"tosa-attach-target", "ModuleOp"> {
];
}
+def TosaDowngrade1p1To1p0Pass
+ : Pass<"tosa-downgrade-1-1-to-1-0", "func::FuncOp"> {
+ let summary = "Downgrade TOSA 1.1 specification constructs to TOSA 1.0";
+ let description = [{
+ Rewrites constructs which are only compatible in TOSA specification 1.1 and
+ above to their TOSA 1.0 counterparts where possible. Downgrading is best-effort
+ and validation should be performed afterwards to ensure compatibility with
+ the TOSA 1.0 specification.
+ }];
+}
+
def TosaNarrowI64ToI32Pass : Pass<"tosa-narrow-i64-to-i32", "func::FuncOp"> {
let summary = "Narrow I64 TOSA operations to I32";
let description = [{
diff --git a/mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt
index e8a76fa3a1d21..1fd18bb5a395b 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt
@@ -2,6 +2,7 @@ add_mlir_dialect_library(MLIRTosaTransforms
TosaAttachTarget.cpp
TosaArithConstantToConst.cpp
TosaConvertIntegerTypeToSignless.cpp
+ TosaDowngrade1p1To1p0.cpp
TosaDecomposeTransposeConv.cpp
TosaDecomposeDepthwise.cpp
TosaFolders.cpp
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
new file mode 100644
index 0000000000000..d48c4fbdf6dd1
--- /dev/null
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
@@ -0,0 +1,217 @@
+//===- TosaDowngrade1_1To1_0.cpp -----------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Rewrites constructs which are only compatible in TOSA specification 1.1 and
+// above to their TOSA 1.0 counterparts where possible. Downgrading is
+// best-effort and validation should be performed afterwards to ensure
+// compatibility with the TOSA 1.0 specification.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Tosa/Transforms/Passes.h"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tosa/IR/TargetEnv.h"
+#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir {
+namespace tosa {
+#define GEN_PASS_DEF_TOSADOWNGRADE1P1TO1P0PASS
+#include "mlir/Dialect/Tosa/Transforms/Passes.h.inc"
+} // namespace tosa
+} // namespace mlir
+
+using namespace mlir;
+using namespace mlir::tosa;
+
+namespace {
+
+class AvgPool2dAdaptiveToAvgPool2d
+ : public OpRewritePattern<tosa::AvgPool2dAdaptiveOp> {
+public:
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tosa::AvgPool2dAdaptiveOp op,
+ PatternRewriter &rewriter) const override {
+ llvm::SmallVector<int64_t> kernel;
+ llvm::SmallVector<int64_t> stride;
+ llvm::SmallVector<int64_t> pad;
+ if (!tosa::getConstShapeValues(op.getKernel().getDefiningOp(), kernel) ||
+ !tosa::getConstShapeValues(op.getStride().getDefiningOp(), stride) ||
+ !tosa::getConstShapeValues(op.getPad().getDefiningOp(), pad))
+ return rewriter.notifyMatchFailure(
+ op, "expected constant kernel, stride, and pad operands");
+
+ auto replacement = tosa::AvgPool2dOp::create(
+ rewriter, op.getLoc(), op.getType(), op.getInput(), op.getInputZp(),
+ op.getOutputZp(), rewriter.getDenseI64ArrayAttr(kernel),
+ rewriter.getDenseI64ArrayAttr(stride),
+ rewriter.getDenseI64ArrayAttr(pad), op.getAccTypeAttr());
+ rewriter.replaceOp(op, replacement.getOutput());
+ return success();
+ }
+};
+
+class MaxPool2dAdaptiveToMaxPool2d
+ : public OpRewritePattern<tosa::MaxPool2dAdaptiveOp> {
+public:
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tosa::MaxPool2dAdaptiveOp op,
+ PatternRewriter &rewriter) const override {
+ llvm::SmallVector<int64_t> kernel;
+ llvm::SmallVector<int64_t> stride;
+ llvm::SmallVector<int64_t> pad;
+ if (!tosa::getConstShapeValues(op.getKernel().getDefiningOp(), kernel) ||
+ !tosa::getConstShapeValues(op.getStride().getDefiningOp(), stride) ||
+ !tosa::getConstShapeValues(op.getPad().getDefiningOp(), pad))
+ return rewriter.notifyMatchFailure(
+ op, "expected constant kernel, stride, and pad operands");
+
+ auto replacement = tosa::MaxPool2dOp::create(
+ rewriter, op.getLoc(), op.getType(), op.getInput(),
+ rewriter.getDenseI64ArrayAttr(kernel),
+ rewriter.getDenseI64ArrayAttr(stride),
+ rewriter.getDenseI64ArrayAttr(pad), op.getNanModeAttr());
+ rewriter.replaceOp(op, replacement.getOutput());
+ return success();
+ }
+};
+
+class BoolFp32CastRewrite : public OpRewritePattern<tosa::CastOp> {
+public:
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tosa::CastOp op,
+ PatternRewriter &rewriter) const override {
+ const Value input = op.getInput();
+
+ const Type i1Type = rewriter.getI1Type();
+ const Type f32Type = rewriter.getF32Type();
+
+ const Type inputElemType = getElementTypeOrSelf(input.getType());
+ const Type outputElemType = getElementTypeOrSelf(op.getType());
+ const bool isFp32ToBool =
+ inputElemType == f32Type && outputElemType == i1Type;
+ const bool isBoolToFp32 =
+ inputElemType == i1Type && outputElemType == f32Type;
+
+ if (!isFp32ToBool && !isBoolToFp32)
+ return rewriter.notifyMatchFailure(op,
+ "expected cast between bool and f32");
+
+ const Type outputType = op.getType();
+ const Type i8Type = rewriter.getI8Type();
+ const Type intermediateType = cast<TensorType>(outputType).clone(i8Type);
+
+ auto inner =
+ tosa::CastOp::create(rewriter, op.getLoc(), intermediateType, input);
+ auto outer = tosa::CastOp::create(rewriter, op.getLoc(), outputType,
+ inner.getOutput());
+ rewriter.replaceOp(op, outer.getOutput());
+ return success();
+ }
+};
+
+class BoolGatherRewrite : public OpRewritePattern<tosa::GatherOp> {
+public:
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tosa::GatherOp op,
+ PatternRewriter &rewriter) const override {
+ const Value values = op.getValues();
+ const Value indices = op.getIndices();
+
+ const Type valuesType = values.getType();
+ const Type resultType = op.getType();
+
+ const Type i1Type = rewriter.getI1Type();
+ const Type i32Type = rewriter.getI32Type();
+ if (getElementTypeOrSelf(valuesType) != i1Type ||
+ getElementTypeOrSelf(indices.getType()) != i32Type)
+ return rewriter.notifyMatchFailure(
+ op, "expected values of bool type and indices of i32 type");
+
+ const Type i8Type = rewriter.getI8Type();
+ const Type valuesI8Type = cast<TensorType>(valuesType).clone(i8Type);
+ const Type resultI8Type = cast<TensorType>(resultType).clone(i8Type);
+
+ auto valuesToI8 =
+ tosa::CastOp::create(rewriter, op.getLoc(), valuesI8Type, values);
+ auto gatherI8 = tosa::GatherOp::create(rewriter, op.getLoc(), resultI8Type,
+ valuesToI8.getOutput(), indices);
+ auto i8ToBool = tosa::CastOp::create(rewriter, op.getLoc(), resultType,
+ gatherI8.getOutput());
+ rewriter.replaceOp(op, i8ToBool.getOutput());
+ return success();
+ }
+};
+
+class BoolScatterRewrite : public OpRewritePattern<tosa::ScatterOp> {
+public:
+ using OpRewritePattern::OpRewritePattern;
+
+ LogicalResult matchAndRewrite(tosa::ScatterOp op,
+ PatternRewriter &rewriter) const override {
+ const Value valuesIn = op.getValuesIn();
+ const Value indices = op.getIndices();
+
+ const Type valuesInType = valuesIn.getType();
+ const Type i1Type = rewriter.getI1Type();
+ const Type i32Type = rewriter.getI32Type();
+ if (getElementTypeOrSelf(valuesInType) != i1Type ||
+ getElementTypeOrSelf(indices.getType()) != i32Type)
+ return rewriter.notifyMatchFailure(
+ op, "expected values of bool type and indices of i32 type");
+
+ const Value input = op.getInput();
+ const Type inputType = input.getType();
+ const Type resultType = op.getType();
+
+ const Type i8Type = rewriter.getI8Type();
+ const Type valuesInI8Type = cast<TensorType>(valuesInType).clone(i8Type);
+ const Type inputI8Type = cast<TensorType>(inputType).clone(i8Type);
+ const Type resultI8Type = cast<TensorType>(resultType).clone(i8Type);
+
+ auto valuesInToI8 =
+ tosa::CastOp::create(rewriter, op.getLoc(), valuesInI8Type, valuesIn);
+ auto inputToI8 =
+ tosa::CastOp::create(rewriter, op.getLoc(), inputI8Type, input);
+ auto scatterI8 = tosa::ScatterOp::create(
+ rewriter, op.getLoc(), resultI8Type, valuesInToI8.getOutput(), indices,
+ inputToI8.getOutput());
+ auto i8ToBool = tosa::CastOp::create(rewriter, op.getLoc(), resultType,
+ scatterI8.getValuesOut());
+ rewriter.replaceOp(op, i8ToBool.getOutput());
+ return success();
+ }
+};
+
+struct TosaDowngrade1p1To1p0Pass
+ : public tosa::impl::TosaDowngrade1p1To1p0PassBase<
+ TosaDowngrade1p1To1p0Pass> {
+ using Base::Base;
+
+ void runOnOperation() override {
+ MLIRContext &context = getContext();
+ func::FuncOp func = getOperation();
+
+ RewritePatternSet patterns(&context);
+ patterns.add<AvgPool2dAdaptiveToAvgPool2d, MaxPool2dAdaptiveToMaxPool2d,
+ BoolFp32CastRewrite, BoolGatherRewrite, BoolScatterRewrite>(
+ &context);
+ FrozenRewritePatternSet frozenPatterns(std::move(patterns));
+
+ if (failed(applyPatternsGreedily(func, frozenPatterns)))
+ return signalPassFailure();
+ }
+};
+
+} // namespace
diff --git a/mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir b/mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir
new file mode 100644
index 0000000000000..44bc7a63f9558
--- /dev/null
+++ b/mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir
@@ -0,0 +1,161 @@
+// RUN: mlir-opt --split-input-file --tosa-downgrade-1-1-to-1-0 %s | FileCheck %s
+
+// CHECK-LABEL: @test_max_pool
+// CHECK: %[[POOL:.+]] = tosa.max_pool2d %arg0 {kernel = array<i64: 1, 1>, nan_mode = IGNORE, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32>
+// CHECK: return %[[POOL]]
+func.func @test_max_pool(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32> {
+ %kernel = tosa.const_shape {values = dense<[1, 1]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %stride = tosa.const_shape {values = dense<[1, 1]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %pad = tosa.const_shape {values = dense<[0, 0, 0, 0]> : tensor<4xindex>} : () -> !tosa.shape<4>
+ %0 = tosa.max_pool2d_adaptive %arg0, %kernel, %stride, %pad {nan_mode = IGNORE} :
+ (tensor<1x32x32x8xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x32x32x8xf32>
+ return %0 : tensor<1x32x32x8xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @test_avg_pool
+// CHECK: %[[POOL:.+]] = tosa.avg_pool2d %arg0, %{{.*}}, %{{.*}} {acc_type = f32, kernel = array<i64: 3, 3>, pad = array<i64: 1, 1, 1, 1>, stride = array<i64: 1, 1>} : (tensor<1x7x7x9xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x7x7x9xf32>
+// CHECK: return %[[POOL]]
+func.func @test_avg_pool(%arg0: tensor<1x7x7x9xf32>, %input_zp: tensor<1xf32>, %output_zp: tensor<1xf32>) -> tensor<1x7x7x9xf32> {
+ %kernel = tosa.const_shape {values = dense<[3, 3]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %stride = tosa.const_shape {values = dense<[1, 1]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %pad = tosa.const_shape {values = dense<[1, 1, 1, 1]> : tensor<4xindex>} : () -> !tosa.shape<4>
+ %0 = tosa.avg_pool2d_adaptive %arg0, %input_zp, %output_zp, %kernel, %stride, %pad {acc_type = f32} :
+ (tensor<1x7x7x9xf32>, tensor<1xf32>, tensor<1xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x7x7x9xf32>
+ return %0 : tensor<1x7x7x9xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @test_bool_to_fp32
+// CHECK: %[[BOOL_TO_I8:.+]] = tosa.cast %arg0 : (tensor<13x21x3xi1>) -> tensor<13x21x3xi8>
+// CHECK: %[[I8_TO_F32:.+]] = tosa.cast %[[BOOL_TO_I8]] : (tensor<13x21x3xi8>) -> tensor<13x21x3xf32>
+// CHECK: return %[[I8_TO_F32]]
+func.func @test_bool_to_fp32(%arg0: tensor<13x21x3xi1>) -> tensor<13x21x3xf32> {
+ %0 = tosa.cast %arg0 : (tensor<13x21x3xi1>) -> tensor<13x21x3xf32>
+ return %0 : tensor<13x21x3xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @test_bool_to_fp32_unranked
+// CHECK: %[[BOOL_TO_I8:.+]] = tosa.cast %arg0 : (tensor<*xi1>) -> tensor<*xi8>
+// CHECK: %[[I8_TO_F32:.+]] = tosa.cast %[[BOOL_TO_I8]] : (tensor<*xi8>) -> tensor<*xf32>
+// CHECK: return %[[I8_TO_F32]]
+func.func @test_bool_to_fp32_unranked(%arg0: tensor<*xi1>) -> tensor<*xf32> {
+ %0 = tosa.cast %arg0 : (tensor<*xi1>) -> tensor<*xf32>
+ return %0 : tensor<*xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @test_fp32_to_bool_ranked_dynamic
+// CHECK: %[[FP32_TO_I8:.+]] = tosa.cast %arg0 : (tensor<13x?x3xf32>) -> tensor<13x?x3xi8>
+// CHECK: %[[I8_TO_BOOL:.+]] = tosa.cast %[[FP32_TO_I8]] : (tensor<13x?x3xi8>) -> tensor<13x?x3xi1>
+// CHECK: return %[[I8_TO_BOOL]]
+func.func @test_fp32_to_bool_ranked_dynamic(%arg0: tensor<13x?x3xf32>) -> tensor<13x?x3xi1> {
+ %0 = tosa.cast %arg0 : (tensor<13x?x3xf32>) -> tensor<13x?x3xi1>
+ return %0 : tensor<13x?x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @test_unranked_fp32_to_bool
+// CHECK: %[[FP32_TO_I8:.+]] = tosa.cast %arg0 : (tensor<*xf32>) -> tensor<*xi8>
+// CHECK: %[[I8_TO_BOOL:.+]] = tosa.cast %[[FP32_TO_I8]] : (tensor<*xi8>) -> tensor<*xi1>
+// CHECK: return %[[I8_TO_BOOL]]
+func.func @test_unranked_fp32_to_bool(%arg0: tensor<*xf32>) -> tensor<*xi1> {
+ %0 = tosa.cast %arg0 : (tensor<*xf32>) -> tensor<*xi1>
+ return %0 : tensor<*xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @test_preserve_bool_to_i8
+// CHECK: %[[CAST:.+]] = tosa.cast %arg0 : (tensor<13x21x3xi1>) -> tensor<13x21x3xi8>
+// CHECK: return %[[CAST]]
+func.func @test_preserve_bool_to_i8(%arg0: tensor<13x21x3xi1>) -> tensor<13x21x3xi8> {
+ %0 = tosa.cast %arg0 : (tensor<13x21x3xi1>) -> tensor<13x21x3xi8>
+ return %0 : tensor<13x21x3xi8>
+}
+
+// -----
+
+// CHECK-LABEL: @test_gather_bool_i32
+// CHECK: %[[VALUES_TO_I8:.+]] = tosa.cast %arg0 : (tensor<13x21x3xi1>) -> tensor<13x21x3xi8>
+// CHECK: %[[GATHER_I8:.+]] = tosa.gather %[[VALUES_TO_I8]], %arg1 : (tensor<13x21x3xi8>, tensor<13x26xi32>) -> tensor<13x26x3xi8>
+// CHECK: %[[I8_TO_BOOL:.+]] = tosa.cast %[[GATHER_I8]] : (tensor<13x26x3xi8>) -> tensor<13x26x3xi1>
+// CHECK: return %[[I8_TO_BOOL]]
+func.func @test_gather_bool_i32(%arg0: tensor<13x21x3xi1>, %arg1: tensor<13x26xi32>) -> tensor<13x26x3xi1> {
+ %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x3xi1>, tensor<13x26xi32>) -> tensor<13x26x3xi1>
+ return %0 : tensor<13x26x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @test_preserve_gather_bool_i64
+// CHECK: %[[GATHER:.+]] = tosa.gather %arg0, %arg1 : (tensor<13x21x3xi1>, tensor<13x26xi64>) -> tensor<13x26x3xi1>
+// CHECK: return %[[GATHER]]
+func.func @test_preserve_gather_bool_i64(%arg0: tensor<13x21x3xi1>, %arg1: tensor<13x26xi64>) -> tensor<13x26x3xi1> {
+ %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x3xi1>, tensor<13x26xi64>) -> tensor<13x26x3xi1>
+ return %0 : tensor<13x26x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @test_preserve_gather_i8_i32
+// CHECK: %[[GATHER:.+]] = tosa.gather %arg0, %arg1 : (tensor<13x21x3xi8>, tensor<13x26xi32>) -> tensor<13x26x3xi8>
+// CHECK: return %[[GATHER]]
+func.func @test_preserve_gather_i8_i32(%arg0: tensor<13x21x3xi8>, %arg1: tensor<13x26xi32>) -> tensor<13x26x3xi8> {
+ %0 = tosa.gather %arg0, %arg1 : (tensor<13x21x3xi8>, tensor<13x26xi32>) -> tensor<13x26x3xi8>
+ return %0 : tensor<13x26x3xi8>
+}
+
+// -----
+
+// CHECK-LABEL: @test_scatter_bool_i32
+// CHECK: %[[VALUES_IN_TO_I8:.+]] = tosa.cast %arg0 : (tensor<13x52x3xi1>) -> tensor<13x52x3xi8>
+// CHECK: %[[INPUT_TO_I8:.+]] = tosa.cast %arg2 : (tensor<13x26x3xi1>) -> tensor<13x26x3xi8>
+// CHECK: %[[SCATTER_I8:.+]] = tosa.scatter %[[VALUES_IN_TO_I8]], %arg1, %[[INPUT_TO_I8]] : (tensor<13x52x3xi8>, tensor<13x26xi32>, tensor<13x26x3xi8>) -> tensor<13x52x3xi8>
+// CHECK: %[[I8_TO_BOOL:.+]] = tosa.cast %[[SCATTER_I8]] : (tensor<13x52x3xi8>) -> tensor<13x52x3xi1>
+// CHECK: return %[[I8_TO_BOOL]]
+func.func @test_scatter_bool_i32(%arg0: tensor<13x52x3xi1>, %arg1: tensor<13x26xi32>, %arg2: tensor<13x26x3xi1>) -> tensor<13x52x3xi1> {
+ %0 = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x52x3xi1>, tensor<13x26xi32>, tensor<13x26x3xi1>) -> tensor<13x52x3xi1>
+ return %0 : tensor<13x52x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @test_preserve_scatter_bool_i64
+// CHECK: %[[SCATTER:.+]] = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x52x3xi1>, tensor<13x26xi64>, tensor<13x26x3xi1>) -> tensor<13x52x3xi1>
+// CHECK: return %[[SCATTER]]
+func.func @test_preserve_scatter_bool_i64(%arg0: tensor<13x52x3xi1>, %arg1: tensor<13x26xi64>, %arg2: tensor<13x26x3xi1>) -> tensor<13x52x3xi1> {
+ %0 = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x52x3xi1>, tensor<13x26xi64>, tensor<13x26x3xi1>) -> tensor<13x52x3xi1>
+ return %0 : tensor<13x52x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @test_preserve_scatter_i8_i32
+// CHECK: %[[SCATTER:.+]] = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x52x3xi8>, tensor<13x26xi32>, tensor<13x26x3xi8>) -> tensor<13x52x3xi8>
+// CHECK: return %[[SCATTER]]
+func.func @test_preserve_scatter_i8_i32(%arg0: tensor<13x52x3xi8>, %arg1: tensor<13x26xi32>, %arg2: tensor<13x26x3xi8>) -> tensor<13x52x3xi8> {
+ %0 = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x52x3xi8>, tensor<13x26xi32>, tensor<13x26x3xi8>) -> tensor<13x52x3xi8>
+ return %0 : tensor<13x52x3xi8>
+}
+
+// -----
+
+// CHECK-LABEL: @test_preserve_non_const_adaptive
+// CHECK: tosa.max_pool2d_adaptive
+func.func @test_preserve_non_const_adaptive(%arg0: tensor<1x?x?x8xf32>) -> tensor<1x?x?x8xf32> {
+ %dim1 = tosa.dim %arg0 {axis = 1 : i32} : (tensor<1x?x?x8xf32>) -> !tosa.shape<1>
+ %dim2 = tosa.dim %arg0 {axis = 2 : i32} : (tensor<1x?x?x8xf32>) -> !tosa.shape<1>
+ %kernel = tosa.concat_shape %dim1, %dim2 : (!tosa.shape<1>, !tosa.shape<1>) -> !tosa.shape<2>
+ %stride = tosa.const_shape {values = dense<[1, 1]> : tensor<2xindex>} : () -> !tosa.shape<2>
+ %pad = tosa.const_shape {values = dense<[0, 0, 0, 0]> : tensor<4xindex>} : () -> !tosa.shape<4>
+ %0 = tosa.max_pool2d_adaptive %arg0, %kernel, %stride, %pad {nan_mode = IGNORE} :
+ (tensor<1x?x?x8xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x?x?x8xf32>
+ return %0 : tensor<1x?x?x8xf32>
+}
>From f17f30ed3a3cf583161b6860f7e09bd7b4dc7e54 Mon Sep 17 00:00:00 2001
From: Luke Hutton <luke.hutton at arm.com>
Date: Wed, 6 May 2026 16:50:25 +0100
Subject: [PATCH 2/2] Remove adaptive pooling patterns
These have been moved to a canonicalization pattern.
Change-Id: I2a7f6ad4eb9e7fca0266c2bc3179da5f490cd0ab
---
.../Tosa/Transforms/TosaDowngrade1p1To1p0.cpp | 57 +------------------
.../Tosa/tosa-downgrade-1-1-to-1-0.mlir | 43 --------------
2 files changed, 1 insertion(+), 99 deletions(-)
diff --git a/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp b/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
index d48c4fbdf6dd1..cfd2dd9d29650 100644
--- a/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
+++ b/mlir/lib/Dialect/Tosa/Transforms/TosaDowngrade1p1To1p0.cpp
@@ -16,8 +16,6 @@
#include "mlir/Dialect/Tosa/Transforms/Passes.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/Tosa/IR/TargetEnv.h"
-#include "mlir/Dialect/Tosa/Utils/ConversionUtils.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -33,58 +31,6 @@ using namespace mlir::tosa;
namespace {
-class AvgPool2dAdaptiveToAvgPool2d
- : public OpRewritePattern<tosa::AvgPool2dAdaptiveOp> {
-public:
- using OpRewritePattern::OpRewritePattern;
-
- LogicalResult matchAndRewrite(tosa::AvgPool2dAdaptiveOp op,
- PatternRewriter &rewriter) const override {
- llvm::SmallVector<int64_t> kernel;
- llvm::SmallVector<int64_t> stride;
- llvm::SmallVector<int64_t> pad;
- if (!tosa::getConstShapeValues(op.getKernel().getDefiningOp(), kernel) ||
- !tosa::getConstShapeValues(op.getStride().getDefiningOp(), stride) ||
- !tosa::getConstShapeValues(op.getPad().getDefiningOp(), pad))
- return rewriter.notifyMatchFailure(
- op, "expected constant kernel, stride, and pad operands");
-
- auto replacement = tosa::AvgPool2dOp::create(
- rewriter, op.getLoc(), op.getType(), op.getInput(), op.getInputZp(),
- op.getOutputZp(), rewriter.getDenseI64ArrayAttr(kernel),
- rewriter.getDenseI64ArrayAttr(stride),
- rewriter.getDenseI64ArrayAttr(pad), op.getAccTypeAttr());
- rewriter.replaceOp(op, replacement.getOutput());
- return success();
- }
-};
-
-class MaxPool2dAdaptiveToMaxPool2d
- : public OpRewritePattern<tosa::MaxPool2dAdaptiveOp> {
-public:
- using OpRewritePattern::OpRewritePattern;
-
- LogicalResult matchAndRewrite(tosa::MaxPool2dAdaptiveOp op,
- PatternRewriter &rewriter) const override {
- llvm::SmallVector<int64_t> kernel;
- llvm::SmallVector<int64_t> stride;
- llvm::SmallVector<int64_t> pad;
- if (!tosa::getConstShapeValues(op.getKernel().getDefiningOp(), kernel) ||
- !tosa::getConstShapeValues(op.getStride().getDefiningOp(), stride) ||
- !tosa::getConstShapeValues(op.getPad().getDefiningOp(), pad))
- return rewriter.notifyMatchFailure(
- op, "expected constant kernel, stride, and pad operands");
-
- auto replacement = tosa::MaxPool2dOp::create(
- rewriter, op.getLoc(), op.getType(), op.getInput(),
- rewriter.getDenseI64ArrayAttr(kernel),
- rewriter.getDenseI64ArrayAttr(stride),
- rewriter.getDenseI64ArrayAttr(pad), op.getNanModeAttr());
- rewriter.replaceOp(op, replacement.getOutput());
- return success();
- }
-};
-
class BoolFp32CastRewrite : public OpRewritePattern<tosa::CastOp> {
public:
using OpRewritePattern::OpRewritePattern;
@@ -204,8 +150,7 @@ struct TosaDowngrade1p1To1p0Pass
func::FuncOp func = getOperation();
RewritePatternSet patterns(&context);
- patterns.add<AvgPool2dAdaptiveToAvgPool2d, MaxPool2dAdaptiveToMaxPool2d,
- BoolFp32CastRewrite, BoolGatherRewrite, BoolScatterRewrite>(
+ patterns.add<BoolFp32CastRewrite, BoolGatherRewrite, BoolScatterRewrite>(
&context);
FrozenRewritePatternSet frozenPatterns(std::move(patterns));
diff --git a/mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir b/mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir
index 44bc7a63f9558..5427d9119b7af 100644
--- a/mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir
+++ b/mlir/test/Dialect/Tosa/tosa-downgrade-1-1-to-1-0.mlir
@@ -1,33 +1,5 @@
// RUN: mlir-opt --split-input-file --tosa-downgrade-1-1-to-1-0 %s | FileCheck %s
-// CHECK-LABEL: @test_max_pool
-// CHECK: %[[POOL:.+]] = tosa.max_pool2d %arg0 {kernel = array<i64: 1, 1>, nan_mode = IGNORE, pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>} : (tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32>
-// CHECK: return %[[POOL]]
-func.func @test_max_pool(%arg0: tensor<1x32x32x8xf32>) -> tensor<1x32x32x8xf32> {
- %kernel = tosa.const_shape {values = dense<[1, 1]> : tensor<2xindex>} : () -> !tosa.shape<2>
- %stride = tosa.const_shape {values = dense<[1, 1]> : tensor<2xindex>} : () -> !tosa.shape<2>
- %pad = tosa.const_shape {values = dense<[0, 0, 0, 0]> : tensor<4xindex>} : () -> !tosa.shape<4>
- %0 = tosa.max_pool2d_adaptive %arg0, %kernel, %stride, %pad {nan_mode = IGNORE} :
- (tensor<1x32x32x8xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x32x32x8xf32>
- return %0 : tensor<1x32x32x8xf32>
-}
-
-// -----
-
-// CHECK-LABEL: @test_avg_pool
-// CHECK: %[[POOL:.+]] = tosa.avg_pool2d %arg0, %{{.*}}, %{{.*}} {acc_type = f32, kernel = array<i64: 3, 3>, pad = array<i64: 1, 1, 1, 1>, stride = array<i64: 1, 1>} : (tensor<1x7x7x9xf32>, tensor<1xf32>, tensor<1xf32>) -> tensor<1x7x7x9xf32>
-// CHECK: return %[[POOL]]
-func.func @test_avg_pool(%arg0: tensor<1x7x7x9xf32>, %input_zp: tensor<1xf32>, %output_zp: tensor<1xf32>) -> tensor<1x7x7x9xf32> {
- %kernel = tosa.const_shape {values = dense<[3, 3]> : tensor<2xindex>} : () -> !tosa.shape<2>
- %stride = tosa.const_shape {values = dense<[1, 1]> : tensor<2xindex>} : () -> !tosa.shape<2>
- %pad = tosa.const_shape {values = dense<[1, 1, 1, 1]> : tensor<4xindex>} : () -> !tosa.shape<4>
- %0 = tosa.avg_pool2d_adaptive %arg0, %input_zp, %output_zp, %kernel, %stride, %pad {acc_type = f32} :
- (tensor<1x7x7x9xf32>, tensor<1xf32>, tensor<1xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x7x7x9xf32>
- return %0 : tensor<1x7x7x9xf32>
-}
-
-// -----
-
// CHECK-LABEL: @test_bool_to_fp32
// CHECK: %[[BOOL_TO_I8:.+]] = tosa.cast %arg0 : (tensor<13x21x3xi1>) -> tensor<13x21x3xi8>
// CHECK: %[[I8_TO_F32:.+]] = tosa.cast %[[BOOL_TO_I8]] : (tensor<13x21x3xi8>) -> tensor<13x21x3xf32>
@@ -144,18 +116,3 @@ func.func @test_preserve_scatter_i8_i32(%arg0: tensor<13x52x3xi8>, %arg1: tensor
%0 = tosa.scatter %arg0, %arg1, %arg2 : (tensor<13x52x3xi8>, tensor<13x26xi32>, tensor<13x26x3xi8>) -> tensor<13x52x3xi8>
return %0 : tensor<13x52x3xi8>
}
-
-// -----
-
-// CHECK-LABEL: @test_preserve_non_const_adaptive
-// CHECK: tosa.max_pool2d_adaptive
-func.func @test_preserve_non_const_adaptive(%arg0: tensor<1x?x?x8xf32>) -> tensor<1x?x?x8xf32> {
- %dim1 = tosa.dim %arg0 {axis = 1 : i32} : (tensor<1x?x?x8xf32>) -> !tosa.shape<1>
- %dim2 = tosa.dim %arg0 {axis = 2 : i32} : (tensor<1x?x?x8xf32>) -> !tosa.shape<1>
- %kernel = tosa.concat_shape %dim1, %dim2 : (!tosa.shape<1>, !tosa.shape<1>) -> !tosa.shape<2>
- %stride = tosa.const_shape {values = dense<[1, 1]> : tensor<2xindex>} : () -> !tosa.shape<2>
- %pad = tosa.const_shape {values = dense<[0, 0, 0, 0]> : tensor<4xindex>} : () -> !tosa.shape<4>
- %0 = tosa.max_pool2d_adaptive %arg0, %kernel, %stride, %pad {nan_mode = IGNORE} :
- (tensor<1x?x?x8xf32>, !tosa.shape<2>, !tosa.shape<2>, !tosa.shape<4>) -> tensor<1x?x?x8xf32>
- return %0 : tensor<1x?x?x8xf32>
-}
More information about the Mlir-commits
mailing list