//===- TosaToSPIRVTosaOps.cpp - TOSA to SPIR-V Graph/TOSA patterns --------===// // // 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 // //===----------------------------------------------------------------------===// // // This file implements patterns to convert TOSA IR to SPIR-V Graph/TOSA. // //===----------------------------------------------------------------------===// #include "mlir/Conversion/TosaToSPIRVTosa/TosaToSPIRVTosa.h" #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h" #include "mlir/Dialect/Tosa/IR/TosaOps.h" #include "mlir/Transforms/DialectConversion.h" #define DEBUG_TYPE "tosa-to-spirv-tosa-ops-pattern" namespace mlir::tosa { namespace { constexpr unsigned maxConcatOpInputs = 64; template spirv::TosaExtNaNPropagationModeType getNanMode(OpAdaptor adaptor) { return static_cast( adaptor.getNanMode()); } template spirv::TosaExtResizeModeType getResizeMode(OpAdaptor adaptor) { return static_cast(adaptor.getMode()); } template spirv::TosaExtRoundingModeType getRoundingMode(OpAdaptor adaptor) { return static_cast(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 values, ConversionPatternRewriter &rewriter) { return DenseIntElementsAttr::get( spirv::TensorArmType::get(static_cast(values.size()), IntegerType::get(rewriter.getContext(), 32)), 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 values, ConversionPatternRewriter &rewriter) { SmallVector i32Values(values.begin(), values.end()); return getI32TensorArmAttr(i32Values, rewriter); } FailureOr 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(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(resultType); if (!resultTensorType) return rewriter.notifyMatchFailure(op, "expected tensor result type"); if (!resultTensorType.hasRank()) return rewriter.notifyMatchFailure(op, "expected ranked tensor result type"); SmallVector concatInputs; SmallVector concatShape(resultTensorType.getShape()); concatShape[axis] = 0; for (auto [index, input] : llvm::enumerate(inputs)) { auto inputType = dyn_cast(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(op, resultType, axis, concatInputs); return success(); } template struct TosaOpConvert final : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; LogicalResult matchAndRewrite(SourceOp op, typename SourceOp::Adaptor adaptor, ConversionPatternRewriter &rewriter) const override { Type type = this->getTypeConverter()->convertType(op.getType()); if (!type) return rewriter.notifyMatchFailure(op, "type conversion failed"); return Replace(op, adaptor, type, rewriter); } }; template struct TosaMultiResultOpConvert final : public OpConversionPattern { using OpConversionPattern::OpConversionPattern; LogicalResult matchAndRewrite(SourceOp op, typename SourceOp::Adaptor adaptor, ConversionPatternRewriter &rewriter) const override { SmallVector types; if (failed(this->getTypeConverter()->convertTypes(op->getResultTypes(), types))) return rewriter.notifyMatchFailure(op, "type conversion failed"); return Replace(op, adaptor, types, rewriter); } }; template LogicalResult replaceUnaryInput1(SourceOp op, typename SourceOp::Adaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp(op, type, adaptor.getInput1()); return success(); } template LogicalResult replaceUnaryInput(SourceOp op, typename SourceOp::Adaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp(op, type, adaptor.getInput()); return success(); } template LogicalResult replaceBinaryElementwise(SourceOp op, typename SourceOp::Adaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp(op, type, adaptor.getInput1(), adaptor.getInput2()); return success(); } template LogicalResult replaceBinaryNanModeElementwise(SourceOp op, typename SourceOp::Adaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( op, type, getNanMode(adaptor), adaptor.getInput1(), adaptor.getInput2()); return success(); } template LogicalResult replaceReduction(SourceOp op, typename SourceOp::Adaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp(op, type, adaptor.getAxis(), adaptor.getInput()); return success(); } template LogicalResult replaceNanModeReduction(SourceOp op, typename SourceOp::Adaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( op, type, adaptor.getAxis(), getNanMode(adaptor), adaptor.getInput()); return success(); } LogicalResult replaceAvgPool2d(tosa::AvgPool2dOp op, tosa::AvgPool2dOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( 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 LogicalResult replaceConvolution(SourceOp op, typename SourceOp::Adaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( 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 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( 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( 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 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( 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( op, type, adaptor.getMinVal(), adaptor.getMaxVal(), getNanMode(adaptor), adaptor.getInput()); return success(); } LogicalResult replaceArithmeticRightShift(tosa::ArithmeticRightShiftOp op, tosa::ArithmeticRightShiftOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( op, type, adaptor.getRound(), adaptor.getInput1(), adaptor.getInput2()); return success(); } LogicalResult replaceMul(tosa::MulOp op, tosa::MulOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( op, type, adaptor.getInput1(), adaptor.getInput2(), adaptor.getShift()); return success(); } LogicalResult replaceTable(tosa::TableOp op, tosa::TableOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp(op, type, adaptor.getInput1(), adaptor.getTable()); return success(); } LogicalResult replaceNegate(tosa::NegateOp op, tosa::NegateOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( op, type, adaptor.getInput1(), adaptor.getInput1Zp(), adaptor.getOutputZp()); return success(); } LogicalResult replaceSelect(tosa::SelectOp op, tosa::SelectOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( op, type, adaptor.getInput1(), adaptor.getInput2(), adaptor.getInput3()); 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(op, type, adaptor.getAxis(), adaptor.getInput1()); return success(); } LogicalResult replacePad(tosa::PadOp op, tosa::PadOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp(op, type, adaptor.getInput1(), adaptor.getPadding(), adaptor.getPadConst()); return success(); } LogicalResult replaceReshape(tosa::ReshapeOp op, tosa::ReshapeOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( op, type, adaptor.getInput1(), adaptor.getShape()); return success(); } LogicalResult replaceReverse(tosa::ReverseOp op, tosa::ReverseOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp(op, type, adaptor.getAxis(), adaptor.getInput1()); return success(); } LogicalResult replaceSlice(tosa::SliceOp op, tosa::SliceOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( op, type, adaptor.getInput1(), adaptor.getStart(), adaptor.getSize()); return success(); } LogicalResult replaceTile(tosa::TileOp op, tosa::TileOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp(op, type, adaptor.getInput1(), adaptor.getMultiples()); return success(); } LogicalResult replaceTranspose(tosa::TransposeOp op, tosa::TransposeOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { DenseIntElementsAttr perms = getI32TensorArmAttr(adaptor.getPerms(), rewriter); rewriter.replaceOpWithNewOp(op, type, perms, adaptor.getInput1()); return success(); } LogicalResult replaceGather(tosa::GatherOp op, tosa::GatherOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( op, type, adaptor.getValues(), adaptor.getIndices()); return success(); } LogicalResult replaceScatter(tosa::ScatterOp op, tosa::ScatterOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( op, type, adaptor.getValuesIn(), adaptor.getIndices(), adaptor.getInput()); return success(); } LogicalResult replaceResize(tosa::ResizeOp op, tosa::ResizeOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( op, type, getResizeMode(adaptor), adaptor.getInput(), adaptor.getScale(), adaptor.getOffset(), adaptor.getBorder()); return success(); } LogicalResult replaceRescale(tosa::RescaleOp op, tosa::RescaleOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOpWithNewOp( 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 LogicalResult replaceConstant(SourceOp op, typename SourceOp::Adaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { auto convertedType = dyn_cast(type); auto values = dyn_cast(adaptor.getValues()); if (!convertedType || !values) return failure(); FailureOr convertedValues = convertDenseElementsAttr(values, convertedType); if (failed(convertedValues)) return failure(); rewriter.replaceOpWithNewOp(op, type, *convertedValues); return success(); } LogicalResult replaceIdentity(tosa::IdentityOp op, tosa::IdentityOpAdaptor adaptor, Type type, ConversionPatternRewriter &rewriter) { rewriter.replaceOp(op, adaptor.getInput1()); return success(); } } // namespace void populateTosaToSPIRVTosaOpsConversionPatterns( SPIRVTypeConverter &typeConverter, RewritePatternSet &patterns) { patterns.add< TosaOpConvert>, TosaOpConvert, TosaOpConvert>, TosaOpConvert>, TosaOpConvert>, TosaMultiResultOpConvert, TosaOpConvert, TosaOpConvert, TosaMultiResultOpConvert, TosaOpConvert, TosaOpConvert, TosaOpConvert>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert, TosaOpConvert>, TosaOpConvert< tosa::BitwiseOrOp, replaceBinaryElementwise>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert< tosa::LogicalOrOp, replaceBinaryElementwise>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert, TosaOpConvert>, TosaOpConvert>, TosaOpConvert, TosaOpConvert>, TosaOpConvert< tosa::BitwiseNotOp, replaceUnaryInput1>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert< tosa::LogicalNotOp, replaceUnaryInput1>, TosaOpConvert, TosaOpConvert< tosa::ReciprocalOp, replaceUnaryInput1>, TosaOpConvert>, TosaOpConvert>, TosaOpConvert, TosaOpConvert>, TosaOpConvert< tosa::GreaterOp, replaceBinaryElementwise>, TosaOpConvert>, TosaOpConvert< tosa::ReduceAllOp, replaceReduction>, TosaOpConvert< tosa::ReduceAnyOp, replaceReduction>, TosaOpConvert< tosa::ReduceMaxOp, replaceNanModeReduction>, TosaOpConvert< tosa::ReduceMinOp, replaceNanModeReduction>, TosaOpConvert< tosa::ReduceProductOp, replaceReduction>, TosaOpConvert< tosa::ReduceSumOp, replaceReduction>, TosaOpConvert, TosaOpConvert, TosaOpConvert, TosaOpConvert, TosaOpConvert, TosaOpConvert, TosaOpConvert, TosaOpConvert, TosaOpConvert, TosaOpConvert, TosaOpConvert>, TosaOpConvert, TosaOpConvert>, TosaOpConvert>, TosaOpConvert>(typeConverter, patterns.getContext()); } } // namespace mlir::tosa