llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[ONNX] LogSoftmax to Torch (#3024) 

This PR adds support for onnx.LogSoftmax both for old versions (<13,
with axis >=0), and new versions (13).
pull/3049/head
zjgarvey 2024-03-22 13:01:39 -05:00 committed by GitHub
parent 50635dd509
commit 6aa481c204
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 154 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

94
lib/Conversion/TorchOnnxToTorch/DefaultDomainGtoP.cpp
View File

@ -195,6 +195,100 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
binder.op, resultType, operand);
return success();
});
patterns.onOp(
"LogSoftmax", 13,
[](OpBinder binder, ConversionPatternRewriter &rewriter) {
Value input;
Torch::ValueTensorType resultType;
if (binder.tensorOperand(input) || binder.tensorResultType(resultType))
return failure();
int64_t axis;
if (binder.s64IntegerAttr(axis, "axis", -1))
return rewriter.notifyMatchFailure(binder.op, "axis bind failure");
Value axisConst = rewriter.create<Torch::ConstantIntOp>(
binder.getLoc(), rewriter.getI64IntegerAttr(axis));
Value none = rewriter.create<Torch::ConstantNoneOp>(binder.getLoc());
rewriter.replaceOpWithNewOp<Torch::AtenLogSoftmaxIntOp>(
binder.op, resultType, input, axisConst, none);
return success();
});
patterns.onOp(
"LogSoftmax", 1,
[](OpBinder binder, ConversionPatternRewriter &rewriter) {
Value input;
Torch::ValueTensorType resultType;
if (binder.tensorOperand(input) || binder.tensorResultType(resultType))
return failure();
int64_t axis;
if (binder.s64IntegerAttr(axis, "axis", 1))
return rewriter.notifyMatchFailure(binder.op, "axis bind failure");
std::optional<unsigned> maybeRank = Torch::getTensorRank(input);
if (!maybeRank)
return rewriter.notifyMatchFailure(binder.op,
"Unsupported: unranked tensor");
int64_t rank = *maybeRank;
// if negative axis is provided, then flip it to a positive axis
if (axis < 0) {
axis = rank + axis;
}
// need input type and sizes to flatten/unflatten later.
auto inputTy = input.getType().cast<Torch::ValueTensorType>();
if (!inputTy || !inputTy.hasSizes())
return rewriter.notifyMatchFailure(
binder.op, "failed to get input type or sizes");
Value axisConst = rewriter.create<Torch::ConstantIntOp>(
binder.getLoc(), rewriter.getI64IntegerAttr(axis));
Value none = rewriter.create<Torch::ConstantNoneOp>(binder.getLoc());
Value cstEnd = rewriter.create<Torch::ConstantIntOp>(
binder.getLoc(), rewriter.getI64IntegerAttr(rank - 1));
// The old version of LogSoftmax flattens post-axis dims, performs
// LogSoftmax on the flattened dim, then unflattens back to the original
// shape.
// this section gets some size information necessary for
// flattening/unflattening
if (!inputTy || !inputTy.hasSizes())
return failure();
llvm::ArrayRef<int64_t> allDims(inputTy.getSizes());
llvm::ArrayRef<int64_t> rightDims(allDims.begin() + axis,
allDims.end());
llvm::SmallVector<int64_t> leftDims(allDims.begin(),
allDims.begin() + axis);
int64_t prodRightSizes = 1;
llvm::SmallVector<Value> rightDimConsts;
for (int64_t n : rightDims) {
rightDimConsts.push_back(rewriter.create<Torch::ConstantIntOp>(
binder.getLoc(), rewriter.getI64IntegerAttr(n)));
if (n == Torch::kUnknownSize) {
prodRightSizes = -1;
break;
}
prodRightSizes *= n;
}
leftDims.push_back(prodRightSizes);
// the following list will be used to unflatten the right side
Value rightDimsPrimList = rewriter.create<Torch::PrimListConstructOp>(
binder.getLoc(),
rewriter.getType<Torch::ListType>(
rewriter.getType<Torch::IntType>()),
rightDimConsts);
auto flatRightTy = rewriter.getType<Torch::ValueTensorType>(
leftDims, inputTy.getOptionalDtype());
// flatten input
Value inputFlatRight = rewriter.create<Torch::AtenFlattenUsingIntsOp>(
binder.getLoc(), flatRightTy, input, axisConst, cstEnd);
// compute lsm over flattened index
Value outputFlatRight = rewriter.create<Torch::AtenLogSoftmaxIntOp>(
binder.getLoc(), flatRightTy, inputFlatRight, axisConst, none);
// unflatten
rewriter.replaceOpWithNewOp<Torch::AtenUnflattenIntOp>(
binder.op, resultType, outputFlatRight, axisConst,
rightDimsPrimList);
return success();
});
patterns.onOp("MatMul", 13,
[](OpBinder binder, ConversionPatternRewriter &rewriter) {
Torch::ValueTensorType resultType;

5
projects/pt1/e2e_testing/xfail_sets.py
View File

@ -1906,11 +1906,6 @@ ONNX_XFAIL_SET = {
"HardswishRandomModule_basic",
"MobilenetV3Module_basic",
# Failure - onnx_lowering: onnx.LogSoftmax
"LogSoftmaxIntModule_basic",
"_LogSoftmaxModuleStable_basic",
"_LogSoftmaxModule_basic",
# Failure - onnx_lowering: onnx.MaxPool
"MaxPool2dWithIndicesAllNegativeValuesModule_basic",
"MaxPool2dWithIndicesNonDefaultPaddingModule_basic",

60
test/Conversion/TorchOnnxToTorch/simple_ops_g_to_p.mlir
View File

@ -748,6 +748,66 @@ func.func @test_mod_int64_no_fmod(%arg0: !torch.vtensor<[6],si64>, %arg1: !torch
// -----
// CHECK-LABEL: func.func @test_log_softmax_default_axis
func.func @test_log_softmax_default_axis(%arg0: !torch.vtensor<[1,3],f32>) -> !torch.vtensor<[1,3],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
// CHECK: %[[CIM1:.*]] = torch.constant.int -1
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[LSM:.*]] = torch.aten.log_softmax.int %arg0, %[[CIM1]], %[[NONE]] : !torch.vtensor<[1,3],f32>, !torch.int, !torch.none -> !torch.vtensor<[1,3],f32>
// CHECK: return %[[LSM]] : !torch.vtensor<[1,3],f32>
%0 = torch.operator "onnx.LogSoftmax"(%arg0) : (!torch.vtensor<[1,3],f32>) -> !torch.vtensor<[1,3],f32>
return %0 : !torch.vtensor<[1,3],f32>
}
// -----
// CHECK-LABEL: func.func @test_log_softmax_axis_2
func.func @test_log_softmax_axis_2(%arg0: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
// CHECK: %[[CI2:.*]] = torch.constant.int 2
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[LSM:.*]] = torch.aten.log_softmax.int %arg0, %[[CI2]], %[[NONE]] : !torch.vtensor<[3,4,5],f32>, !torch.int, !torch.none -> !torch.vtensor<[3,4,5],f32>
// CHECK: return %[[LSM]] : !torch.vtensor<[3,4,5],f32>
%0 = torch.operator "onnx.LogSoftmax"(%arg0) {torch.onnx.axis = 2 : si64} : (!torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],f32>
return %0 : !torch.vtensor<[3,4,5],f32>
}
// -----
// CHECK-LABEL: func.func @test_logsoftmax_old_axis_1_dynamic_dim
func.func @test_logsoftmax_old_axis_1_dynamic_dim(%arg0: !torch.vtensor<[3,4,?],f32>) -> !torch.vtensor<[3,4,?],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 1 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
// CHECK: %[[CI1:.*]] = torch.constant.int 1
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[CI2:.*]] = torch.constant.int 2
// CHECK: %[[CI4:.*]] = torch.constant.int 4
// CHECK: %[[CIM1:.*]] = torch.constant.int -1
// CHECK: %[[LIST:.*]] = torch.prim.ListConstruct %[[CI4]], %[[CIM1]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: %[[FLAT_IN:.*]] = torch.aten.flatten.using_ints %arg0, %[[CI1]], %[[CI2]] : !torch.vtensor<[3,4,?],f32>, !torch.int, !torch.int -> !torch.vtensor<[3,?],f32>
// CHECK: %[[LSM:.*]] = torch.aten.log_softmax.int %[[FLAT_IN]], %[[CI1]], %[[NONE]] : !torch.vtensor<[3,?],f32>, !torch.int, !torch.none -> !torch.vtensor<[3,?],f32>
// CHECK: %[[UNFLAT:.*]] = torch.aten.unflatten.int %[[LSM]], %[[CI1]], %[[LIST]] : !torch.vtensor<[3,?],f32>, !torch.int, !torch.list<int> -> !torch.vtensor<[3,4,?],f32>
// CHECK: return %[[UNFLAT]] : !torch.vtensor<[3,4,?],f32>
%0 = torch.operator "onnx.LogSoftmax"(%arg0) {torch.onnx.axis = 1 : si64} : (!torch.vtensor<[3,4,?],f32>) -> !torch.vtensor<[3,4,?],f32>
return %0 : !torch.vtensor<[3,4,?],f32>
}
// -----
// CHECK-LABEL: func.func @test_logsoftmax_old_axis_1_static
func.func @test_logsoftmax_old_axis_1_static(%arg0: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 1 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
// CHECK: %[[CI1:.*]] = torch.constant.int 1
// CHECK: %[[NONE:.*]] = torch.constant.none
// CHECK: %[[CI2:.*]] = torch.constant.int 2
// CHECK: %[[CI4:.*]] = torch.constant.int 4
// CHECK: %[[CI5:.*]] = torch.constant.int 5
// CHECK: %[[LIST:.*]] = torch.prim.ListConstruct %[[CI4]], %[[CI5]] : (!torch.int, !torch.int) -> !torch.list<int>
// CHECK: %[[FLAT_IN:.*]] = torch.aten.flatten.using_ints %arg0, %[[CI1]], %[[CI2]] : !torch.vtensor<[3,4,5],f32>, !torch.int, !torch.int -> !torch.vtensor<[3,20],f32>
// CHECK: %[[LSM:.*]] = torch.aten.log_softmax.int %[[FLAT_IN]], %[[CI1]], %[[NONE]] : !torch.vtensor<[3,20],f32>, !torch.int, !torch.none -> !torch.vtensor<[3,20],f32>
// CHECK: %[[UNFLAT:.*]] = torch.aten.unflatten.int %[[LSM]], %[[CI1]], %[[LIST]] : !torch.vtensor<[3,20],f32>, !torch.int, !torch.list<int> -> !torch.vtensor<[3,4,5],f32>
// CHECK: return %[[UNFLAT]] : !torch.vtensor<[3,4,5],f32>
%0 = torch.operator "onnx.LogSoftmax"(%arg0) {torch.onnx.axis = 1 : si64} : (!torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],f32>
return %0 : !torch.vtensor<[3,4,5],f32>
}
// -----
// CHECK-LABEL: func.func @test_neg
func.func @test_neg(%arg0: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],f32> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 13 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
// CHECK: torch.aten.neg %arg0 : !torch.vtensor<[3,4,5],f32> -> !torch.vtensor<[3,4,5],f32>