Skip to content

[TorchToLinalg] Support lowering AtenSoftmaxIntOp for linalg backend #4293

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[TorchToLinalg] Support lowering AtenSoftmaxIntOp for linalg backend #4293

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 57 additions & 0 deletions lib/Conversion/TorchToLinalg/Linear.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1759,6 +1759,61 @@ struct ConvertAtenFftRfftOp final : OpConversionPattern<AtenFftRfftOp> {

} // namespace

namespace {
class ConvertAtenSoftmaxIntOp : public OpConversionPattern<AtenSoftmaxIntOp> {
public:
using OpConversionPattern::OpConversionPattern;
LogicalResult
matchAndRewrite(AtenSoftmaxIntOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {

Location loc = op.getLoc();
if (failed(verifyLinalgCompatibleTypes(op, rewriter)))
return failure();

Value input = adaptor.getSelf();

Value result = op.getResult();
Value dimValue = op.getDim();
auto constOp = dimValue.getDefiningOp<Torch::ConstantIntOp>();
if (!constOp) {
return rewriter.notifyMatchFailure(
op, "dimension must be a constant integer");
}

int64_t dimInt = constOp.getValue();

// Handle negative dimensions by converting to positive
if (auto tensorType = cast<RankedTensorType>(input.getType())) {
int64_t rank = tensorType.getRank();
if (dimInt < 0) {
dimInt += rank;
}
if (dimInt < 0 || dimInt >= rank) {
return rewriter.notifyMatchFailure(op, "dimension out of bounds");
}
}

IntegerAttr dimAttr = rewriter.getI64IntegerAttr(dimInt);

Type newResultType =
getTypeConverter()->convertType(op.getResult().getType());
auto resultType = cast<RankedTensorType>(newResultType);
Value result_tensor = rewriter.create<tensor::EmptyOp>(
loc, resultType.getShape(), resultType.getElementType());

auto softmax = rewriter.create<linalg::SoftmaxOp>(
loc, TypeRange{resultType}, input, result_tensor, dimAttr);

rewriter.replaceOp(op, softmax.getResult());
// if we know constop is only used by this softmax op, erase it
if (constOp.getResult().hasOneUse()) {
rewriter.eraseOp(constOp);
}
return success();
}
};
} // namespace
void mlir::torch::torch_to_linalg::populateLinearPatternsAndLegality(
TypeConverter &typeConverter, RewritePatternSet &patterns,
ConversionTarget &target) {
Expand All @@ -1775,4 +1830,6 @@ void mlir::torch::torch_to_linalg::populateLinearPatternsAndLegality(
patterns.add<ConvertAtenConvolutionOp>(typeConverter, context);
target.addIllegalOp<AtenFftRfftOp>();
patterns.add<ConvertAtenFftRfftOp>(typeConverter, context);
target.addIllegalOp<AtenSoftmaxIntOp>();
patterns.add<ConvertAtenSoftmaxIntOp>(typeConverter, context);
}