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

Implement `onnx.Hardmax` lowering (#3342) 

Co-authored-by: Ubuntu <xunli@wsno1.judsoscro3wupi0qm4bjlj5m3b.bx.internal.cloudapp.net>
Co-authored-by: Hasekawa-Takumi <bewater.private476@passmail.net>
pull/3368/head
lialan 2024-05-20 11:26:24 -04:00 committed by GitHub
parent cc28d566ff
commit 99511cef82
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 62 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -1816,4 +1816,56 @@ void mlir::torch::onnx_c::populateDefaultDomainGtoP(
binder.op, resultType, input);
return success();
});
patterns.onOp(
"Hardmax", 13, [](OpBinder binder, ConversionPatternRewriter &rewriter) {
// onnx.Hardmax can be expanded into the following python code:
//
// import torch.nn.functional as F
// def hardmax(tensor, dim=-1):
// maximums = torch.argmax(tensor, dim=dim, keepdim=False)
// return F.one_hot(maximums)
//
// Given an example input:
// tensor([[1, 2, 3],
// [4, 6, 5],
// [9, 8, 7]])
// Above code yields the following:
// tensor([[0, 0, 1],
// [0, 1, 0],
// [1, 0, 0]])
Torch::ValueTensorType resultType;
int64_t axisValue;
Value input, axis;
if (binder.tensorOperand(input) ||
binder.s64IntegerAttr(axisValue, "axis") ||
binder.tensorResultType(resultType))
return failure();
auto loc = binder.getLoc();
std::optional<int64_t> axisIntTorch =
onnxDtypeIntToTorchDtypeInt(axisValue);
if (!axisIntTorch.has_value())
return rewriter.notifyMatchFailure(
binder.op, "unimplemented support for the given axis conversion");
axis = rewriter.create<Torch::ConstantIntOp>(
loc, rewriter.getI64IntegerAttr(axisIntTorch.value()));
// torch.argmax
Value constKeepDims = rewriter.create<Torch::ConstantBoolOp>(
loc, rewriter.getType<Torch::BoolType>(),
rewriter.getBoolAttr(false));
Value argmax = rewriter.create<Torch::AtenArgmaxOp>(
loc, resultType, input, axis, constKeepDims);
// one_hot
Value oneInt = rewriter.create<Torch::ConstantIntOp>(
loc, rewriter.getI64IntegerAttr(1));
rewriter.replaceOpWithNewOp<Torch::AtenOneHotOp>(binder.op, resultType,
argmax, oneInt);
return success();
});
}

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

@ -1025,3 +1025,13 @@ func.func @test_hardswish(%arg0: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<
%0 = torch.operator "onnx.HardSwish"(%arg0) : (!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_hardmax
func.func @test_hardmax(%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 = 14 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
// CHECK: %[[ARGMAX:.+]] = torch.aten.argmax %arg0, %int6, %false : !torch.vtensor<[3,4,5],f32>, !torch.int, !torch.bool -> !torch.vtensor<[3,4,5],f32>
// CHECK: torch.aten.one_hot %[[ARGMAX]], %int1 : !torch.vtensor<[3,4,5],f32>, !torch.int -> !torch.vtensor<[3,4,5],f32>
%0 = torch.operator "onnx.Hardmax"(%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>
}