[Mlir-commits] [mlir] [mlir][spirv][tosa] Add remaining TOSA 1.0 SPIR-V TOSA ops (PR #200383)
Davide Grohmann
llvmlistbot at llvm.org
Fri May 29 06:59:55 PDT 2026
https://github.com/davidegrohmann updated https://github.com/llvm/llvm-project/pull/200383
>From 43484aaaa8b6f1fcf674ae247b9cc1cfaa7a2d99 Mon Sep 17 00:00:00 2001
From: Davide Grohmann <davide.grohmann at arm.com>
Date: Thu, 28 May 2026 20:28:41 +0200
Subject: [PATCH] [mlir][spirv][tosa] Add remaining TOSA 1.0 SPIR-V TOSA ops
Add conversion patterns for additional TOSA 1.0 operations targeting
the SPIR-V TOSA extended instruction set.
This covers pooling and convolution ops, FFT/RFFT, matmul, concat,
pad, rescale, const, const_shape, and identity. Concat is split into
conservative chunks to avoid producing SPIR-V instructions with too
many operands.
Add a multi-result conversion pattern for FFT/RFFT and share the
convolution replacement logic for conv2d, conv3d, and depthwise_conv2d
while keeping transpose_conv2d explicit because it has different
attributes.
Also share constant attribute conversion for const and const_shape,
including integer element type conversions such as index to i32, i4 to
i8, and i48 to i64, and preserve the empty const_shape edge case.
Add conversion tests for the newly covered operations.
Signed-off-by: Davide Grohmann <davide.grohmann at arm.com>
Change-Id: Idd4b7b5e7cb289f5e87bc91a0b324cc3a4c343d4
---
.../TosaToSPIRVTosa/TosaToSPIRVTosaOps.cpp | 285 ++++++++++++++++-
.../TosaToSPIRVTosa/tosa-to-spirv.mlir | 288 ++++++++++++++++++
2 files changed, 563 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaOps.cpp b/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaOps.cpp
index 948a0b277cd86..16ea459d5edbe 100644
--- a/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaOps.cpp
+++ b/mlir/lib/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosaOps.cpp
@@ -20,6 +20,8 @@
namespace mlir::tosa {
namespace {
+constexpr unsigned maxConcatOpInputs = 64;
+
template <typename OpAdaptor>
spirv::TosaExtNaNPropagationModeType getNanMode(OpAdaptor adaptor) {
return static_cast<spirv::TosaExtNaNPropagationModeType>(
@@ -31,6 +33,23 @@ spirv::TosaExtResizeModeType getResizeMode(OpAdaptor adaptor) {
return static_cast<spirv::TosaExtResizeModeType>(adaptor.getMode());
}
+template <typename OpAdaptor>
+spirv::TosaExtRoundingModeType getRoundingMode(OpAdaptor adaptor) {
+ return static_cast<spirv::TosaExtRoundingModeType>(adaptor.getRoundingMode());
+}
+
+spirv::TosaExtAccType getAccType(Type accType) {
+ if (accType.isInteger(32))
+ return spirv::TosaExtAccType::INT32;
+ else if (accType.isF16())
+ return spirv::TosaExtAccType::FP16;
+ else if (accType.isF32())
+ return spirv::TosaExtAccType::FP32;
+ else if (accType.isInteger(48))
+ return spirv::TosaExtAccType::INT48;
+ llvm_unreachable("unknown accumulator type");
+}
+
DenseIntElementsAttr getI32TensorArmAttr(ArrayRef<int32_t> values,
ConversionPatternRewriter &rewriter) {
return DenseIntElementsAttr::get(
@@ -39,6 +58,87 @@ DenseIntElementsAttr getI32TensorArmAttr(ArrayRef<int32_t> values,
values);
}
+// TOSA stores many integer array attributes as i64 in MLIR, while the
+// SPIR-V TOSA extended instruction set models the same attributes as i32.
+DenseIntElementsAttr getI32TensorArmAttr(ArrayRef<int64_t> values,
+ ConversionPatternRewriter &rewriter) {
+ SmallVector<int32_t> i32Values(values.begin(), values.end());
+ return getI32TensorArmAttr(i32Values, rewriter);
+}
+
+FailureOr<DenseElementsAttr>
+convertDenseElementsAttr(DenseElementsAttr values, ShapedType convertedType) {
+ Type convertedElementType = convertedType.getElementType();
+ if (values.getElementType() == convertedElementType)
+ return values.reshape(convertedType);
+
+ // Constant attributes still have the source TOSA element type. Rebuild them
+ // for the converted SPIR-V tensor type, including integer type-converter
+ // changes such as index to i32, i4 to i8, and i48 to i64.
+ auto integerType = dyn_cast<IntegerType>(convertedElementType);
+ if (!integerType)
+ return failure();
+
+ // The SPIR-V ARM tensor type does not represent scalar shape constants
+ // directly, so model tensor<0xindex> as a one-element i32 tensor.
+ if (values.empty() && values.getElementType().isIndex())
+ return DenseIntElementsAttr::get(convertedType, {1});
+
+ DenseElementsAttr convertedValues =
+ values.mapValues(integerType, [&](const APInt &value) {
+ return value.sextOrTrunc(integerType.getWidth());
+ });
+ return convertedValues.reshape(convertedType);
+}
+
+// Split a large concat into smaller concat operations so the generated SPIR-V
+// instructions stay below the binary operand count limit.
+LogicalResult splitConcat(tosa::ConcatOp op, Type resultType, int32_t axis,
+ ValueRange inputs,
+ ConversionPatternRewriter &rewriter) {
+ auto resultTensorType = dyn_cast<spirv::TensorArmType>(resultType);
+ if (!resultTensorType)
+ return rewriter.notifyMatchFailure(op, "expected tensor result type");
+ if (!resultTensorType.hasRank())
+ return rewriter.notifyMatchFailure(op,
+ "expected ranked tensor result type");
+
+ SmallVector<Value> concatInputs;
+ SmallVector<int64_t> concatShape(resultTensorType.getShape());
+ concatShape[axis] = 0;
+
+ for (auto [index, input] : llvm::enumerate(inputs)) {
+ auto inputType = dyn_cast<spirv::TensorArmType>(input.getType());
+ if (!inputType)
+ return rewriter.notifyMatchFailure(op, "expected tensor input type");
+ if (!inputType.hasRank())
+ return rewriter.notifyMatchFailure(op,
+ "expected ranked tensor input type");
+
+ int64_t inputAxisDim = inputType.getShape()[axis];
+ if (ShapedType::isDynamic(inputAxisDim) ||
+ ShapedType::isDynamic(concatShape[axis]))
+ concatShape[axis] = ShapedType::kDynamic;
+ else
+ concatShape[axis] += inputAxisDim;
+
+ concatInputs.push_back(input);
+ if (concatInputs.size() != maxConcatOpInputs || index == inputs.size() - 1)
+ continue;
+
+ Type concatType = spirv::TensorArmType::get(
+ concatShape, resultTensorType.getElementType());
+ auto concat = spirv::TosaConcatOp::create(rewriter, op.getLoc(), concatType,
+ axis, concatInputs);
+ concatInputs.clear();
+ concatInputs.push_back(concat.getOutput());
+ }
+
+ rewriter.replaceOpWithNewOp<spirv::TosaConcatOp>(op, resultType, axis,
+ concatInputs);
+ return success();
+}
+
template <typename SourceOp, auto Replace>
struct TosaOpConvert final : public OpConversionPattern<SourceOp> {
using OpConversionPattern<SourceOp>::OpConversionPattern;
@@ -53,6 +153,21 @@ struct TosaOpConvert final : public OpConversionPattern<SourceOp> {
}
};
+template <typename SourceOp, auto Replace>
+struct TosaMultiResultOpConvert final : public OpConversionPattern<SourceOp> {
+ using OpConversionPattern<SourceOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(SourceOp op, typename SourceOp::Adaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ SmallVector<Type> types;
+ if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(),
+ types)))
+ return rewriter.notifyMatchFailure(op, "type conversion failed");
+ return Replace(op, adaptor, types, rewriter);
+ }
+};
+
template <typename SourceOp, typename TargetOp>
LogicalResult replaceUnaryInput1(SourceOp op,
typename SourceOp::Adaptor adaptor, Type type,
@@ -105,6 +220,94 @@ replaceNanModeReduction(SourceOp op, typename SourceOp::Adaptor adaptor,
return success();
}
+LogicalResult replaceAvgPool2d(tosa::AvgPool2dOp op,
+ tosa::AvgPool2dOpAdaptor adaptor, Type type,
+ ConversionPatternRewriter &rewriter) {
+ rewriter.replaceOpWithNewOp<spirv::TosaAvgPool2DOp>(
+ op, type, getI32TensorArmAttr(adaptor.getKernel(), rewriter),
+ getI32TensorArmAttr(adaptor.getStride(), rewriter),
+ getI32TensorArmAttr(adaptor.getPad(), rewriter),
+ getAccType(adaptor.getAccType()), adaptor.getInput(),
+ adaptor.getInputZp(), adaptor.getOutputZp());
+ return success();
+}
+
+template <typename SourceOp, typename TargetOp>
+LogicalResult replaceConvolution(SourceOp op,
+ typename SourceOp::Adaptor adaptor, Type type,
+ ConversionPatternRewriter &rewriter) {
+ rewriter.replaceOpWithNewOp<TargetOp>(
+ op, type, getI32TensorArmAttr(adaptor.getPad(), rewriter),
+ getI32TensorArmAttr(adaptor.getStride(), rewriter),
+ getI32TensorArmAttr(adaptor.getDilation(), rewriter),
+ getAccType(adaptor.getAccType()), adaptor.getLocalBound(),
+ adaptor.getInput(), adaptor.getWeight(), adaptor.getBias(),
+ adaptor.getInputZp(), adaptor.getWeightZp());
+ return success();
+}
+
+LogicalResult replaceFFT2d(tosa::FFT2dOp op, tosa::FFT2dOpAdaptor adaptor,
+ ArrayRef<Type> types,
+ ConversionPatternRewriter &rewriter) {
+ auto structType = spirv::StructType::get(types);
+ auto result = spirv::TosaFFT2DOp::create(
+ rewriter, op.getLoc(), structType, adaptor.getInverse(),
+ adaptor.getLocalBound(), adaptor.getInputReal(), adaptor.getInputImag());
+ auto outputReal =
+ spirv::CompositeExtractOp::create(rewriter, op.getLoc(), result, {0});
+ auto outputImag =
+ spirv::CompositeExtractOp::create(rewriter, op.getLoc(), result, {1});
+ rewriter.replaceOp(op, {outputReal, outputImag});
+ return success();
+}
+
+LogicalResult replaceMatMul(tosa::MatMulOp op, tosa::MatMulOpAdaptor adaptor,
+ Type type, ConversionPatternRewriter &rewriter) {
+ rewriter.replaceOpWithNewOp<spirv::TosaMatMulOp>(
+ op, type, adaptor.getA(), adaptor.getB(), adaptor.getAZp(),
+ adaptor.getBZp());
+ return success();
+}
+
+LogicalResult replaceMaxPool2d(tosa::MaxPool2dOp op,
+ tosa::MaxPool2dOpAdaptor adaptor, Type type,
+ ConversionPatternRewriter &rewriter) {
+ rewriter.replaceOpWithNewOp<spirv::TosaMaxPool2DOp>(
+ op, type, getI32TensorArmAttr(adaptor.getKernel(), rewriter),
+ getI32TensorArmAttr(adaptor.getStride(), rewriter),
+ getI32TensorArmAttr(adaptor.getPad(), rewriter), getNanMode(adaptor),
+ adaptor.getInput());
+ return success();
+}
+
+LogicalResult replaceRFFT2d(tosa::RFFT2dOp op, tosa::RFFT2dOpAdaptor adaptor,
+ ArrayRef<Type> types,
+ ConversionPatternRewriter &rewriter) {
+ auto structType = spirv::StructType::get(types);
+ auto result = spirv::TosaRFFT2DOp::create(rewriter, op.getLoc(), structType,
+ adaptor.getLocalBound(),
+ adaptor.getInputReal());
+ auto outputReal =
+ spirv::CompositeExtractOp::create(rewriter, op.getLoc(), result, {0});
+ auto outputImag =
+ spirv::CompositeExtractOp::create(rewriter, op.getLoc(), result, {1});
+ rewriter.replaceOp(op, {outputReal, outputImag});
+ return success();
+}
+
+LogicalResult replaceTransposeConv2d(tosa::TransposeConv2DOp op,
+ tosa::TransposeConv2DOpAdaptor adaptor,
+ Type type,
+ ConversionPatternRewriter &rewriter) {
+ rewriter.replaceOpWithNewOp<spirv::TosaTransposeConv2DOp>(
+ op, type, getI32TensorArmAttr(adaptor.getOutPad(), rewriter),
+ getI32TensorArmAttr(adaptor.getStride(), rewriter),
+ getAccType(adaptor.getAccType()), adaptor.getLocalBound(),
+ adaptor.getInput(), adaptor.getWeight(), adaptor.getBias(),
+ adaptor.getInputZp(), adaptor.getWeightZp());
+ return success();
+}
+
LogicalResult replaceClamp(tosa::ClampOp op, tosa::ClampOpAdaptor adaptor,
Type type, ConversionPatternRewriter &rewriter) {
rewriter.replaceOpWithNewOp<spirv::TosaClampOp>(
@@ -151,6 +354,28 @@ LogicalResult replaceSelect(tosa::SelectOp op, tosa::SelectOpAdaptor adaptor,
return success();
}
+LogicalResult replaceConcat(tosa::ConcatOp op, tosa::ConcatOpAdaptor adaptor,
+ Type type, ConversionPatternRewriter &rewriter) {
+ // Large TOSA concats can produce SPIR-V instructions with too many
+ // operands and fail validation. Split them into conservative 64-input
+ // chunks to keep the generated SPIR-V valid.
+ if (adaptor.getInput1().size() > maxConcatOpInputs)
+ return splitConcat(op, type, adaptor.getAxis(), adaptor.getInput1(),
+ rewriter);
+
+ rewriter.replaceOpWithNewOp<spirv::TosaConcatOp>(op, type, adaptor.getAxis(),
+ adaptor.getInput1());
+ return success();
+}
+
+LogicalResult replacePad(tosa::PadOp op, tosa::PadOpAdaptor adaptor, Type type,
+ ConversionPatternRewriter &rewriter) {
+ rewriter.replaceOpWithNewOp<spirv::TosaPadOp>(op, type, adaptor.getInput1(),
+ adaptor.getPadding(),
+ adaptor.getPadConst());
+ return success();
+}
+
LogicalResult replaceReshape(tosa::ReshapeOp op, tosa::ReshapeOpAdaptor adaptor,
Type type, ConversionPatternRewriter &rewriter) {
rewriter.replaceOpWithNewOp<spirv::TosaReshapeOp>(
@@ -212,15 +437,37 @@ LogicalResult replaceResize(tosa::ResizeOp op, tosa::ResizeOpAdaptor adaptor,
return success();
}
-LogicalResult replaceConstShape(tosa::ConstShapeOp op,
- tosa::ConstShapeOpAdaptor adaptor, Type type,
- ConversionPatternRewriter &rewriter) {
- SmallVector<int32_t> values;
- for (const APInt &value : adaptor.getValues().getValues<APInt>())
- values.push_back(value.getSExtValue());
+LogicalResult replaceRescale(tosa::RescaleOp op, tosa::RescaleOpAdaptor adaptor,
+ Type type, ConversionPatternRewriter &rewriter) {
+ rewriter.replaceOpWithNewOp<spirv::TosaRescaleOp>(
+ op, type, adaptor.getScale32(), getRoundingMode(adaptor),
+ adaptor.getPerChannel(), adaptor.getInputUnsigned(),
+ adaptor.getOutputUnsigned(), adaptor.getInput(), adaptor.getMultiplier(),
+ adaptor.getShift(), adaptor.getInputZp(), adaptor.getOutputZp());
+ return success();
+}
+
+template <typename SourceOp>
+LogicalResult replaceConstant(SourceOp op, typename SourceOp::Adaptor adaptor,
+ Type type, ConversionPatternRewriter &rewriter) {
+ auto convertedType = dyn_cast<ShapedType>(type);
+ auto values = dyn_cast<DenseElementsAttr>(adaptor.getValues());
+ if (!convertedType || !values)
+ return failure();
+
+ FailureOr<DenseElementsAttr> convertedValues =
+ convertDenseElementsAttr(values, convertedType);
+ if (failed(convertedValues))
+ return failure();
+
+ rewriter.replaceOpWithNewOp<spirv::ConstantOp>(op, type, *convertedValues);
+ return success();
+}
- rewriter.replaceOpWithNewOp<spirv::ConstantOp>(
- op, type, getI32TensorArmAttr(values, rewriter));
+LogicalResult replaceIdentity(tosa::IdentityOp op,
+ tosa::IdentityOpAdaptor adaptor, Type type,
+ ConversionPatternRewriter &rewriter) {
+ rewriter.replaceOp(op, adaptor.getInput1());
return success();
}
@@ -231,6 +478,19 @@ void populateTosaToSPIRVTosaOpsConversionPatterns(
patterns.add<
TosaOpConvert<tosa::ArgMaxOp, replaceNanModeReduction<
tosa::ArgMaxOp, spirv::TosaArgMaxOp>>,
+ TosaOpConvert<tosa::AvgPool2dOp, replaceAvgPool2d>,
+ TosaOpConvert<tosa::Conv2DOp,
+ replaceConvolution<tosa::Conv2DOp, spirv::TosaConv2DOp>>,
+ TosaOpConvert<tosa::Conv3DOp,
+ replaceConvolution<tosa::Conv3DOp, spirv::TosaConv3DOp>>,
+ TosaOpConvert<tosa::DepthwiseConv2DOp,
+ replaceConvolution<tosa::DepthwiseConv2DOp,
+ spirv::TosaDepthwiseConv2DOp>>,
+ TosaMultiResultOpConvert<tosa::FFT2dOp, replaceFFT2d>,
+ TosaOpConvert<tosa::MatMulOp, replaceMatMul>,
+ TosaOpConvert<tosa::MaxPool2dOp, replaceMaxPool2d>,
+ TosaMultiResultOpConvert<tosa::RFFT2dOp, replaceRFFT2d>,
+ TosaOpConvert<tosa::TransposeConv2DOp, replaceTransposeConv2d>,
TosaOpConvert<tosa::ClampOp, replaceClamp>,
TosaOpConvert<tosa::ErfOp,
replaceUnaryInput<tosa::ErfOp, spirv::TosaErfOp>>,
@@ -334,6 +594,8 @@ void populateTosaToSPIRVTosaOpsConversionPatterns(
TosaOpConvert<
tosa::ReduceSumOp,
replaceReduction<tosa::ReduceSumOp, spirv::TosaReduceSumOp>>,
+ TosaOpConvert<tosa::ConcatOp, replaceConcat>,
+ TosaOpConvert<tosa::PadOp, replacePad>,
TosaOpConvert<tosa::ReshapeOp, replaceReshape>,
TosaOpConvert<tosa::ReverseOp, replaceReverse>,
TosaOpConvert<tosa::SliceOp, replaceSlice>,
@@ -344,8 +606,11 @@ void populateTosaToSPIRVTosaOpsConversionPatterns(
TosaOpConvert<tosa::ResizeOp, replaceResize>,
TosaOpConvert<tosa::CastOp,
replaceUnaryInput<tosa::CastOp, spirv::TosaCastOp>>,
- TosaOpConvert<tosa::ConstShapeOp, replaceConstShape>>(
- typeConverter, patterns.getContext());
+ TosaOpConvert<tosa::RescaleOp, replaceRescale>,
+ TosaOpConvert<tosa::ConstOp, replaceConstant<tosa::ConstOp>>,
+ TosaOpConvert<tosa::ConstShapeOp, replaceConstant<tosa::ConstShapeOp>>,
+ TosaOpConvert<tosa::IdentityOp, replaceIdentity>>(typeConverter,
+ patterns.getContext());
}
} // namespace mlir::tosa
diff --git a/mlir/test/Conversion/TosaToSPIRVTosa/tosa-to-spirv.mlir b/mlir/test/Conversion/TosaToSPIRVTosa/tosa-to-spirv.mlir
index a175baf62eda1..84edc10e8832e 100644
--- a/mlir/test/Conversion/TosaToSPIRVTosa/tosa-to-spirv.mlir
+++ b/mlir/test/Conversion/TosaToSPIRVTosa/tosa-to-spirv.mlir
@@ -13,6 +13,127 @@ func.func @argmax_int(%arg0: tensor<2x3x4xi8>) -> tensor<2x4xi32> {
// -----
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.AvgPool2D
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @avg_pool2d_int
+func.func @avg_pool2d_int(%arg0: tensor<1x4x4x1xi8>, %arg1: tensor<1xi8>, %arg2: tensor<1xi8>) -> tensor<1x2x2x1xi8> {
+ // CHECK: %[[AVG_POOL:.*]] = spirv.Tosa.AvgPool2D kernel = [2, 2], stride = [2, 2], pad = [0, 0, 0, 0], acc_type = <INT32>, %arg0, %arg1, %arg2 : !spirv.arm.tensor<1x4x4x1xi8>, !spirv.arm.tensor<1xi8>, !spirv.arm.tensor<1xi8> -> !spirv.arm.tensor<1x2x2x1xi8>
+ %res = tosa.avg_pool2d %arg0, %arg1, %arg2 {kernel = array<i64: 2, 2>, stride = array<i64: 2, 2>, pad = array<i64: 0, 0, 0, 0>, acc_type = i32} : (tensor<1x4x4x1xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<1x2x2x1xi8>
+ return %res : tensor<1x2x2x1xi8>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.Conv2D
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @conv2d_int
+func.func @conv2d_int(%arg0: tensor<1x4x4x1xi8>, %arg1: tensor<2x3x3x1xi8>, %arg2: tensor<2xi32>, %arg3: tensor<1xi8>, %arg4: tensor<1xi8>) -> tensor<1x2x2x2xi32> {
+ // CHECK: %[[CONV2D:.*]] = spirv.Tosa.Conv2D pad = [0, 0, 0, 0], stride = [1, 1], dilation = [1, 1], acc_type = <INT32>, local_bound = false, %arg0, %arg1, %arg2, %arg3, %arg4 : !spirv.arm.tensor<1x4x4x1xi8>, !spirv.arm.tensor<2x3x3x1xi8>, !spirv.arm.tensor<2xi32>, !spirv.arm.tensor<1xi8>, !spirv.arm.tensor<1xi8> -> !spirv.arm.tensor<1x2x2x2xi32>
+ %res = tosa.conv2d %arg0, %arg1, %arg2, %arg3, %arg4 {pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>, dilation = array<i64: 1, 1>, acc_type = i32, local_bound = false} : (tensor<1x4x4x1xi8>, tensor<2x3x3x1xi8>, tensor<2xi32>, tensor<1xi8>, tensor<1xi8>) -> tensor<1x2x2x2xi32>
+ return %res : tensor<1x2x2x2xi32>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.Conv3D
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @conv3d_int
+func.func @conv3d_int(%arg0: tensor<1x4x4x4x1xi8>, %arg1: tensor<2x2x2x2x1xi8>, %arg2: tensor<2xi32>, %arg3: tensor<1xi8>, %arg4: tensor<1xi8>) -> tensor<1x3x3x3x2xi32> {
+ // CHECK: %[[CONV3D:.*]] = spirv.Tosa.Conv3D pad = [0, 0, 0, 0, 0, 0], stride = [1, 1, 1], dilation = [1, 1, 1], acc_type = <INT32>, local_bound = false, %arg0, %arg1, %arg2, %arg3, %arg4 : !spirv.arm.tensor<1x4x4x4x1xi8>, !spirv.arm.tensor<2x2x2x2x1xi8>, !spirv.arm.tensor<2xi32>, !spirv.arm.tensor<1xi8>, !spirv.arm.tensor<1xi8> -> !spirv.arm.tensor<1x3x3x3x2xi32>
+ %res = tosa.conv3d %arg0, %arg1, %arg2, %arg3, %arg4 {pad = array<i64: 0, 0, 0, 0, 0, 0>, stride = array<i64: 1, 1, 1>, dilation = array<i64: 1, 1, 1>, acc_type = i32, local_bound = false} : (tensor<1x4x4x4x1xi8>, tensor<2x2x2x2x1xi8>, tensor<2xi32>, tensor<1xi8>, tensor<1xi8>) -> tensor<1x3x3x3x2xi32>
+ return %res : tensor<1x3x3x3x2xi32>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.DepthwiseConv2D
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @depthwise_conv2d_int
+func.func @depthwise_conv2d_int(%arg0: tensor<1x4x4x1xi8>, %arg1: tensor<3x3x1x2xi8>, %arg2: tensor<2xi32>, %arg3: tensor<1xi8>, %arg4: tensor<1xi8>) -> tensor<1x2x2x2xi32> {
+ // CHECK: %[[DEPTHWISE_CONV2D:.*]] = spirv.Tosa.DepthwiseConv2D pad = [0, 0, 0, 0], stride = [1, 1], dilation = [1, 1], acc_type = <INT32>, local_bound = false, %arg0, %arg1, %arg2, %arg3, %arg4 : !spirv.arm.tensor<1x4x4x1xi8>, !spirv.arm.tensor<3x3x1x2xi8>, !spirv.arm.tensor<2xi32>, !spirv.arm.tensor<1xi8>, !spirv.arm.tensor<1xi8> -> !spirv.arm.tensor<1x2x2x2xi32>
+ %res = tosa.depthwise_conv2d %arg0, %arg1, %arg2, %arg3, %arg4 {pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>, dilation = array<i64: 1, 1>, acc_type = i32, local_bound = false} : (tensor<1x4x4x1xi8>, tensor<3x3x1x2xi8>, tensor<2xi32>, tensor<1xi8>, tensor<1xi8>) -> tensor<1x2x2x2xi32>
+ return %res : tensor<1x2x2x2xi32>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.FFT2D
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @fft2d_fp
+func.func @fft2d_fp(%arg0: tensor<1x32x32xf32>, %arg1: tensor<1x32x32xf32>) -> (tensor<1x32x32xf32>, tensor<1x32x32xf32>) {
+ // CHECK: %[[FFT:.*]] = spirv.Tosa.FFT2D inverse = true, local_bound = false, %arg0, %arg1 : !spirv.arm.tensor<1x32x32xf32>, !spirv.arm.tensor<1x32x32xf32> -> !spirv.struct<(!spirv.arm.tensor<1x32x32xf32>, !spirv.arm.tensor<1x32x32xf32>)>
+ // CHECK: %[[REAL:.*]] = spirv.CompositeExtract %[[FFT]][0 : i32] : !spirv.struct<(!spirv.arm.tensor<1x32x32xf32>, !spirv.arm.tensor<1x32x32xf32>)>
+ // CHECK: %[[IMAG:.*]] = spirv.CompositeExtract %[[FFT]][1 : i32] : !spirv.struct<(!spirv.arm.tensor<1x32x32xf32>, !spirv.arm.tensor<1x32x32xf32>)>
+ %real, %imag = "tosa.fft2d"(%arg0, %arg1) <{inverse = true, local_bound = false}> : (tensor<1x32x32xf32>, tensor<1x32x32xf32>) -> (tensor<1x32x32xf32>, tensor<1x32x32xf32>)
+ return %real, %imag : tensor<1x32x32xf32>, tensor<1x32x32xf32>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.MatMul
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @matmul_int
+func.func @matmul_int(%arg0: tensor<1x2x3xi8>, %arg1: tensor<1x3x4xi8>, %arg2: tensor<1xi8>, %arg3: tensor<1xi8>) -> tensor<1x2x4xi32> {
+ // CHECK: %[[MATMUL:.*]] = spirv.Tosa.MatMul %arg0, %arg1, %arg2, %arg3 : !spirv.arm.tensor<1x2x3xi8>, !spirv.arm.tensor<1x3x4xi8>, !spirv.arm.tensor<1xi8>, !spirv.arm.tensor<1xi8> -> !spirv.arm.tensor<1x2x4xi32>
+ %res = tosa.matmul %arg0, %arg1, %arg2, %arg3 : (tensor<1x2x3xi8>, tensor<1x3x4xi8>, tensor<1xi8>, tensor<1xi8>) -> tensor<1x2x4xi32>
+ return %res : tensor<1x2x4xi32>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.MaxPool2D
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @max_pool2d_fp
+func.func @max_pool2d_fp(%arg0: tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32> {
+ // CHECK: %[[MAX_POOL:.*]] = spirv.Tosa.MaxPool2D kernel = [2, 2], stride = [2, 2], pad = [0, 0, 0, 0], nan_mode = <Propagate>, %arg0 : !spirv.arm.tensor<1x4x4x1xf32> -> !spirv.arm.tensor<1x2x2x1xf32>
+ %res = tosa.max_pool2d %arg0 {kernel = array<i64: 2, 2>, stride = array<i64: 2, 2>, pad = array<i64: 0, 0, 0, 0>, nan_mode = PROPAGATE} : (tensor<1x4x4x1xf32>) -> tensor<1x2x2x1xf32>
+ return %res : tensor<1x2x2x1xf32>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.RFFT2D
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @rfft2d_fp
+func.func @rfft2d_fp(%arg0: tensor<1x32x32xf32>) -> (tensor<1x32x17xf32>, tensor<1x32x17xf32>) {
+ // CHECK: %[[RFFT:.*]] = spirv.Tosa.RFFT2D local_bound = false, %arg0 : !spirv.arm.tensor<1x32x32xf32> -> !spirv.struct<(!spirv.arm.tensor<1x32x17xf32>, !spirv.arm.tensor<1x32x17xf32>)>
+ // CHECK: %[[REAL:.*]] = spirv.CompositeExtract %[[RFFT]][0 : i32] : !spirv.struct<(!spirv.arm.tensor<1x32x17xf32>, !spirv.arm.tensor<1x32x17xf32>)>
+ // CHECK: %[[IMAG:.*]] = spirv.CompositeExtract %[[RFFT]][1 : i32] : !spirv.struct<(!spirv.arm.tensor<1x32x17xf32>, !spirv.arm.tensor<1x32x17xf32>)>
+ %real, %imag = "tosa.rfft2d"(%arg0) <{local_bound = false}> : (tensor<1x32x32xf32>) -> (tensor<1x32x17xf32>, tensor<1x32x17xf32>)
+ return %real, %imag : tensor<1x32x17xf32>, tensor<1x32x17xf32>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.TransposeConv2D
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @transpose_conv2d_int
+func.func @transpose_conv2d_int(%arg0: tensor<1x2x2x1xi8>, %arg1: tensor<2x3x3x1xi8>, %arg2: tensor<2xi32>, %arg3: tensor<1xi8>, %arg4: tensor<1xi8>) -> tensor<1x4x4x2xi32> {
+ // CHECK: %[[TRANSPOSE_CONV2D:.*]] = spirv.Tosa.TransposeConv2D out_pad = [0, 0, 0, 0], stride = [1, 1], acc_type = <INT32>, local_bound = false, %arg0, %arg1, %arg2, %arg3, %arg4 : !spirv.arm.tensor<1x2x2x1xi8>, !spirv.arm.tensor<2x3x3x1xi8>, !spirv.arm.tensor<2xi32>, !spirv.arm.tensor<1xi8>, !spirv.arm.tensor<1xi8> -> !spirv.arm.tensor<1x4x4x2xi32>
+ %res = tosa.transpose_conv2d %arg0, %arg1, %arg2, %arg3, %arg4 {out_pad = array<i64: 0, 0, 0, 0>, stride = array<i64: 1, 1>, acc_type = i32, local_bound = false} : (tensor<1x2x2x1xi8>, tensor<2x3x3x1xi8>, tensor<2xi32>, tensor<1xi8>, tensor<1xi8>) -> tensor<1x4x4x2xi32>
+ return %res : tensor<1x4x4x2xi32>
+}
+
+// -----
+
//===----------------------------------------------------------------------===//
// spirv.TOSA.Clamp
//===----------------------------------------------------------------------===//
@@ -585,6 +706,98 @@ func.func @reduce_sum_int(%arg0: tensor<2x3x4xi32>) -> tensor<2x1x4xi32> {
// -----
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.Concat
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @concat_int
+func.func @concat_int(%arg0: tensor<2x3xi8>, %arg1: tensor<2x4xi8>) -> tensor<2x7xi8> {
+ // CHECK: %[[CONCAT:.*]] = spirv.Tosa.Concat axis = 1, %arg0, %arg1 : !spirv.arm.tensor<2x3xi8>, !spirv.arm.tensor<2x4xi8> -> !spirv.arm.tensor<2x7xi8>
+ %res = tosa.concat %arg0, %arg1 {axis = 1 : i32} : (tensor<2x3xi8>, tensor<2x4xi8>) -> tensor<2x7xi8>
+ return %res : tensor<2x7xi8>
+}
+
+// -----
+
+// CHECK-LABEL: spirv.ARM.Graph @concat_split_int
+func.func @concat_split_int(%arg0: tensor<1xi8>, %arg1: tensor<1xi8>, %arg2: tensor<1xi8>, %arg3: tensor<1xi8>,
+ %arg4: tensor<1xi8>, %arg5: tensor<1xi8>, %arg6: tensor<1xi8>, %arg7: tensor<1xi8>,
+ %arg8: tensor<1xi8>, %arg9: tensor<1xi8>, %arg10: tensor<1xi8>, %arg11: tensor<1xi8>,
+ %arg12: tensor<1xi8>, %arg13: tensor<1xi8>, %arg14: tensor<1xi8>, %arg15: tensor<1xi8>,
+ %arg16: tensor<1xi8>, %arg17: tensor<1xi8>, %arg18: tensor<1xi8>, %arg19: tensor<1xi8>,
+ %arg20: tensor<1xi8>, %arg21: tensor<1xi8>, %arg22: tensor<1xi8>, %arg23: tensor<1xi8>,
+ %arg24: tensor<1xi8>, %arg25: tensor<1xi8>, %arg26: tensor<1xi8>, %arg27: tensor<1xi8>,
+ %arg28: tensor<1xi8>, %arg29: tensor<1xi8>, %arg30: tensor<1xi8>, %arg31: tensor<1xi8>,
+ %arg32: tensor<1xi8>, %arg33: tensor<1xi8>, %arg34: tensor<1xi8>, %arg35: tensor<1xi8>,
+ %arg36: tensor<1xi8>, %arg37: tensor<1xi8>, %arg38: tensor<1xi8>, %arg39: tensor<1xi8>,
+ %arg40: tensor<1xi8>, %arg41: tensor<1xi8>, %arg42: tensor<1xi8>, %arg43: tensor<1xi8>,
+ %arg44: tensor<1xi8>, %arg45: tensor<1xi8>, %arg46: tensor<1xi8>, %arg47: tensor<1xi8>,
+ %arg48: tensor<1xi8>, %arg49: tensor<1xi8>, %arg50: tensor<1xi8>, %arg51: tensor<1xi8>,
+ %arg52: tensor<1xi8>, %arg53: tensor<1xi8>, %arg54: tensor<1xi8>, %arg55: tensor<1xi8>,
+ %arg56: tensor<1xi8>, %arg57: tensor<1xi8>, %arg58: tensor<1xi8>, %arg59: tensor<1xi8>,
+ %arg60: tensor<1xi8>, %arg61: tensor<1xi8>, %arg62: tensor<1xi8>, %arg63: tensor<1xi8>,
+ %arg64: tensor<1xi8>) -> tensor<65xi8> {
+ // CHECK: %[[CONCAT0:.*]] = spirv.Tosa.Concat axis = 0, %arg0, %arg1{{.*}} -> !spirv.arm.tensor<64xi8>
+ // CHECK: %[[CONCAT1:.*]] = spirv.Tosa.Concat axis = 0, %[[CONCAT0]], %arg64 : !spirv.arm.tensor<64xi8>, !spirv.arm.tensor<1xi8> -> !spirv.arm.tensor<65xi8>
+ %res = tosa.concat
+ %arg0, %arg1, %arg2, %arg3, %arg4, %arg5, %arg6, %arg7,
+ %arg8, %arg9, %arg10, %arg11, %arg12, %arg13, %arg14, %arg15,
+ %arg16, %arg17, %arg18, %arg19, %arg20, %arg21, %arg22, %arg23,
+ %arg24, %arg25, %arg26, %arg27, %arg28, %arg29, %arg30, %arg31,
+ %arg32, %arg33, %arg34, %arg35, %arg36, %arg37, %arg38, %arg39,
+ %arg40, %arg41, %arg42, %arg43, %arg44, %arg45, %arg46, %arg47,
+ %arg48, %arg49, %arg50, %arg51, %arg52, %arg53, %arg54, %arg55,
+ %arg56, %arg57, %arg58, %arg59, %arg60, %arg61, %arg62, %arg63,
+ %arg64 {axis = 0 : i32}
+ : (tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>, tensor<1xi8>, tensor<1xi8>, tensor<1xi8>,
+ tensor<1xi8>) -> tensor<65xi8>
+ return %res : tensor<65xi8>
+}
+
+// -----
+
+// CHECK-LABEL: spirv.ARM.Graph @concat_split_dynamic
+func.func @concat_split_dynamic(%arg0: tensor<?xi8>) -> tensor<?xi8> {
+ // CHECK: %[[CONCAT0:.*]] = spirv.Tosa.Concat axis = 0, %arg0, %arg0{{.*}} -> !spirv.arm.tensor<?xi8>
+ // CHECK: %[[CONCAT1:.*]] = spirv.Tosa.Concat axis = 0, %[[CONCAT0]], %arg0 : !spirv.arm.tensor<?xi8>, !spirv.arm.tensor<?xi8> -> !spirv.arm.tensor<?xi8>
+ %res = tosa.concat %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0 {axis = 0 : i32} : (tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>, tensor<?xi8>) -> tensor<?xi8>
+ return %res : tensor<?xi8>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.Pad
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @pad_int
+func.func @pad_int(%arg0: tensor<1x2xi8>) -> tensor<4x8xi8> {
+ %padding = "tosa.const_shape"() <{values = dense<[1, 2, 3, 3]> : tensor<4xindex>}> : () -> !tosa.shape<4>
+ %pad_const = "tosa.const"() <{values = dense<7> : tensor<1xi8>}> : () -> tensor<1xi8>
+ // CHECK: %[[PADDING:.*]] = spirv.Constant dense<[1, 2, 3, 3]> : !spirv.arm.tensor<4xi32>
+ // CHECK: %[[PAD_CONST:.*]] = spirv.Constant dense<7> : !spirv.arm.tensor<1xi8>
+ // CHECK: %[[PAD:.*]] = spirv.Tosa.Pad %arg0, %[[PADDING]], %[[PAD_CONST]] : !spirv.arm.tensor<1x2xi8>, !spirv.arm.tensor<4xi32>, !spirv.arm.tensor<1xi8> -> !spirv.arm.tensor<4x8xi8>
+ %res = tosa.pad %arg0, %padding, %pad_const : (tensor<1x2xi8>, !tosa.shape<4>, tensor<1xi8>) -> tensor<4x8xi8>
+ return %res : tensor<4x8xi8>
+}
+
+// -----
+
//===----------------------------------------------------------------------===//
// spirv.TOSA.Reshape
//===----------------------------------------------------------------------===//
@@ -713,3 +926,78 @@ func.func @cast_int(%arg0: tensor<2x3xi8>) -> tensor<2x3xi32> {
%res = tosa.cast %arg0 : (tensor<2x3xi8>) -> tensor<2x3xi32>
return %res : tensor<2x3xi32>
}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.Rescale
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @rescale_int
+func.func @rescale_int(%arg0: tensor<2x3xi16>) -> tensor<2x3xi16> {
+ %multiplier = "tosa.const"() <{values = dense<1073741824> : tensor<1xi32>}> : () -> tensor<1xi32>
+ %shift = "tosa.const"() <{values = dense<30> : tensor<1xi8>}> : () -> tensor<1xi8>
+ %input_zp = "tosa.const"() <{values = dense<0> : tensor<1xi16>}> : () -> tensor<1xi16>
+ %output_zp = "tosa.const"() <{values = dense<0> : tensor<1xi16>}> : () -> tensor<1xi16>
+ // CHECK: %[[MULTIPLIER:.*]] = spirv.Constant dense<1073741824> : !spirv.arm.tensor<1xi32>
+ // CHECK: %[[SHIFT:.*]] = spirv.Constant dense<30> : !spirv.arm.tensor<1xi8>
+ // CHECK: %[[INPUT_ZP:.*]] = spirv.Constant dense<0> : !spirv.arm.tensor<1xi16>
+ // CHECK: %[[OUTPUT_ZP:.*]] = spirv.Constant dense<0> : !spirv.arm.tensor<1xi16>
+ // CHECK: %[[RESCALE:.*]] = spirv.Tosa.Rescale scale32 = true, rounding_mode = <DoubleRound>, per_channel = false, input_unsigned = false, output_unsigned = false, %arg0, %[[MULTIPLIER]], %[[SHIFT]], %[[INPUT_ZP]], %[[OUTPUT_ZP]] : !spirv.arm.tensor<2x3xi16>, !spirv.arm.tensor<1xi32>, !spirv.arm.tensor<1xi8>, !spirv.arm.tensor<1xi16>, !spirv.arm.tensor<1xi16> -> !spirv.arm.tensor<2x3xi16>
+ %res = tosa.rescale %arg0, %multiplier, %shift, %input_zp, %output_zp {scale32 = true, rounding_mode = DOUBLE_ROUND, per_channel = false, input_unsigned = false, output_unsigned = false} : (tensor<2x3xi16>, tensor<1xi32>, tensor<1xi8>, tensor<1xi16>, tensor<1xi16>) -> tensor<2x3xi16>
+ return %res : tensor<2x3xi16>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.Const
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @const_int
+func.func @const_int() -> tensor<2x3xi8> {
+ // CHECK: %[[CONST:.*]] = spirv.Constant dense<{{.*}}> : !spirv.arm.tensor<2x3xi8>
+ %res = "tosa.const"() <{values = dense<[[1, 2, 3], [4, 5, 6]]> : tensor<2x3xi8>}> : () -> tensor<2x3xi8>
+ return %res : tensor<2x3xi8>
+}
+
+// -----
+
+// CHECK-LABEL: spirv.ARM.Graph @const_i4
+func.func @const_i4() -> tensor<2xi4> {
+ // CHECK: %[[CONST:.*]] = spirv.Constant dense<[7, -8]> : !spirv.arm.tensor<2xi8>
+ %res = "tosa.const"() <{values = dense<[7, -8]> : tensor<2xi4>}> : () -> tensor<2xi4>
+ return %res : tensor<2xi4>
+}
+
+// -----
+
+// CHECK-LABEL: spirv.ARM.Graph @const_i48
+func.func @const_i48() -> tensor<2xi48> {
+ // CHECK: %[[CONST:.*]] = spirv.Constant dense<[140737488355327, -140737488355328]> : !spirv.arm.tensor<2xi64>
+ %res = "tosa.const"() <{values = dense<[140737488355327, -140737488355328]> : tensor<2xi48>}> : () -> tensor<2xi48>
+ return %res : tensor<2xi48>
+}
+
+// -----
+
+// CHECK-LABEL: spirv.ARM.Graph @const_shape_empty
+func.func @const_shape_empty() -> !tosa.shape<0> {
+ // CHECK: %[[SHAPE:.*]] = spirv.Constant dense<1> : !spirv.arm.tensor<1xi32>
+ %res = "tosa.const_shape"() <{values = dense<> : tensor<0xindex>}> : () -> !tosa.shape<0>
+ return %res : !tosa.shape<0>
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// spirv.TOSA.Identity
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: spirv.ARM.Graph @identity_int
+func.func @identity_int(%arg0: tensor<2x3xi8>) -> tensor<2x3xi8> {
+ // CHECK-NOT: spirv.Tosa.Identity
+ // CHECK: spirv.ARM.GraphOutputs %arg0 : !spirv.arm.tensor<2x3xi8>
+ %res = tosa.identity %arg0 : (tensor<2x3xi8>) -> tensor<2x3xi8>
+ return %res : tensor<2x3xi8>
+}
More information about the Mlir-commits
mailing list