112 lines · cpp
1//===- TosaArithConstantToConst.cpp ---------------------------------------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8//9// This file implements a pass that converts tensor-valued arith.constant ops10// into tosa.const so that TOSA pipelines operate on a uniform constant form.11//12//===----------------------------------------------------------------------===//13 14#include "mlir/Dialect/Tosa/Transforms/Passes.h"15 16#include "mlir/Dialect/Arith/IR/Arith.h"17#include "mlir/Dialect/Func/IR/FuncOps.h"18#include "mlir/Dialect/Quant/IR/QuantTypes.h"19#include "mlir/Dialect/Tosa/IR/TosaOps.h"20#include "mlir/IR/BuiltinAttributes.h"21#include "mlir/IR/BuiltinTypes.h"22#include "mlir/IR/PatternMatch.h"23#include "mlir/Transforms/GreedyPatternRewriteDriver.h"24 25namespace mlir {26namespace tosa {27#define GEN_PASS_DEF_TOSAARITHCONSTANTTOTOSACONSTPASS28#include "mlir/Dialect/Tosa/Transforms/Passes.h.inc"29} // namespace tosa30} // namespace mlir31 32using namespace mlir;33using namespace mlir::tosa;34 35namespace {36 37// NOTE: TOSA pipelines already lower their constants through shared Arith38// folding passes, so tensor literals often come back as `arith.constant` even39// after the IR is otherwise TOSA-only. Keep this normalization with the rest of40// the TOSA transforms so any client can re-establish a canonical `tosa.const`41// representation without needing a full Arith->TOSA conversion library.42 43/// Returns true when `elementType` is natively representable by tosa.const.44static bool isSupportedElementType(Type elementType) {45 if (isa<FloatType>(elementType))46 return true;47 48 if (auto intType = dyn_cast<IntegerType>(elementType))49 return intType.isSignless() || intType.isUnsigned();50 51 if (isa<quant::QuantizedType>(elementType))52 return true;53 54 if (isa<tosa::mxint8Type>(elementType))55 return true;56 57 return false;58}59 60class ArithConstantToTosaConst : public OpRewritePattern<arith::ConstantOp> {61public:62 using OpRewritePattern::OpRewritePattern;63 64 LogicalResult matchAndRewrite(arith::ConstantOp constOp,65 PatternRewriter &rewriter) const override {66 // TOSA constant verification requires a ranked, statically shaped tensor.67 auto resultType = dyn_cast<RankedTensorType>(constOp.getResult().getType());68 if (!resultType || !resultType.hasStaticShape())69 return failure();70 71 if (!isSupportedElementType(resultType.getElementType()))72 return failure();73 74 Attribute attr = constOp.getValueAttr();75 auto elementsAttr = dyn_cast<ElementsAttr>(attr);76 if (!elementsAttr)77 return failure();78 79 auto attrType = dyn_cast<RankedTensorType>(elementsAttr.getType());80 if (!attrType || !attrType.hasStaticShape())81 return failure();82 if (attrType != resultType)83 return failure();84 85 auto newConst = tosa::ConstOp::create(rewriter, constOp.getLoc(),86 resultType, elementsAttr);87 rewriter.replaceOp(constOp, newConst.getResult());88 return success();89 }90};91 92struct TosaArithConstantToTosaConstPass93 : public tosa::impl::TosaArithConstantToTosaConstPassBase<94 TosaArithConstantToTosaConstPass> {95 using Base::Base;96 97 void getDependentDialects(DialectRegistry ®istry) const override {98 registry.insert<arith::ArithDialect, tosa::TosaDialect>();99 }100 101 void runOnOperation() override {102 auto *ctx = &getContext();103 RewritePatternSet patterns(ctx);104 patterns.add<ArithConstantToTosaConst>(ctx);105 106 if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))107 signalPassFailure();108 }109};110 111} // namespace112