Skip to content

Commit

Permalink
Add support for the onnx.SequenceConstruct op. (llvm#3316)
Browse files Browse the repository at this point in the history
  • Loading branch information
AWoloszyn authored and Branko Trifkovic committed May 24, 2024
1 parent 015feff commit fe0017e
Show file tree
Hide file tree
Showing 3 changed files with 50 additions and 0 deletions.
12 changes: 12 additions & 0 deletions include/torch-mlir/Conversion/TorchOnnxToTorch/Patterns.h
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,18 @@ struct OpBinder {
return success();
}

ParseResult tensorListResultType(Torch::ListType &type0) {
if (op->getNumResults() != 1)
return failure();
auto tt = dyn_cast<Torch::ListType>(op->getResult(0).getType());
if (!tt)
return failure();
if (!toValidTensorType(tt.getContainedType()))
return failure();
type0 = tt;
return success();
}

ParseResult tensorResultTypes(llvm::SmallVector<mlir::Type> &typeList) {
for (auto result : op->getResults()) {
auto t = toValidTensorType(result.getType());
Expand Down
14 changes: 14 additions & 0 deletions lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -518,6 +518,20 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ(
cstStrReduction);
return success();
});
patterns.onOp(
"SequenceConstruct", 11,
[](OpBinder binder, ConversionPatternRewriter &rewriter) {
SmallVector<Value> operands;
Torch::ListType resultType;

if (binder.tensorOperands(operands, binder.getNumOperands()) ||
binder.tensorListResultType(resultType))
return failure();

rewriter.replaceOpWithNewOp<Torch::PrimListConstructOp>(
binder.op, resultType, operands);
return success();
});
patterns.onOp(
"Sigmoid", 1, [](OpBinder binder, ConversionPatternRewriter &rewriter) {
Torch::ValueTensorType resultType;
Expand Down
24 changes: 24 additions & 0 deletions test/Conversion/TorchOnnxToTorch/simple_ops_q_to_z.mlir
Original file line number Diff line number Diff line change
Expand Up @@ -2075,6 +2075,30 @@ func.func @test_random_uniform_like(%arg0: !torch.vtensor<[10],f32>) -> !torch.v

// -----

// CHECK-LABEL: func.func @test_sequence_construct_3
module {
func.func @test_sequence_construct_3(%arg0: !torch.vtensor<[2,3,4],f32>, %arg1: !torch.vtensor<[2,3,4],f32>, %arg2: !torch.vtensor<[2,3,4],f32>) -> !torch.list<vtensor<[2,3,4],f32>> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 12 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
// CHECK: %[[SEQ:.+]] = torch.prim.ListConstruct %arg0, %arg1, %arg2 : (!torch.vtensor<[2,3,4],f32>, !torch.vtensor<[2,3,4],f32>, !torch.vtensor<[2,3,4],f32>) -> !torch.list<vtensor<[2,3,4],f32>>
// CHECK: return %[[SEQ]] : !torch.list<vtensor<[2,3,4],f32>>
%0 = torch.operator "onnx.SequenceConstruct"(%arg0, %arg1, %arg2) : (!torch.vtensor<[2,3,4],f32>, !torch.vtensor<[2,3,4],f32>, !torch.vtensor<[2,3,4],f32>) -> !torch.list<vtensor<[2,3,4],f32>>
return %0 : !torch.list<vtensor<[2,3,4],f32>>
}
}

// -----

// CHECK-LABEL: func.func @test_sequence_construct_1
module {
func.func @test_sequence_construct_1(%arg0: !torch.vtensor<[2,3,4],f32>) -> !torch.list<vtensor<[2,3,4],f32>> attributes {torch.onnx_meta.ir_version = 7 : si64, torch.onnx_meta.opset_version = 12 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} {
// CHECK: %[[SEQ:.+]] = torch.prim.ListConstruct %arg0 : (!torch.vtensor<[2,3,4],f32>) -> !torch.list<vtensor<[2,3,4],f32>>
// CHECK: return %[[SEQ]] : !torch.list<vtensor<[2,3,4],f32>>
%0 = torch.operator "onnx.SequenceConstruct"(%arg0) : (!torch.vtensor<[2,3,4],f32>) -> !torch.list<vtensor<[2,3,4],f32>>
return %0 : !torch.list<vtensor<[2,3,4],f32>>
}
}

// -----

// CHECK-LABEL: func.func @test_sce_mean_3d
func.func @test_sce_mean_3d(%arg0: !torch.vtensor<[3,5,2],f32>, %arg1: !torch.vtensor<[3,2],si64>) -> !torch.vtensor<[],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: %[[NONE:.+]] = torch.constant.none
Expand Down

0 comments on commit fe0017e

Please sign in to comment.