mirror of https://github.com/llvm/torch-mlir
1573 lines
61 KiB
C++
1573 lines
61 KiB
C++
//===----------------------------------------------------------------------===//
|
|
//
|
|
// 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
|
|
// Also available under a BSD-style license. See LICENSE.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "torch-mlir/Conversion/TorchToTosa/TorchToTosa.h"
|
|
#include "torch-mlir/Conversion/TorchToTosa/TosaLegalizeCommon.h"
|
|
#include "torch-mlir/Conversion/TorchToTosa/TosaLegalizeUtils.h"
|
|
|
|
#include "../PassDetail.h"
|
|
#include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
|
|
#include "mlir/Dialect/Tensor/IR/Tensor.h"
|
|
#include "mlir/Dialect/Tosa/IR/TosaOps.h"
|
|
#include "mlir/Dialect/Traits.h"
|
|
#include "mlir/IR/Matchers.h"
|
|
#include "mlir/Transforms/DialectConversion.h"
|
|
#include "torch-mlir/Dialect/Torch/IR/TorchOps.h"
|
|
#include "torch-mlir/Dialect/TorchConversion/IR/TorchConversionDialect.h"
|
|
#include "torch-mlir/Dialect/TorchConversion/Transforms/BackendTypeConversion.h"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::torch;
|
|
using namespace mlir::torch::Torch;
|
|
|
|
namespace {
|
|
|
|
// These legalizations are for unary ops with only for floating point datatypes.
|
|
// There is no supported quantized integer mode for these.
|
|
template <typename AtenOpT, typename TosaOpT>
|
|
class ConvertAtenUnaryFPOnlyOp : public OpConversionPattern<AtenOpT> {
|
|
public:
|
|
using OpConversionPattern<AtenOpT>::OpConversionPattern;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
LogicalResult
|
|
matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
Value self = adaptor.self();
|
|
auto selfTy = self.getType().cast<TensorType>();
|
|
|
|
if (!selfTy)
|
|
return op.emitError("Only Tensor types supported in TOSA");
|
|
|
|
if (selfTy.getElementType().isa<mlir::FloatType>()) {
|
|
rewriter.replaceOpWithNewOp<TosaOpT>(
|
|
op,
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
op.getType()),
|
|
self);
|
|
return success();
|
|
} else {
|
|
return op.emitError(
|
|
"Only floating-point datatype legalization supported");
|
|
}
|
|
}
|
|
};
|
|
|
|
// These unary op legalizations are identical for floating-point
|
|
// or quantized types
|
|
template <typename AtenOpT, typename TosaOpT>
|
|
class ConvertAtenUnaryOp : public OpConversionPattern<AtenOpT> {
|
|
public:
|
|
using OpConversionPattern<AtenOpT>::OpConversionPattern;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
LogicalResult
|
|
matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
rewriter.replaceOpWithNewOp<TosaOpT>(
|
|
op,
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
op.getType()),
|
|
adaptor.self());
|
|
return success();
|
|
}
|
|
};
|
|
|
|
// These binary op legalizations are identical for floating-point
|
|
// or quantized types
|
|
template <typename AtenOpT, typename TosaOpT>
|
|
class ConvertAtenBinaryOp : public OpConversionPattern<AtenOpT> {
|
|
public:
|
|
using OpConversionPattern<AtenOpT>::OpConversionPattern;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
LogicalResult
|
|
matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
Value lhs = adaptor.self();
|
|
auto lhsTy = lhs.getType().cast<TensorType>();
|
|
Value rhs = adaptor.other();
|
|
auto rhsTy = rhs.getType().cast<TensorType>();
|
|
|
|
if (!lhsTy || !rhsTy)
|
|
return op.emitError("Only Tensor types supported in TOSA");
|
|
|
|
auto lhsElemTy = lhsTy.getElementType();
|
|
auto rhsElemTy = rhsTy.getElementType();
|
|
|
|
if (lhsElemTy != rhsElemTy)
|
|
return op.emitError("Add: input datatypes mismatched");
|
|
|
|
rewriter.replaceOpWithNewOp<TosaOpT>(
|
|
op,
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
op.getType()),
|
|
lhs, rhs);
|
|
return success();
|
|
}
|
|
};
|
|
|
|
// FIXME: This will eventually go into a Tosa*Utils file.
|
|
LogicalResult torchScalarToTosaTensor(ConversionPatternRewriter &rewriter,
|
|
Operation *op, Value torchScalarValue,
|
|
Value &tosaTensor, Type dtype) {
|
|
if (dtype.isa<mlir::FloatType>()) {
|
|
double scalarValue;
|
|
|
|
if (!matchPattern(torchScalarValue, m_TorchConstantFloat(&scalarValue)))
|
|
return failure();
|
|
|
|
tosaTensor =
|
|
mlir::tosa::getTosaConstTensorSingleF32(rewriter, op, scalarValue);
|
|
} else if (auto intType = dtype.dyn_cast<mlir::IntegerType>()) {
|
|
int64_t scalarValue;
|
|
|
|
if (!matchPattern(torchScalarValue, m_TorchConstantInt(&scalarValue)))
|
|
return failure();
|
|
|
|
auto w = intType.getWidth();
|
|
if (w != 32 && w != 64)
|
|
return op->emitError("Unsupported integer type") << intType;
|
|
|
|
if (w == 32) {
|
|
tosaTensor = tosa::getConstTensor<int32_t>(
|
|
rewriter, op, {static_cast<int32_t>(scalarValue)}, {})
|
|
.getValue();
|
|
} else if (w == 64) {
|
|
tosaTensor =
|
|
tosa::getConstTensor<int64_t>(rewriter, op, {scalarValue}, {})
|
|
.getValue();
|
|
}
|
|
return success();
|
|
} else
|
|
return op->emitError("Usupported element type");
|
|
|
|
return success();
|
|
}
|
|
|
|
LogicalResult torchAlphaToTosaTensor(ConversionPatternRewriter &rewriter,
|
|
Operation *op, Value alphaScalar,
|
|
Value &alphaTensor, Type dtype,
|
|
bool checkForUnity) {
|
|
if (succeeded(torchScalarToTosaTensor(rewriter, op, alphaScalar, alphaTensor,
|
|
dtype)))
|
|
return success();
|
|
|
|
// `alpha` has not been specified.
|
|
int64_t alphaValue;
|
|
if (!matchPattern(alphaScalar, m_TorchConstantInt(&alphaValue)))
|
|
return op->emitError("Currently only scalar constants are supported for "
|
|
"alpha in TOSA operation");
|
|
// When no alpha has been specified, this must be 1.
|
|
if (checkForUnity && alphaValue != 1)
|
|
return op->emitError("Unsupported integer value for alpha");
|
|
|
|
alphaTensor =
|
|
mlir::tosa::getTosaConstTensorSingleF32(rewriter, op, alphaValue);
|
|
|
|
return success();
|
|
}
|
|
|
|
// These binary op legalizations are specific to add/sub which have an
|
|
// alpha multiplier.
|
|
template <typename AtenOpT, typename TosaOpT>
|
|
class ConvertAtenAddSubOp : public OpConversionPattern<AtenOpT> {
|
|
public:
|
|
using OpConversionPattern<AtenOpT>::OpConversionPattern;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
LogicalResult
|
|
matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
Value lhs = adaptor.self();
|
|
auto lhsTy = lhs.getType().dyn_cast<TensorType>();
|
|
Value rhs = adaptor.other();
|
|
auto rhsTy = rhs.getType().dyn_cast<TensorType>();
|
|
|
|
if (!lhsTy)
|
|
return op.emitError("Only Tensor types supported in TOSA");
|
|
|
|
auto lhsElemTy = lhsTy.getElementType();
|
|
if (!lhsElemTy.isIntOrFloat())
|
|
return op.emitError(
|
|
"Only floating-point or integer datatype legalization supported");
|
|
|
|
Value rhsAsTensor;
|
|
if (!rhsTy) {
|
|
if (failed(torchScalarToTosaTensor(rewriter, op.getOperation(),
|
|
op.other(), rhsAsTensor, lhsElemTy)))
|
|
return op.emitError("Currently only scalar constants are supported for "
|
|
"conversion in TOSA operation");
|
|
}
|
|
auto rhsTensor = rhsTy ? rhs : rhsAsTensor;
|
|
|
|
// Handle alpha.
|
|
Value alphaTensor;
|
|
if (failed(torchAlphaToTosaTensor(rewriter, op.getOperation(), op.alpha(),
|
|
alphaTensor, lhsElemTy, false)))
|
|
return op.emitError("Currently only scalar constants are supported for "
|
|
"alpha in conversion to TOSA operation");
|
|
|
|
auto multTensor = rewriter.create<tosa::MulOp>(
|
|
op.getLoc(), rhsTy ? rhsTy : RankedTensorType::get({}, lhsElemTy),
|
|
rhsTensor, alphaTensor, /*shift*/ 0);
|
|
|
|
if (lhsElemTy.isa<mlir::FloatType>()) {
|
|
rewriter.replaceOpWithNewOp<TosaOpT>(
|
|
op,
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
op.getType()),
|
|
lhs, multTensor);
|
|
return success();
|
|
} else {
|
|
return op.emitError(
|
|
"Only floating-point datatype legalization supported");
|
|
}
|
|
}
|
|
}; // namespace
|
|
|
|
// Binary op legalizations for comparator ops.
|
|
template <typename AtenOpT, typename TosaOpT>
|
|
class ConvertAtenCompareOp : public OpConversionPattern<AtenOpT> {
|
|
public:
|
|
using OpConversionPattern<AtenOpT>::OpConversionPattern;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
LogicalResult
|
|
matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
Value lhs = adaptor.self();
|
|
auto lhsTy = lhs.getType().dyn_cast<TensorType>();
|
|
Value rhs = adaptor.other();
|
|
auto rhsTy = rhs.getType().dyn_cast<TensorType>();
|
|
|
|
if (!lhsTy)
|
|
return op.emitError("Only Tensor types supported in TOSA");
|
|
|
|
auto lhsElemTy = lhsTy.getElementType();
|
|
if (!lhsElemTy.isIntOrFloat())
|
|
return op.emitError(
|
|
"Only floating-point or integer datatype legalization supported");
|
|
|
|
Value rhsAsTensor;
|
|
if (!rhsTy) {
|
|
if (failed(torchScalarToTosaTensor(rewriter, op.getOperation(),
|
|
op.other(), rhsAsTensor, lhsElemTy)))
|
|
return op.emitError("Currently only scalar constants are supported for "
|
|
"conversion in TOSA operation");
|
|
}
|
|
auto rhsTensor = rhsTy ? rhs : rhsAsTensor;
|
|
// There is no Lesser operator in TOSA
|
|
auto swapLhsRhs = (std::is_same<AtenOpT, AtenLtTensorOp>() ||
|
|
std::is_same<AtenOpT, AtenLtScalarOp>());
|
|
|
|
rewriter.replaceOpWithNewOp<TosaOpT>(
|
|
op,
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
op.getType()),
|
|
(swapLhsRhs ? rhsTensor : lhs), (swapLhsRhs ? lhs : rhsTensor));
|
|
return success();
|
|
}
|
|
};
|
|
|
|
// Binary op legalizations for Mul variants.
|
|
template <typename AtenOpT>
|
|
class ConvertAtenMulOp : public OpConversionPattern<AtenOpT> {
|
|
public:
|
|
using OpConversionPattern<AtenOpT>::OpConversionPattern;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
LogicalResult
|
|
matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
Value lhs = adaptor.self();
|
|
auto lhsTy = lhs.getType().dyn_cast<TensorType>();
|
|
Value rhs = adaptor.other();
|
|
auto rhsTy = rhs.getType().dyn_cast<TensorType>();
|
|
|
|
if (!lhsTy)
|
|
return op.emitError("Only Tensor types supported in TOSA");
|
|
|
|
auto lhsElemTy = lhsTy.getElementType();
|
|
if (!lhsElemTy.isIntOrFloat())
|
|
return op.emitError(
|
|
"Only floating-point or integer datatype legalization supported");
|
|
|
|
Value rhsAsTensor;
|
|
if (!rhsTy) {
|
|
if (failed(torchScalarToTosaTensor(rewriter, op.getOperation(),
|
|
op.other(), rhsAsTensor, lhsElemTy)))
|
|
return op.emitError("Currently only scalar constants are supported for "
|
|
"conversion in TOSA operation");
|
|
}
|
|
auto rhsTensor = rhsTy ? rhs : rhsAsTensor;
|
|
|
|
if (lhsElemTy.isa<mlir::FloatType>() ||
|
|
lhsElemTy.isa<mlir::IntegerType>()) {
|
|
rewriter.replaceOpWithNewOp<tosa::MulOp>(
|
|
op,
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
op.getType()),
|
|
lhs, rhsTensor,
|
|
/*shift=*/0);
|
|
return success();
|
|
} else {
|
|
// Quantized multiplication may need to rescale inputs.
|
|
return op.emitError("Only floating-point or integer datatype "
|
|
"legalization currently supported");
|
|
}
|
|
}
|
|
};
|
|
|
|
template <typename AtenOpT>
|
|
class ConvertAtenDivOp : public OpConversionPattern<AtenOpT> {
|
|
public:
|
|
using OpConversionPattern<AtenOpT>::OpConversionPattern;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
LogicalResult
|
|
matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
Value lhs = adaptor.self();
|
|
auto lhsTy = lhs.getType().dyn_cast<TensorType>();
|
|
Value rhs = adaptor.other();
|
|
auto rhsTy = rhs.getType().dyn_cast<TensorType>();
|
|
|
|
if (!lhsTy)
|
|
return op.emitError("Only Tensor types supported in TOSA");
|
|
|
|
auto lhsElemTy = lhsTy.getElementType();
|
|
if (!lhsElemTy.isIntOrFloat())
|
|
return op.emitError(
|
|
"Only floating-point or integer datatype legalization supported");
|
|
|
|
Value rhsAsTensor;
|
|
if (!rhsTy) {
|
|
if (failed(torchScalarToTosaTensor(rewriter, op.getOperation(),
|
|
op.other(), rhsAsTensor, lhsElemTy)))
|
|
return op.emitError("Currently only scalar constants are supported for "
|
|
"conversion in TOSA operation");
|
|
}
|
|
auto rhsTensor = rhsTy ? rhs : rhsAsTensor;
|
|
|
|
if (lhsElemTy.isa<mlir::FloatType>()) {
|
|
auto rcpOp = rewriter.create<tosa::ReciprocalOp>(
|
|
op->getLoc(), rhsTy ? rhsTy : RankedTensorType::get({}, lhsElemTy),
|
|
rhsTensor);
|
|
rewriter.replaceOpWithNewOp<tosa::MulOp>(
|
|
op,
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
op.getType()),
|
|
lhs, rcpOp.getResult(), /*shift=*/0);
|
|
} else {
|
|
rewriter.replaceOpWithNewOp<tosa::DivOp>(
|
|
op,
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
op.getType()),
|
|
lhs, rhsTensor);
|
|
}
|
|
return success();
|
|
}
|
|
};
|
|
|
|
// This defines a template to construct ops whose legalizations are
|
|
// specialized.
|
|
template <typename AtenOpT>
|
|
class ConvertAtenOp : public OpConversionPattern<AtenOpT> {
|
|
public:
|
|
using OpConversionPattern<AtenOpT>::OpConversionPattern;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
LogicalResult
|
|
matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override;
|
|
};
|
|
|
|
template <>
|
|
LogicalResult ConvertAtenOp<AtenTanhOp>::matchAndRewrite(
|
|
AtenTanhOp op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
Value self = adaptor.self();
|
|
auto selfTy = self.getType().cast<TensorType>();
|
|
if (selfTy && selfTy.getElementType().isa<mlir::FloatType>()) {
|
|
rewriter.replaceOpWithNewOp<tosa::TanhOp>(
|
|
op, getTypeConverter()->convertType(op.getType()), self);
|
|
return success();
|
|
} else {
|
|
// Sigmoid legalization in TOSA for quantized element-type uses
|
|
// specialized tosa.table construct.
|
|
return op.emitError(
|
|
"Only floating-point datatype legalization currently supported");
|
|
}
|
|
}
|
|
|
|
template <>
|
|
LogicalResult ConvertAtenOp<AtenSigmoidOp>::matchAndRewrite(
|
|
AtenSigmoidOp op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
Value self = adaptor.self();
|
|
auto selfTy = self.getType().cast<TensorType>();
|
|
if (selfTy && selfTy.getElementType().isa<mlir::FloatType>()) {
|
|
rewriter.replaceOpWithNewOp<tosa::SigmoidOp>(
|
|
op, getTypeConverter()->convertType(op.getType()), self);
|
|
return success();
|
|
} else {
|
|
// Sigmoid legalization in TOSA for quantized element-type uses
|
|
// specialized tosa.table construct.
|
|
return op.emitError(
|
|
"Only floating-point datatype legalization currently supported");
|
|
}
|
|
}
|
|
|
|
template <>
|
|
LogicalResult ConvertAtenOp<AtenReluOp>::matchAndRewrite(
|
|
AtenReluOp op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
Value self = adaptor.self();
|
|
auto selfTy = self.getType().cast<TensorType>();
|
|
|
|
// Maps to tosa.clamp which has both int and fp limits.
|
|
int64_t clampMin = 0;
|
|
Value clampIn = self;
|
|
if (selfTy) {
|
|
// Rescale the clampIn for quantized types. TBD
|
|
if (!selfTy.getElementType().isa<mlir::FloatType>()) {
|
|
return op.emitError(
|
|
"Only floating-point datatype legalization currently supported");
|
|
}
|
|
rewriter.replaceOpWithNewOp<tosa::ClampOp>(
|
|
op, getTypeConverter()->convertType(op.getType()), clampIn,
|
|
rewriter.getI64IntegerAttr(clampMin),
|
|
rewriter.getI64IntegerAttr(std::numeric_limits<int32_t>::max()),
|
|
rewriter.getF32FloatAttr(0.0f),
|
|
rewriter.getF32FloatAttr(std::numeric_limits<float>::max()));
|
|
return success();
|
|
} else {
|
|
return op.emitError("Only Tensor types supported in TOSA");
|
|
}
|
|
}
|
|
|
|
using ReductionConvFunc = llvm::Optional<Value> (*)(PatternRewriter &,
|
|
Operation *,
|
|
RankedTensorType, Value,
|
|
ElementsAttr, bool);
|
|
|
|
// They all constitute a common form invoking the appropriate
|
|
// converion function in TosaLegalizeCommon.cpp
|
|
template <typename AtenOpT, ReductionConvFunc ConversionFuncT>
|
|
class ConvertAtenReductionOp : public OpConversionPattern<AtenOpT> {
|
|
public:
|
|
using OpConversionPattern<AtenOpT>::OpConversionPattern;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
|
|
// Each variant must implement corresponding parameter parsing options
|
|
virtual LogicalResult readReduceDimsAndKeepDims(
|
|
AtenOpT op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter,
|
|
ElementsAttr &reduceDimsAttr, bool &keepDims) const {
|
|
return rewriter.notifyMatchFailure(
|
|
op, "Unimplemented reduce_dims and keep_dims parsing function");
|
|
}
|
|
|
|
// Common rewriter for all reduction ops, calls the specific implementation of
|
|
// readReduceDimsAndKeepDims() needed for the op variant.
|
|
LogicalResult
|
|
matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
Value self = adaptor.self();
|
|
auto selfTy = self.getType().cast<TensorType>();
|
|
|
|
if (!selfTy)
|
|
return op.emitError("Only Tensor types supported in TOSA");
|
|
|
|
auto outputTy = OpConversionPattern<AtenOpT>::getTypeConverter()
|
|
->convertType(op.getType())
|
|
.template cast<RankedTensorType>();
|
|
if (!outputTy)
|
|
return op.emitError(
|
|
"Only ranked tensor type outputs permitted for reduce_mean");
|
|
|
|
ElementsAttr reduceDimsAttr;
|
|
bool keepDims;
|
|
|
|
if (failed(readReduceDimsAndKeepDims(op, adaptor, rewriter, reduceDimsAttr,
|
|
keepDims)))
|
|
return failure();
|
|
|
|
llvm::Optional<Value> result =
|
|
ConversionFuncT(rewriter, op, outputTy, self, reduceDimsAttr, keepDims);
|
|
|
|
if (!result)
|
|
return failure();
|
|
|
|
// TBD - support dtype casting.
|
|
|
|
rewriter.replaceOp(op, {result.getValue()});
|
|
|
|
return success();
|
|
}
|
|
};
|
|
|
|
// This reduction op legalization template handles op variants that have
|
|
// explicit reduce_dims dimensions (provided as a list) and keep_dims
|
|
// parameters.
|
|
template <typename AtenOpT, ReductionConvFunc ConversionFuncT>
|
|
class ConvertAtenMultipleDimsReductionOp
|
|
: public ConvertAtenReductionOp<AtenOpT, ConversionFuncT> {
|
|
using ConvertAtenReductionOp<AtenOpT,
|
|
ConversionFuncT>::ConvertAtenReductionOp;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
LogicalResult readReduceDimsAndKeepDims(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter,
|
|
ElementsAttr &reduceDimsAttr,
|
|
bool &keepDims) const override {
|
|
SmallVector<int64_t, 4> reduceDims;
|
|
if (!matchPattern(op.dim(), m_TorchConstantIntList(reduceDims)))
|
|
return rewriter.notifyMatchFailure(op,
|
|
"non-const dim parameter unsupported");
|
|
int64_t N = reduceDims.size();
|
|
auto reduceDimsType = RankedTensorType::get({N}, rewriter.getI64Type());
|
|
reduceDimsAttr = DenseIntElementsAttr::get(reduceDimsType,
|
|
llvm::makeArrayRef(reduceDims));
|
|
|
|
keepDims = false;
|
|
if (!matchPattern(op.keepdim(), m_TorchConstantBool(&keepDims)))
|
|
return rewriter.notifyMatchFailure(
|
|
op, "non-const keepdim parameter unsupported");
|
|
|
|
return success();
|
|
}
|
|
};
|
|
|
|
// This reduction op legalization template handles op variants that reduce in
|
|
// only one explicit dim which is provided as a number (rather than a list), and
|
|
// a keep_dims parameter.
|
|
template <typename AtenOpT, ReductionConvFunc ConversionFuncT>
|
|
class ConvertAtenOneDimReductionOp
|
|
: public ConvertAtenReductionOp<AtenOpT, ConversionFuncT> {
|
|
using ConvertAtenReductionOp<AtenOpT,
|
|
ConversionFuncT>::ConvertAtenReductionOp;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
LogicalResult readReduceDimsAndKeepDims(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter,
|
|
ElementsAttr &reduceDimsAttr,
|
|
bool &keepDims) const override {
|
|
int64_t reduceDim;
|
|
if (!matchPattern(op.dim(), m_TorchConstantInt(&reduceDim)))
|
|
return rewriter.notifyMatchFailure(op,
|
|
"non-const dim parameter unsupported");
|
|
auto reduceDimsType = RankedTensorType::get({1}, rewriter.getI64Type());
|
|
reduceDimsAttr = DenseIntElementsAttr::get(reduceDimsType,
|
|
llvm::makeArrayRef({reduceDim}));
|
|
|
|
keepDims = false;
|
|
if (!matchPattern(op.keepdim(), m_TorchConstantBool(&keepDims)))
|
|
return rewriter.notifyMatchFailure(
|
|
op, "non-const keepdim parameter unsupported");
|
|
|
|
return success();
|
|
}
|
|
};
|
|
|
|
// This reduction op legalization template handles op variants that reduce all
|
|
// dims does not keep dims.
|
|
template <typename AtenOpT, ReductionConvFunc ConversionFuncT>
|
|
class ConvertAtenAllDimsReductionOp
|
|
: public ConvertAtenReductionOp<AtenOpT, ConversionFuncT> {
|
|
public:
|
|
using ConvertAtenReductionOp<AtenOpT,
|
|
ConversionFuncT>::ConvertAtenReductionOp;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
LogicalResult readReduceDimsAndKeepDims(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter,
|
|
ElementsAttr &reduceDimsAttr,
|
|
bool &keepDims) const override {
|
|
auto self = adaptor.self();
|
|
auto selfTy = self.getType().template cast<RankedTensorType>();
|
|
|
|
// Select all dims to reduce
|
|
SmallVector<int64_t, 4> reduceDims;
|
|
for (int64_t i = 0; i < selfTy.getRank(); i++)
|
|
reduceDims.push_back(i);
|
|
int64_t N = selfTy.getRank();
|
|
auto reduceDimsType = RankedTensorType::get({N}, rewriter.getI64Type());
|
|
reduceDimsAttr = DenseIntElementsAttr::get(reduceDimsType,
|
|
llvm::makeArrayRef(reduceDims));
|
|
keepDims = false;
|
|
|
|
return success();
|
|
}
|
|
};
|
|
|
|
template <>
|
|
LogicalResult ConvertAtenOp<AtenArgmaxOp>::matchAndRewrite(
|
|
AtenArgmaxOp op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
|
|
Value self = adaptor.self();
|
|
auto selfTy = self.getType().template cast<RankedTensorType>();
|
|
|
|
if (!selfTy)
|
|
return op.emitError("Only ranked tensor types supported in TOSA argmax");
|
|
|
|
int64_t reduceDim;
|
|
if (!matchPattern(op.dim(), m_TorchConstantInt(&reduceDim))) {
|
|
// NoneType indicates reduce on all dims
|
|
reduceDim = -1;
|
|
}
|
|
|
|
bool keepDim = false;
|
|
if (!matchPattern(op.keepdim(), m_TorchConstantBool(&keepDim)))
|
|
return rewriter.notifyMatchFailure(
|
|
op, "non-const keepdim parameter unsupported");
|
|
|
|
auto resultTy = getTypeConverter()
|
|
->convertType(op.getResult().getType())
|
|
.cast<RankedTensorType>();
|
|
auto outputETy = resultTy.getElementType();
|
|
|
|
// Create a single instance of tosa.argmax.
|
|
// Multiple dims require chained construct.
|
|
auto buildArgmax = [&](int64_t reduceDim, Value input) -> Value {
|
|
auto inputTy = input.getType().cast<RankedTensorType>();
|
|
auto inputShape = inputTy.getShape();
|
|
SmallVector<int64_t> outputShapeArr = {};
|
|
int32_t i = 0;
|
|
|
|
for (auto &dim : inputShape) {
|
|
if (i++ != reduceDim) {
|
|
outputShapeArr.push_back(dim);
|
|
} else {
|
|
if (keepDim)
|
|
outputShapeArr.push_back(1);
|
|
}
|
|
}
|
|
|
|
// Tosa argmax output is i32, while Torch backend mandates i64.
|
|
auto outputReduceTy = RankedTensorType::get(
|
|
ArrayRef<int64_t>(outputShapeArr), rewriter.getI32Type());
|
|
auto reduceDimAttr =
|
|
rewriter.getIntegerAttr(rewriter.getI64Type(), reduceDim);
|
|
return rewriter
|
|
.create<tosa::ArgMaxOp>(op->getLoc(),
|
|
getTypeConverter()->convertType(outputReduceTy),
|
|
input, reduceDimAttr)
|
|
.getResult();
|
|
};
|
|
|
|
// Convert the final index to i64 for backend finalization, However, i64
|
|
// is not a defined type for tosa.cast, so using arith.extsi instead.
|
|
auto castToInt64 = [&](Value result) -> LogicalResult {
|
|
auto resTy = result.getType().cast<ShapedType>();
|
|
if (!resTy)
|
|
return op.emitError("Argmax: Result is not a shaped type");
|
|
|
|
auto resShape = resTy.getShape();
|
|
auto outTy =
|
|
RankedTensorType::get(resShape, outputETy); // rewriter.getI64Type());
|
|
|
|
rewriter.replaceOpWithNewOp<arith::ExtSIOp>(
|
|
op, getTypeConverter()->convertType(outTy), result);
|
|
|
|
return success();
|
|
};
|
|
|
|
if (reduceDim == -1) { // reducing on all dims
|
|
Value input = self;
|
|
for (int dim = 0; dim < selfTy.getRank(); dim++) {
|
|
// progressively reduce each 0-th dim
|
|
input = buildArgmax(0, input);
|
|
}
|
|
return castToInt64(input);
|
|
} else {
|
|
return castToInt64(buildArgmax(reduceDim, self));
|
|
}
|
|
|
|
return success();
|
|
}
|
|
|
|
template <typename AtenOpT>
|
|
class ConvertAtenSqueezeOp : public OpConversionPattern<AtenOpT> {
|
|
public:
|
|
using OpConversionPattern<AtenOpT>::OpConversionPattern;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
|
|
// Each variant must implement corresponding parameter parsing options
|
|
virtual LogicalResult
|
|
generateSqueezedShape(AtenOpT op, RankedTensorType selfTy,
|
|
ConversionPatternRewriter &rewriter,
|
|
SmallVector<int64_t> &squeezedShape) const {
|
|
return rewriter.notifyMatchFailure(
|
|
op, "Unimplemented dim/dim-list parsing function");
|
|
}
|
|
|
|
// Common rewriter for all squeeze ops, calls the specific implementation of
|
|
// generateSqueezedShape() needed for the op variant.
|
|
LogicalResult
|
|
matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
Value self = adaptor.self();
|
|
auto selfTy = self.getType().template cast<RankedTensorType>();
|
|
|
|
if (!selfTy)
|
|
return op.emitError("Only ranked tensor types supported in TOSA argmax");
|
|
|
|
SmallVector<int64_t> newOutputShape;
|
|
if (failed(generateSqueezedShape(op, selfTy, rewriter, newOutputShape)))
|
|
return op.emitError("Squeeze could not compute new shape");
|
|
|
|
auto resultTy = OpConversionPattern<AtenOpT>::getTypeConverter()
|
|
->convertType(op.getResult().getType())
|
|
.template cast<RankedTensorType>();
|
|
auto resultElemTy = resultTy.getElementType();
|
|
|
|
auto newOutputTy = RankedTensorType::get(newOutputShape, resultElemTy);
|
|
|
|
auto reshapeOp = rewriter.create<tosa::ReshapeOp>(
|
|
op->getLoc(),
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
newOutputTy),
|
|
self, rewriter.getI64ArrayAttr(newOutputShape));
|
|
rewriter.replaceOpWithNewOp<tensor::CastOp>(
|
|
op,
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
newOutputTy),
|
|
reshapeOp);
|
|
|
|
return success();
|
|
}
|
|
};
|
|
|
|
template <typename AtenOpT>
|
|
class ConvertAtenSqueezeOneDimOp : public ConvertAtenSqueezeOp<AtenOpT> {
|
|
using ConvertAtenSqueezeOp<AtenOpT>::ConvertAtenSqueezeOp;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
|
|
LogicalResult
|
|
generateSqueezedShape(AtenOpT op, RankedTensorType selfTy,
|
|
ConversionPatternRewriter &rewriter,
|
|
SmallVector<int64_t> &squeezedShape) const override {
|
|
int64_t squeezeDim;
|
|
if (!matchPattern(op.dim(), m_TorchConstantInt(&squeezeDim)))
|
|
return rewriter.notifyMatchFailure(op,
|
|
"non-const dim parameter unsupported");
|
|
|
|
// Handle negative dim
|
|
if (squeezeDim < 0)
|
|
squeezeDim = squeezeDim + selfTy.getRank();
|
|
|
|
auto selfShape = selfTy.getShape();
|
|
|
|
// Only dims statically known to have size=1 are reduced.
|
|
// Dynamic dims are treated as unknowns and will not be squeezed
|
|
// even if dim parameter says it should be.
|
|
uint32_t dimNum = 0;
|
|
for (auto &dim : selfShape) {
|
|
if (dim != 1 || squeezeDim != dimNum)
|
|
squeezedShape.push_back(dim);
|
|
dimNum++;
|
|
}
|
|
|
|
return success();
|
|
}
|
|
};
|
|
|
|
template <typename AtenOpT>
|
|
class ConvertAtenSqueezeAllDimsOp : public ConvertAtenSqueezeOp<AtenOpT> {
|
|
using ConvertAtenSqueezeOp<AtenOpT>::ConvertAtenSqueezeOp;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
|
|
LogicalResult
|
|
generateSqueezedShape(AtenOpT op, RankedTensorType selfTy,
|
|
ConversionPatternRewriter &rewriter,
|
|
SmallVector<int64_t> &squeezedShape) const override {
|
|
auto selfShape = selfTy.getShape();
|
|
|
|
// Dims that may dynamically resolve to 1 are not reduced here. Only
|
|
// compile-time resolvable dims are handled here.
|
|
for (auto &dim : selfShape) {
|
|
if (dim != 1)
|
|
squeezedShape.push_back(dim);
|
|
}
|
|
return success();
|
|
}
|
|
};
|
|
|
|
template <>
|
|
LogicalResult ConvertAtenOp<AtenPowTensorScalarOp>::matchAndRewrite(
|
|
AtenPowTensorScalarOp op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
|
|
Value self = adaptor.self();
|
|
auto selfTy = self.getType().template cast<RankedTensorType>();
|
|
|
|
if (!selfTy)
|
|
return op.emitError("Only ranked tensor types supported in TOSA Pow");
|
|
|
|
if (!selfTy.getElementType().isa<mlir::FloatType>())
|
|
return op.emitError("Only floating-point datatype legalization supported");
|
|
|
|
Value expTensor;
|
|
Value expScalar = op.exponent();
|
|
if (failed(torchScalarToTosaTensor(rewriter, op.getOperation(), expScalar,
|
|
expTensor, selfTy.getElementType())))
|
|
return op.emitError("Currently only scalar constants are supported for "
|
|
"conversion in TOSA Pow operation");
|
|
|
|
rewriter.replaceOpWithNewOp<tosa::PowOp>(
|
|
op, getTypeConverter()->convertType(op.getType()), self, expTensor);
|
|
|
|
return success();
|
|
}
|
|
|
|
// Perform torch matmul, mm and bmm
|
|
template <typename AtenOpT>
|
|
class ConvertAtenMatMulOp : public OpConversionPattern<AtenOpT> {
|
|
public:
|
|
using OpConversionPattern<AtenOpT>::OpConversionPattern;
|
|
using OpAdaptor = typename AtenOpT::Adaptor;
|
|
LogicalResult
|
|
matchAndRewrite(AtenOpT op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const override {
|
|
Value lhs = adaptor.self();
|
|
auto lhsTy = lhs.getType().cast<RankedTensorType>();
|
|
|
|
// Aten matmul, mm and bmm call operand2 by different names.
|
|
Value rhs = adaptor.getOperands()[1];
|
|
auto rhsTy = rhs.getType().cast<RankedTensorType>();
|
|
|
|
if (!lhsTy || !rhsTy)
|
|
return op.emitError("Only ranked tensor types supported in TOSA matmul");
|
|
|
|
auto lhsRank = lhsTy.getRank();
|
|
auto rhsRank = rhsTy.getRank();
|
|
|
|
// Mm takes two 2D tensors
|
|
if (isa<AtenMmOp>(op)) {
|
|
assert(lhsRank == 2 && rhsRank == 2 &&
|
|
"aten.mm called but matrix rank != 2");
|
|
}
|
|
|
|
// Bmm takes two 2D tensors
|
|
if (isa<AtenBmmOp>(op)) {
|
|
assert(lhsRank == 3 && rhsRank == 3 &&
|
|
"aten.bmm called but matrix rank != 2");
|
|
}
|
|
|
|
auto lhsShape = lhsTy.getShape();
|
|
auto rhsShape = rhsTy.getShape();
|
|
|
|
auto lhsElemTy = lhsTy.getElementType();
|
|
auto rhsElemTy = rhsTy.getElementType();
|
|
|
|
if (lhsElemTy != rhsElemTy)
|
|
return op.emitError("Matmul: input datatypes mismatched");
|
|
|
|
// Legalization constructs may offer input shapes but expect output shapes
|
|
// to be inferred, e.g.
|
|
// func @forward(%arg0: !torch.vtensor<[14,19],f32>,
|
|
// %arg1: !torch.vtensor<[19,28],f32>) ->
|
|
// !torch.vtensor<[?,?],f32>
|
|
// This is tricky with matmul, since TOSA matmul is on 3D inputs.
|
|
// This means the need to reshape potentially both inputs and outputs,
|
|
// and reshape to unknown shape is undefined.
|
|
|
|
auto maxInputRank = lhsRank > rhsRank ? lhsRank : rhsRank;
|
|
// If performing dot product on vectors, the RHS is synthetically transposed
|
|
if (maxInputRank == 1)
|
|
maxInputRank++;
|
|
|
|
// Obtaining the rank broadcasted shapes of tensors makes it easier to
|
|
// construct the input and output reshaping logic.
|
|
auto getRankBroadcastedShape = [&](Value tensor,
|
|
bool isRHS) -> SmallVector<int64_t> {
|
|
auto tensorTy = tensor.getType().cast<TensorType>();
|
|
auto tensorShape = tensorTy.getShape();
|
|
auto tensorRank = tensorTy.getRank();
|
|
|
|
SmallVector<int64_t> bcastedShape;
|
|
|
|
auto bcastDims = maxInputRank - tensorRank;
|
|
|
|
if (isRHS && (tensorRank == 1) && bcastDims) {
|
|
// RHS with rank1 is special. It be synthetically transposed to dim[:-2]
|
|
for (int32_t i = 0; i < bcastDims - 1; i++)
|
|
bcastedShape.push_back(1);
|
|
bcastedShape.push_back(tensorShape[0]);
|
|
bcastedShape.push_back(1);
|
|
} else {
|
|
if (bcastDims > 0) { // rank broadcast
|
|
for (uint32_t i = 0; i < bcastDims; i++)
|
|
bcastedShape.push_back(1);
|
|
}
|
|
for (auto &dim : tensorShape)
|
|
bcastedShape.push_back(dim);
|
|
}
|
|
return bcastedShape;
|
|
};
|
|
|
|
// Step: Rank broadcast the two inputs.
|
|
auto lhsBroadcastedShape = getRankBroadcastedShape(lhs, false);
|
|
auto lhsBroadcastedTy =
|
|
RankedTensorType::get(lhsBroadcastedShape, lhsElemTy);
|
|
auto rhsBroadcastedShape = getRankBroadcastedShape(rhs, true);
|
|
auto rhsBroadcastedTy =
|
|
RankedTensorType::get(rhsBroadcastedShape, rhsElemTy);
|
|
|
|
auto rankBroadcastedLhs =
|
|
lhsRank == maxInputRank
|
|
? lhs
|
|
: rewriter.create<tosa::ReshapeOp>(
|
|
op->getLoc(),
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
lhsBroadcastedTy),
|
|
lhs, rewriter.getI64ArrayAttr(lhsBroadcastedShape));
|
|
|
|
auto rankBroadcastedRhs =
|
|
rhsRank == maxInputRank
|
|
? rhs
|
|
: rewriter.create<tosa::ReshapeOp>(
|
|
op->getLoc(),
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
rhsBroadcastedTy),
|
|
rhs, rewriter.getI64ArrayAttr(rhsBroadcastedShape));
|
|
|
|
// TOSA matmul is performed on two 3D inputs and generates a 3D output.
|
|
// Lower ranked tensors are dim-1 reshaped up to 3D
|
|
auto reshapeUpTo3DTensor = [&](Value tensor) -> Value {
|
|
auto tensorTy = tensor.getType().cast<TensorType>();
|
|
auto rank = tensorTy.getRank();
|
|
|
|
assert(rank <= 3 && "reshapeUpTo3D tensor must receive rank <= 3");
|
|
if (rank == 3)
|
|
return tensor;
|
|
|
|
auto shape = tensorTy.getShape();
|
|
SmallVector<int64_t> newShape({1, 1, 1});
|
|
|
|
if (rank == 2) { // batchsize = 1
|
|
newShape[1] = shape[0];
|
|
newShape[2] = shape[1];
|
|
} else { // rank 1
|
|
newShape[2] = shape[0];
|
|
}
|
|
auto newType = RankedTensorType::get(newShape, tensorTy.getElementType());
|
|
|
|
return rewriter.create<tosa::ReshapeOp>(
|
|
op->getLoc(),
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
newType),
|
|
tensor, rewriter.getI64ArrayAttr(newShape));
|
|
};
|
|
|
|
// Where broadcasting is required in one or more batch dims, the following
|
|
// is done.
|
|
// Where all batch dims are involved in broadcasting:
|
|
// Given A: 3x1x5x6 and B: 1x4x6x7
|
|
// 1. Reshape A to 1x15x6 (squeeze all batchdims into dim1)
|
|
// 2. Transpose B to 6x1x4x7, Reshape to 1x6x28
|
|
// 3. tosa.Matmul 1x15x6 1x6x28 = 1x15x28
|
|
// 4. Reshape out to 3x5x4x7, Transpose to 3x4x5x7
|
|
// Where there are batch dimensions that are broadcast and not, the
|
|
// treatment is to have dim0 correspond to product of all non-broadcast
|
|
// dimsizes:
|
|
// Given A: 4x8x16x32 B: 8x32x17
|
|
// 1. Reshape A to 8x64x32 (squeeze all unbroadcasted dims into dim0,
|
|
// broadcasted dims into dim1)
|
|
// 2. No transpose or reshape of B as its batchdims are not broadcast to.
|
|
// 3. tosa.Matmul 8x64x32 8x32x17 = 8x64x17
|
|
// 4. Reshape to 8x4x16x17, Transpose to 4x8x16x17
|
|
|
|
// Check if we need to perform the broadcast on batch dim
|
|
// Not needed if max rank < 3, or if maxrank == 3 and dim[0] matches
|
|
auto needsBatchDimBroadcast = [&]() -> bool {
|
|
if (maxInputRank < 3) {
|
|
return false;
|
|
} else {
|
|
if (maxInputRank == 3 &&
|
|
lhsBroadcastedShape[0] == rhsBroadcastedShape[0]) {
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
};
|
|
|
|
auto performBatchDimBroadcast = needsBatchDimBroadcast();
|
|
|
|
// Inputs to the tosa.matmul
|
|
Value matmulLhs, matmulRhs;
|
|
|
|
using TensorShape_t = struct {
|
|
int64_t dim;
|
|
int64_t shape;
|
|
};
|
|
|
|
// Transpose needs to done if transposeDims are not non-monotonically
|
|
// increasing. E.g. [0, 1, 2, 3]: No transpose [1, 0, 2, 3]: Transpose dim0
|
|
// and dim1 The order need not be sequential, since one or more dims may
|
|
// have been removed due to broadcasting.
|
|
auto isTransposeRequired = [](SmallVector<int32_t> transposedDims) -> bool {
|
|
int32_t lastDim = -1;
|
|
for (auto &dim : transposedDims) {
|
|
if (lastDim > dim)
|
|
return true;
|
|
lastDim = dim;
|
|
}
|
|
return false;
|
|
};
|
|
|
|
SmallVector<TensorShape_t> commonElems, lhsSqueezedElems, rhsSqueezedElems;
|
|
|
|
if (!performBatchDimBroadcast) {
|
|
// Simple with no broadcasting artifacts. Just reshape up to 3D
|
|
matmulLhs = reshapeUpTo3DTensor(rankBroadcastedLhs);
|
|
matmulRhs = reshapeUpTo3DTensor(rankBroadcastedRhs);
|
|
|
|
} else {
|
|
// In this case, either or both input matrices involve broadcasting on
|
|
// their batch dimensions. For example:
|
|
// 4x5x6, 1x6x7 -> 4x5x7
|
|
// 4x1x5x6, 1x3x6x7 -> 4x3x5x7
|
|
// Though maxInputRank is necessarily >=3 here, individual matrices may be
|
|
// lower rank.
|
|
// E.g. 3x4x5x6, 6 -> 3x4x5
|
|
|
|
// These are the accumulated products of the shape of each dim:
|
|
// 1. common dimensions: upper dimensions (dims other than two rightmost)
|
|
// whose shapes are the same for both LHS and RHS.
|
|
// 2. LHS squeezed dimensions: all dimensions of LHS that involve
|
|
// broadcasting in either direction, plus the LHS[-2] shape
|
|
// 3. RHS squeezed dimensions: all dimensions of RHS that involve
|
|
// broadcasting in either direction, plus the RHS[-1] shape
|
|
int64_t commonValue = 1, lhsSqueezedValue = 1, rhsSqueezedValue = 1;
|
|
|
|
// For both LHS and RHS, the dimensions are separated into the common,
|
|
// squeezed and remaining dim. E.g. given
|
|
// LHS = 3x4x5x6
|
|
// RHS = 1x4x6x7
|
|
// common = {{dim=1, shape=4}}
|
|
// lhs squeezed = {{dim=0, shape=3},
|
|
// {dim=2, shape=5}}
|
|
// rhs squeezed = {{dim=0, shape=1},
|
|
// {dim=2, shape=7}}
|
|
// The matmul dim is LHS[-1] and RHS[-2], i.e. 6.
|
|
// Once this is obtained, LHS and RHS are expressed as:
|
|
// LHS = {common, lhs_squeezed, matmul_dim}
|
|
// RHS = {common, matmul_dim, rhs_squeezed}
|
|
// The matmul is then performed to obtain output:
|
|
// matmul_out = {common, lhs_squeezed, rhs_squeezed}
|
|
// Finally, we reshape to 'unsqueeze' the LHS and RHS parts and transpose
|
|
// them back to their correct positions.
|
|
|
|
SmallVector<int64_t> transposedLhsShape;
|
|
SmallVector<int32_t> transposedLhsDims;
|
|
|
|
// Step: generate the common dim/shape information
|
|
for (uint32_t dim = 0; dim < maxInputRank - 2; dim++) {
|
|
bool isDynamicDim =
|
|
lhsBroadcastedTy.isDynamic(lhsBroadcastedShape[dim]);
|
|
if (isDynamicDim ||
|
|
lhsBroadcastedShape[dim] == rhsBroadcastedShape[dim]) {
|
|
commonValue *= lhsBroadcastedShape[dim];
|
|
commonElems.push_back({dim, lhsBroadcastedShape[dim]});
|
|
}
|
|
}
|
|
|
|
// Step: generate the LHS squeezed dim/shape information.
|
|
bool hasDynamicDims = false;
|
|
for (uint32_t dim = 0; dim < maxInputRank - 2; dim++) {
|
|
bool isDynamicDim =
|
|
lhsBroadcastedTy.isDynamic(lhsBroadcastedShape[dim]);
|
|
hasDynamicDims |= isDynamicDim;
|
|
if (!isDynamicDim &&
|
|
lhsBroadcastedShape[dim] != rhsBroadcastedShape[dim]) {
|
|
lhsSqueezedValue *= lhsBroadcastedShape[dim];
|
|
lhsSqueezedElems.push_back({dim, lhsBroadcastedShape[dim]});
|
|
}
|
|
}
|
|
// including LHS[-2]
|
|
lhsSqueezedElems.push_back(
|
|
{maxInputRank - 2, lhsBroadcastedShape[maxInputRank - 2]});
|
|
lhsSqueezedValue *= lhsBroadcastedShape[maxInputRank - 2];
|
|
|
|
// Step: Create the tosa.transpose array. If this array has a
|
|
// non-monotonic series of dims, perform transpose.
|
|
// First the common_elems
|
|
for (uint32_t i = 0; i < commonElems.size(); i++) {
|
|
transposedLhsShape.push_back(commonElems[i].shape);
|
|
transposedLhsDims.push_back(commonElems[i].dim);
|
|
}
|
|
// then the lhs_squeezed elems
|
|
for (uint32_t i = 0; i < lhsSqueezedElems.size(); i++) {
|
|
transposedLhsShape.push_back(lhsSqueezedElems[i].shape);
|
|
transposedLhsDims.push_back(lhsSqueezedElems[i].dim);
|
|
}
|
|
// then the final dim
|
|
transposedLhsDims.push_back(maxInputRank - 1);
|
|
transposedLhsShape.push_back(lhsBroadcastedShape[maxInputRank - 1]);
|
|
|
|
bool lhsNeedsTranspose = isTransposeRequired(transposedLhsDims);
|
|
|
|
auto lhsReshapeInput = rankBroadcastedLhs;
|
|
|
|
if (lhsNeedsTranspose) {
|
|
auto transposedLhsType =
|
|
RankedTensorType::get(transposedLhsShape, rhsElemTy);
|
|
|
|
llvm::Optional<Value> transposedLhsDimsConst =
|
|
tosa::getConstTensor<int32_t>(
|
|
rewriter, op,
|
|
/*vec=*/transposedLhsDims,
|
|
/*shape=*/{static_cast<int32_t>(transposedLhsDims.size())});
|
|
|
|
lhsReshapeInput =
|
|
rewriter
|
|
.create<tosa::TransposeOp>(
|
|
op->getLoc(),
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()
|
|
->convertType(transposedLhsType),
|
|
rankBroadcastedLhs, transposedLhsDimsConst.getValue())
|
|
.getResult();
|
|
}
|
|
|
|
// LHS = {common, lhs_squeezed, matmul_dim}
|
|
SmallVector<int64_t> newLhsShape(
|
|
{1, 1, lhsBroadcastedShape[maxInputRank - 1]});
|
|
newLhsShape[0] = commonValue;
|
|
newLhsShape[1] =
|
|
hasDynamicDims ? ShapedType::kDynamicSize : lhsSqueezedValue;
|
|
|
|
auto newLhsType = RankedTensorType::get(newLhsShape, lhsElemTy);
|
|
|
|
matmulLhs = rewriter.create<tosa::ReshapeOp>(
|
|
op->getLoc(),
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
newLhsType),
|
|
lhsReshapeInput, rewriter.getI64ArrayAttr(newLhsShape));
|
|
|
|
SmallVector<int64_t> transposedRhsShape;
|
|
SmallVector<int32_t> transposedRhsDims;
|
|
|
|
// Step: Create the RHS transpose sequence
|
|
// RHS = {common, matmul_dim, rhs_squeezed}
|
|
// first the common_dims
|
|
for (uint32_t i = 0; i < commonElems.size(); i++) {
|
|
transposedRhsShape.push_back(commonElems[i].shape);
|
|
transposedRhsDims.push_back(commonElems[i].dim);
|
|
}
|
|
// The matmul_dim of RHS
|
|
transposedRhsDims.push_back(maxInputRank - 2);
|
|
transposedRhsShape.push_back(rhsBroadcastedShape[maxInputRank - 2]);
|
|
// finally all the rhs_squeeze dims
|
|
hasDynamicDims = false;
|
|
for (uint32_t dim = 0; dim < maxInputRank - 2; dim++) {
|
|
bool isDynamicDim =
|
|
rhsBroadcastedTy.isDynamic(rhsBroadcastedShape[dim]);
|
|
hasDynamicDims |= isDynamicDim;
|
|
if (!isDynamicDim &&
|
|
rhsBroadcastedShape[dim] != lhsBroadcastedShape[dim]) {
|
|
rhsSqueezedElems.push_back({dim, rhsBroadcastedShape[dim]});
|
|
rhsSqueezedValue *= rhsBroadcastedShape[dim];
|
|
}
|
|
}
|
|
rhsSqueezedElems.push_back(
|
|
{maxInputRank - 1, rhsBroadcastedShape[maxInputRank - 1]});
|
|
rhsSqueezedValue *= rhsBroadcastedShape[maxInputRank - 1];
|
|
for (uint32_t i = 0; i < rhsSqueezedElems.size(); i++) {
|
|
transposedRhsShape.push_back(rhsSqueezedElems[i].shape);
|
|
transposedRhsDims.push_back(rhsSqueezedElems[i].dim);
|
|
}
|
|
|
|
auto transposedRhsType =
|
|
RankedTensorType::get(transposedRhsShape, rhsElemTy);
|
|
|
|
if (hasDynamicDims)
|
|
rhsSqueezedValue = ShapedType::kDynamicSize;
|
|
|
|
SmallVector<int64_t> newRhsShape({commonValue,
|
|
rhsBroadcastedShape[maxInputRank - 2],
|
|
rhsSqueezedValue});
|
|
auto newRhsType = RankedTensorType::get(newRhsShape, rhsElemTy);
|
|
|
|
bool rhsNeedsTranspose = isTransposeRequired(transposedRhsDims);
|
|
|
|
auto transposedRhsValue = rankBroadcastedRhs;
|
|
|
|
if (rhsNeedsTranspose) {
|
|
llvm::Optional<Value> transposedRhsDimsConst =
|
|
tosa::getConstTensor<int32_t>(
|
|
rewriter, op,
|
|
/*vec=*/transposedRhsDims,
|
|
/*shape=*/{static_cast<int32_t>(transposedRhsDims.size())});
|
|
|
|
transposedRhsValue =
|
|
rewriter
|
|
.create<tosa::TransposeOp>(
|
|
op->getLoc(),
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()
|
|
->convertType(transposedRhsType),
|
|
rankBroadcastedRhs, transposedRhsDimsConst.getValue())
|
|
.getResult();
|
|
}
|
|
|
|
// reshape
|
|
matmulRhs = rewriter.create<tosa::ReshapeOp>(
|
|
op->getLoc(),
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
newRhsType),
|
|
transposedRhsValue, rewriter.getI64ArrayAttr(newRhsShape));
|
|
}
|
|
|
|
auto matmulLhsShape =
|
|
matmulLhs.getType().template cast<RankedTensorType>().getShape();
|
|
auto matmulRhsShape =
|
|
matmulRhs.getType().template cast<RankedTensorType>().getShape();
|
|
|
|
// The reshape/transpose should ensure the tosa.matmul always has same
|
|
// batch size for either matrix. If if shapes are dynamic, they'll be
|
|
// appropriately handled.
|
|
assert(matmulLhsShape[0] == matmulRhsShape[0] &&
|
|
"tosa.matmul needs same batchsize on LHS and RHS");
|
|
|
|
SmallVector<int64_t> matmulOutputShape(
|
|
{matmulLhsShape[0], matmulLhsShape[1], matmulRhsShape[2]});
|
|
Type outputElemTy;
|
|
if (lhsElemTy.isa<mlir::FloatType>()) {
|
|
outputElemTy = lhsElemTy;
|
|
} else { // qint8 emits i32 matmul output
|
|
outputElemTy = rewriter.getIntegerType(32);
|
|
}
|
|
|
|
auto mmOutputTy = RankedTensorType::get(matmulOutputShape, outputElemTy);
|
|
auto mmOpResult =
|
|
rewriter
|
|
.create<tosa::MatMulOp>(
|
|
op->getLoc(),
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
mmOutputTy),
|
|
matmulLhs, matmulRhs)
|
|
.getResult();
|
|
|
|
// Perform the reshape to output shape. This is always required unless both
|
|
// inputs are rank=3, in which case the tosa.matmul output itself is
|
|
// correctly shaped.
|
|
bool performOpReshape = !(lhsRank == 3 && rhsRank == 3);
|
|
|
|
auto outputTy = OpConversionPattern<AtenOpT>::getTypeConverter()
|
|
->convertType(op.getType())
|
|
.template cast<RankedTensorType>();
|
|
|
|
if (performOpReshape) {
|
|
// Since the output shape may be unknown, we construct it
|
|
// independently and reshape. Otherwise reshape may be expressed for
|
|
// an unknown to-be-inferred output shape. The final tensor.cast
|
|
// reshapes the known shape to the desired output shape.
|
|
auto computeOpShape = [&](SmallVector<int64_t> &reshapedOpShape,
|
|
SmallVector<int32_t> &transposedOpDims,
|
|
SmallVector<int64_t> &transposedOpShapes) {
|
|
if (maxInputRank == 1)
|
|
return;
|
|
|
|
if (maxInputRank == 2) {
|
|
if (lhsRank == 2)
|
|
reshapedOpShape.push_back(lhsShape[0]);
|
|
if (rhsRank == 2)
|
|
reshapedOpShape.push_back(rhsShape[1]);
|
|
return;
|
|
}
|
|
|
|
// Step: Construct the output transpose/reshape information
|
|
// First the common_dims
|
|
for (uint32_t i = 0; i < commonElems.size(); i++) {
|
|
reshapedOpShape.push_back(commonElems[i].shape);
|
|
transposedOpDims.push_back(commonElems[i].dim);
|
|
}
|
|
|
|
// Then the LHS squeezed dims
|
|
for (uint32_t i = 0; i < lhsSqueezedElems.size() - 1; i++) {
|
|
// Only dims that don't broadcast - broadcasting ones come from the
|
|
// other input.
|
|
if (lhsSqueezedElems[i].shape != 1) {
|
|
reshapedOpShape.push_back(lhsSqueezedElems[i].shape);
|
|
transposedOpDims.push_back(lhsSqueezedElems[i].dim);
|
|
}
|
|
}
|
|
// The last squeezed dim is lhs[-2] which needs to be
|
|
// checked separately for broadcasting
|
|
if (lhsRank > 1) {
|
|
reshapedOpShape.push_back(lhsBroadcastedShape[maxInputRank - 2]);
|
|
transposedOpDims.push_back(maxInputRank - 2);
|
|
}
|
|
|
|
// then the RHS squeezed dims except rhs[-1] which is handled like
|
|
// lhs[-2]
|
|
for (uint32_t i = 0; i < rhsSqueezedElems.size() - 1; i++) {
|
|
if (rhsSqueezedElems[i].shape != 1) {
|
|
reshapedOpShape.push_back(rhsSqueezedElems[i].shape);
|
|
transposedOpDims.push_back(rhsSqueezedElems[i].dim);
|
|
}
|
|
}
|
|
// rhs[-1]
|
|
if (rhsRank > 1) {
|
|
reshapedOpShape.push_back(rhsBroadcastedShape[maxInputRank - 1]);
|
|
transposedOpDims.push_back(maxInputRank - 1);
|
|
}
|
|
|
|
// Final transposed output shape construction
|
|
for (uint32_t i = 0; i < maxInputRank - 2; i++) {
|
|
if (lhsBroadcastedTy.isDynamicDim(i)) {
|
|
transposedOpShapes.push_back(ShapedType::kDynamicSize);
|
|
} else {
|
|
if (lhsBroadcastedShape[i] == rhsBroadcastedShape[i]) {
|
|
transposedOpShapes.push_back(lhsBroadcastedShape[i]);
|
|
} else {
|
|
transposedOpShapes.push_back(lhsBroadcastedShape[i] == 1
|
|
? rhsBroadcastedShape[i]
|
|
: lhsBroadcastedShape[i]);
|
|
}
|
|
}
|
|
}
|
|
if (lhsRank > 1)
|
|
transposedOpShapes.push_back(lhsBroadcastedShape[maxInputRank - 2]);
|
|
if (rhsRank > 1)
|
|
transposedOpShapes.push_back(rhsBroadcastedShape[maxInputRank - 1]);
|
|
|
|
return;
|
|
};
|
|
|
|
SmallVector<int64_t> reshapedOpShape, transposedOpShape;
|
|
SmallVector<int32_t> transposedOpDims;
|
|
|
|
computeOpShape(reshapedOpShape, transposedOpDims, transposedOpShape);
|
|
|
|
bool opNeedsTranspose = isTransposeRequired(transposedOpDims);
|
|
|
|
// Perform reshape
|
|
auto reshapedOpType =
|
|
RankedTensorType::get(reshapedOpShape, outputElemTy);
|
|
auto reshapedOp = rewriter.create<tosa::ReshapeOp>(
|
|
op->getLoc(),
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
reshapedOpType),
|
|
mmOpResult, rewriter.getI64ArrayAttr(reshapedOpShape));
|
|
|
|
if (opNeedsTranspose) {
|
|
|
|
llvm::Optional<Value> transposedOpShapeConst =
|
|
tosa::getConstTensor<int32_t>(
|
|
rewriter, op,
|
|
/*vec=*/transposedOpDims,
|
|
/*shape=*/{static_cast<int32_t>(transposedOpDims.size())});
|
|
|
|
auto transposedOpType =
|
|
RankedTensorType::get(transposedOpShape, outputElemTy);
|
|
auto transposedOp = rewriter.create<tosa::TransposeOp>(
|
|
op->getLoc(),
|
|
OpConversionPattern<AtenOpT>::getTypeConverter()->convertType(
|
|
transposedOpType),
|
|
reshapedOp.getResult(), transposedOpShapeConst.getValue());
|
|
|
|
rewriter.replaceOpWithNewOp<tensor::CastOp>(op, outputTy, transposedOp);
|
|
} else {
|
|
rewriter.replaceOpWithNewOp<tensor::CastOp>(op, outputTy, reshapedOp);
|
|
}
|
|
} else {
|
|
rewriter.replaceOpWithNewOp<tensor::CastOp>(op, outputTy, mmOpResult);
|
|
}
|
|
|
|
return success();
|
|
}
|
|
};
|
|
|
|
template <>
|
|
LogicalResult ConvertAtenOp<AtenRsubScalarOp>::matchAndRewrite(
|
|
AtenRsubScalarOp op, OpAdaptor adaptor,
|
|
ConversionPatternRewriter &rewriter) const {
|
|
|
|
auto self = adaptor.self();
|
|
auto otherScalar = op.other();
|
|
auto alphaScalar = op.alpha();
|
|
|
|
auto selfTy = self.getType().template cast<RankedTensorType>();
|
|
if (!selfTy)
|
|
return op.emitError("Only ranked tensor types supported in TOSA Rsub");
|
|
|
|
if (!selfTy.getElementType().isa<mlir::FloatType>())
|
|
return op.emitError("Only floating-point datatype legalization supported");
|
|
|
|
Value otherTensor, alphaTensor;
|
|
|
|
if (failed(torchScalarToTosaTensor(rewriter, op.getOperation(), otherScalar,
|
|
otherTensor, selfTy.getElementType())))
|
|
return op.emitError("Currently only scalar constants are supported for "
|
|
"conversion in TOSA Rsub operation");
|
|
|
|
if (failed(torchAlphaToTosaTensor(rewriter, op.getOperation(), alphaScalar,
|
|
alphaTensor, selfTy.getElementType(),
|
|
true)))
|
|
return failure();
|
|
|
|
auto multTensor = rewriter.create<tosa::MulOp>(
|
|
op->getLoc(), getTypeConverter()->convertType(op.getType()), self,
|
|
alphaTensor, /*shift*/ 0);
|
|
|
|
rewriter.replaceOpWithNewOp<tosa::SubOp>(
|
|
op, getTypeConverter()->convertType(op.getType()), otherTensor,
|
|
multTensor);
|
|
|
|
return success();
|
|
}
|
|
|
|
} // namespace
|
|
|
|
// -----------------------------------------------------------------------------
|
|
// TorchToTosa Pass
|
|
// -----------------------------------------------------------------------------
|
|
|
|
namespace {
|
|
class ConvertTorchToTosa : public ConvertTorchToTosaBase<ConvertTorchToTosa> {
|
|
public:
|
|
void getDependentDialects(DialectRegistry ®istry) const override {
|
|
registry.insert<tosa::TosaDialect>();
|
|
registry.insert<tensor::TensorDialect>();
|
|
registry.insert<arith::ArithmeticDialect>();
|
|
TorchConversion::getBackendTypeConversionDependentDialects(registry);
|
|
}
|
|
|
|
void runOnOperation() override {
|
|
MLIRContext *context = &getContext();
|
|
ConversionTarget target(*context);
|
|
target.addLegalDialect<tosa::TosaDialect, tensor::TensorDialect,
|
|
arith::ArithmeticDialect>();
|
|
|
|
TypeConverter typeConverter;
|
|
typeConverter.addConversion([](Type type) { return type; });
|
|
TorchConversion::setupBackendTypeConversion(target, typeConverter);
|
|
|
|
RewritePatternSet patterns(context);
|
|
|
|
#define INSERT_UNARY_FPONLY_PATTERN(AtenOp, TosaOp) \
|
|
target.addIllegalOp<AtenOp>(); \
|
|
patterns.add<ConvertAtenUnaryFPOnlyOp<AtenOp, TosaOp>>(typeConverter, \
|
|
context);
|
|
INSERT_UNARY_FPONLY_PATTERN(AtenLogOp, tosa::LogOp)
|
|
INSERT_UNARY_FPONLY_PATTERN(AtenExpOp, tosa::ExpOp)
|
|
#undef INSERT_UNARY_FPONLY_PATTERN
|
|
|
|
#define INSERT_UNARY_PATTERN(AtenOp, TosaOp) \
|
|
target.addIllegalOp<AtenOp>(); \
|
|
patterns.add<ConvertAtenUnaryOp<AtenOp, TosaOp>>(typeConverter, context);
|
|
INSERT_UNARY_PATTERN(AtenNegOp, tosa::NegateOp)
|
|
INSERT_UNARY_PATTERN(AtenFloorOp, tosa::FloorOp)
|
|
INSERT_UNARY_PATTERN(AtenRsqrtOp, tosa::RsqrtOp)
|
|
INSERT_UNARY_PATTERN(AtenBitwiseNotOp, tosa::BitwiseNotOp)
|
|
INSERT_UNARY_PATTERN(AtenCeilOp, tosa::CeilOp)
|
|
INSERT_UNARY_PATTERN(AtenReciprocalOp, tosa::ReciprocalOp)
|
|
#undef INSERT_UNARY_PATTERN
|
|
|
|
#define INSERT_BINARY_PATTERN(AtenOp, TosaOp) \
|
|
target.addIllegalOp<AtenOp>(); \
|
|
patterns.add<ConvertAtenBinaryOp<AtenOp, TosaOp>>(typeConverter, context);
|
|
INSERT_BINARY_PATTERN(AtenMaximumOp, tosa::MaximumOp)
|
|
INSERT_BINARY_PATTERN(AtenMinimumOp, tosa::MinimumOp)
|
|
#undef INSERT_BINARY_PATTERN
|
|
|
|
#define INSERT_BINARY_ADDSUB_PATTERN(AtenOp, TosaOp) \
|
|
target.addIllegalOp<AtenOp>(); \
|
|
patterns.add<ConvertAtenAddSubOp<AtenOp, TosaOp>>(typeConverter, context);
|
|
INSERT_BINARY_ADDSUB_PATTERN(AtenAddTensorOp, tosa::AddOp)
|
|
INSERT_BINARY_ADDSUB_PATTERN(AtenAddScalarOp, tosa::AddOp)
|
|
INSERT_BINARY_ADDSUB_PATTERN(AtenSubTensorOp, tosa::SubOp)
|
|
INSERT_BINARY_ADDSUB_PATTERN(AtenSubScalarOp, tosa::SubOp)
|
|
#undef INSERT_BINARY_ADDSUB_PATTERN
|
|
|
|
#define INSERT_BINARY_COMPARE_PATTERN(AtenOp, TosaOp) \
|
|
target.addIllegalOp<AtenOp>(); \
|
|
patterns.add<ConvertAtenCompareOp<AtenOp, TosaOp>>(typeConverter, context);
|
|
INSERT_BINARY_COMPARE_PATTERN(AtenGtTensorOp, tosa::GreaterOp)
|
|
INSERT_BINARY_COMPARE_PATTERN(AtenGtScalarOp, tosa::GreaterOp)
|
|
INSERT_BINARY_COMPARE_PATTERN(AtenLtTensorOp, tosa::GreaterOp)
|
|
INSERT_BINARY_COMPARE_PATTERN(AtenLtScalarOp, tosa::GreaterOp)
|
|
INSERT_BINARY_COMPARE_PATTERN(AtenEqTensorOp, tosa::EqualOp)
|
|
INSERT_BINARY_COMPARE_PATTERN(AtenEqScalarOp, tosa::EqualOp)
|
|
#undef INSERT_BINARY_COMPARE_PATTERN
|
|
|
|
#define INSERT_BINARY_MUL_PATTERN(AtenOp) \
|
|
target.addIllegalOp<AtenOp>(); \
|
|
patterns.add<ConvertAtenMulOp<AtenOp>>(typeConverter, context);
|
|
INSERT_BINARY_MUL_PATTERN(AtenMulTensorOp);
|
|
INSERT_BINARY_MUL_PATTERN(AtenMulScalarOp);
|
|
#undef INSERT_BINARY_MUL_PATTERN
|
|
|
|
#define INSERT_BINARY_DIV_PATTERN(AtenOp) \
|
|
target.addIllegalOp<AtenOp>(); \
|
|
patterns.add<ConvertAtenDivOp<AtenOp>>(typeConverter, context);
|
|
INSERT_BINARY_DIV_PATTERN(AtenDivTensorOp);
|
|
INSERT_BINARY_DIV_PATTERN(AtenDivScalarOp);
|
|
#undef INSERT_BINARY_DIV_PATTERN
|
|
|
|
#define INSERT_NDIMS_REDUCTION_OP_PATTERN(AtenOp, ConversionFunc) \
|
|
target.addIllegalOp<AtenOp>(); \
|
|
patterns.add<ConvertAtenMultipleDimsReductionOp<AtenOp, ConversionFunc>>( \
|
|
typeConverter, context);
|
|
INSERT_NDIMS_REDUCTION_OP_PATTERN(AtenMeanDimOp,
|
|
mlir::tosa::convertReduceMeanOp)
|
|
INSERT_NDIMS_REDUCTION_OP_PATTERN(AtenSumDimIntListOp,
|
|
mlir::tosa::convertReduceSumOp)
|
|
#undef INSERT_NDIMS_REDUCTION_OP_PATTERN
|
|
|
|
#define INSERT_ONEDIM_REDUCTION_OP_PATTERN(AtenOp, ConversionFunc) \
|
|
target.addIllegalOp<AtenOp>(); \
|
|
patterns.add<ConvertAtenOneDimReductionOp<AtenOp, ConversionFunc>>( \
|
|
typeConverter, context);
|
|
INSERT_ONEDIM_REDUCTION_OP_PATTERN(AtenAnyDimOp,
|
|
mlir::tosa::convertReduceAnyOp)
|
|
#undef INSERT_ONEDIM_REDUCTION_OP_PATTERN
|
|
|
|
#define INSERT_ALLDIMS_REDUCTION_OP_PATTERN(AtenOp, ConversionFunc) \
|
|
target.addIllegalOp<AtenOp>(); \
|
|
patterns.add<ConvertAtenAllDimsReductionOp<AtenOp, ConversionFunc>>( \
|
|
typeConverter, context);
|
|
INSERT_ALLDIMS_REDUCTION_OP_PATTERN(AtenAllOp,
|
|
mlir::tosa::convertReduceAllOp)
|
|
INSERT_ALLDIMS_REDUCTION_OP_PATTERN(AtenAnyOp,
|
|
mlir::tosa::convertReduceAnyOp)
|
|
INSERT_ALLDIMS_REDUCTION_OP_PATTERN(AtenSumOp,
|
|
mlir::tosa::convertReduceSumOp)
|
|
#undef INSERT_ALLDIMS_REDUCTION_OP_PATTERN
|
|
|
|
#define INSERT_SQUEEZE_OP_PATTERN(AtenOp, TemplateForm) \
|
|
target.addIllegalOp<AtenOp>(); \
|
|
patterns.add<TemplateForm<AtenOp>>(typeConverter, context);
|
|
INSERT_SQUEEZE_OP_PATTERN(AtenSqueezeOp, ConvertAtenSqueezeAllDimsOp)
|
|
INSERT_SQUEEZE_OP_PATTERN(AtenSqueezeDimOp, ConvertAtenSqueezeOneDimOp)
|
|
#undef INSERT_SQUEEZE_OP_PATTERN
|
|
|
|
#define INSERT_MATMUL_ATENOP_PATTERN(AtenOp) \
|
|
target.addIllegalOp<AtenOp>(); \
|
|
patterns.add<ConvertAtenMatMulOp<AtenOp>>(typeConverter, context);
|
|
INSERT_MATMUL_ATENOP_PATTERN(AtenMatmulOp);
|
|
INSERT_MATMUL_ATENOP_PATTERN(AtenMmOp);
|
|
INSERT_MATMUL_ATENOP_PATTERN(AtenBmmOp);
|
|
#undef INSERT_MATMUL_ATEMOP_PATTERN
|
|
|
|
#define INSERT_ATENOP_PATTERN(AtenOp) \
|
|
target.addIllegalOp<AtenOp>(); \
|
|
patterns.add<ConvertAtenOp<AtenOp>>(typeConverter, context);
|
|
INSERT_ATENOP_PATTERN(AtenTanhOp);
|
|
INSERT_ATENOP_PATTERN(AtenSigmoidOp);
|
|
INSERT_ATENOP_PATTERN(AtenReluOp);
|
|
INSERT_ATENOP_PATTERN(AtenArgmaxOp);
|
|
INSERT_ATENOP_PATTERN(AtenPowTensorScalarOp);
|
|
INSERT_ATENOP_PATTERN(AtenRsubScalarOp);
|
|
#undef INSERT_ATENOP_PATTERN
|
|
|
|
if (failed(applyPartialConversion(getOperation(), target,
|
|
std::move(patterns))))
|
|
return signalPassFailure();
|
|
}
|
|
};
|
|
} // namespace
|
|
|
|
std::unique_ptr<OperationPass<FuncOp>>
|
|
mlir::torch::createConvertTorchToTosaPass() {
|
|
return std::make_unique<ConvertTorchToTosa>();
|
|
}
|