diff --git a/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp b/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp index f943f288f..c0e9af22c 100644 --- a/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp +++ b/lib/Conversion/TorchOnnxToTorch/DefaultDomainQtoZ.cpp @@ -794,6 +794,148 @@ void mlir::torch::onnx_c::populateDefaultDomainQtoZ( return success(); }); + // split with fixed-size parts + // Arguments: + // - input: the tensor to split + // Attributes: + // - axis: the axis along which to split the input + // - num_outputs: the number of outputs to produce + // Outputs: + // - outputs: the produced outputs. Variadic with num_outputs elements. + // Note: torch.aten gives a list of tensors, but ONNX gives a variadic list of + // tensors + // so we need to unpack the list + patterns.onOp( + "Split", 1, [](OpBinder binder, ConversionPatternRewriter &rewriter) { + Value self; + int64_t axis; + int64_t num_outputs; + if (binder.tensorOperand(self)) + return rewriter.notifyMatchFailure( + binder.op, "Not converting to AtenSplitTensorOp due to input " + "tensor mismatch"); + if (binder.s64IntegerAttr(axis, "axis", 0)) + return rewriter.notifyMatchFailure(binder.op, + "Failed to get axis attribute"); + if (binder.s64IntegerAttr(num_outputs, "num_outputs", 0)) + return rewriter.notifyMatchFailure( + binder.op, "Failed to get num_outputs attribute"); + + auto result0Ty = + binder.op->getResult(0).getType().cast(); + auto selfTy = self.getType().cast(); + + int64_t dim = axis; + if (dim < 0) + dim += selfTy.getSizes().size(); + + // set intermediate shape to the shape of the first result + // if the results are of different shapes + // set the splitted axis to variable shape + llvm::SmallVector intermediateShape(result0Ty.getSizes()); + for (auto result : binder.op->getResultTypes()) { + int64_t d = result.cast().getSizes()[dim]; + intermediateShape[dim] = d == intermediateShape[dim] ? d : -1; + } + + Value dimValue = rewriter.create( + binder.getLoc(), rewriter.getType(), + rewriter.getIntegerAttr(rewriter.getIntegerType(64), dim)); + + Value splitSize = rewriter.create( + binder.getLoc(), rewriter.getType(), + rewriter.getIntegerAttr(rewriter.getIntegerType(64), num_outputs)); + + // TODO: Attempting to use the shape expected by the ONNX mlir as ground + // truth. For now just use dynamic shapes. + auto resultOuterType = + Torch::ListType::get(rewriter.getType( + /*std::optional>=*/intermediateShape, + result0Ty.getOptionalDtype())); + Torch::AtenSplitTensorOp new_op = + rewriter.create( + binder.getLoc(), resultOuterType, self, splitSize, dimValue); + + // the onnx op is variadic with multiple results, but AtenSplitWithSizes + // outputs a list so we need to unpack the list + rewriter.replaceOpWithNewOp( + binder.op, binder.op->getResults().getType(), new_op.getResult()); + + return success(); + }); + + // split with variable parts + // Arguments: + // - input: the tensor to split + // - split: the sizes of the splits to be produced + // Attributes: + // - axis: the axis along which to split the input + // - num_outputs: the number of outputs to produce + // Outputs: + // - outputs: the produced outputs. Variadic with num_outputs elements. + // Note: torch.aten gives a list of tensors, but ONNX gives a variadic list of + // tensors + // so we need to unpack the list + patterns.onOp( + "Split", 1, [](OpBinder binder, ConversionPatternRewriter &rewriter) { + Value self; + Value split; + int64_t axis; + int64_t num_outputs; + if (binder.tensorOperandAtIndex(self, 0) || + binder.tensorOperandAtIndex(split, 1)) + return rewriter.notifyMatchFailure( + binder.op, "Not converting to AtenSplitWithSizesOp due to input " + "tensor mismatch"); + if (binder.s64IntegerAttr(axis, "axis", 0)) + return rewriter.notifyMatchFailure(binder.op, + "Failed to get axis attribute"); + if (binder.s64IntegerAttr(num_outputs, "num_outputs", 0)) + return rewriter.notifyMatchFailure( + binder.op, "Failed to get num_outputs attribute"); + + auto result0Ty = + binder.op->getResult(0).getType().cast(); + auto selfTy = + cast(binder.op->getOperand(0).getType()); + + int64_t dim = axis; + if (dim < 0) + dim += selfTy.getSizes().size(); + + llvm::SmallVector intermediateShape(result0Ty.getSizes()); + for (auto result : binder.op->getResultTypes()) { + int64_t d = result.cast().getSizes()[dim]; + intermediateShape[dim] = d == intermediateShape[dim] ? d : -1; + } + + Torch::PrimTolistOp splitToList = rewriter.create( + binder.getLoc(), + Torch::ListType::get(rewriter.getType()), split); + + Value dimValue = rewriter.create( + binder.getLoc(), rewriter.getType(), + rewriter.getIntegerAttr(rewriter.getIntegerType(64), dim)); + + // TODO: Attempting to use the shape expected by the ONNX mlir as ground + // truth. For now just use dynamic shapes. + auto resultOuterType = + Torch::ListType::get(rewriter.getType( + /*std::optional>=*/intermediateShape, + result0Ty.getOptionalDtype())); + Torch::AtenSplitWithSizesOp new_op = + rewriter.create( + binder.getLoc(), resultOuterType, self, + splitToList.getResult(0), dimValue); + + // the onnx op is variadic with multiple results, but AtenSplitWithSizes + // outputs a list so we need to unpack the list + rewriter.replaceOpWithNewOp( + binder.op, binder.op->getResults().getType(), new_op.getResult()); + + return success(); + }); + patterns.onOp("Tan", 7, [](OpBinder binder, ConversionPatternRewriter &rewriter) { Torch::ValueTensorType resultType; diff --git a/test/Conversion/TorchOnnxToTorch/simple_ops_q_to_z.mlir b/test/Conversion/TorchOnnxToTorch/simple_ops_q_to_z.mlir index 5aca8688d..b2a19334a 100644 --- a/test/Conversion/TorchOnnxToTorch/simple_ops_q_to_z.mlir +++ b/test/Conversion/TorchOnnxToTorch/simple_ops_q_to_z.mlir @@ -795,6 +795,36 @@ func.func @test_sinh_example(%arg0: !torch.vtensor<[3],f32>) -> !torch.vtensor<[ // ----- +// CHECK-LABEL: func.func @test_split_variable_parts_2d_opset18( +// CHECK-SAME: %[[VAL_INPUT:.*]]: !torch.vtensor<[2,6],f32>, +// CHECK-SAME: %[[VAL_SPLIT:.*]]: !torch.vtensor<[2],si64> +// CHECK: %[[VAL_SPLIT_LIST:.*]] = torch.prim.tolist(%[[VAL_SPLIT]]) : !torch.vtensor<[2],si64> -> !torch.list +// CHECK: %[[VAL_AXIS:.*]] = torch.constant.int 1 +// CHECK: %[[VAL_RESULT_LIST:.*]] = torch.aten.split_with_sizes %[[VAL_INPUT]], %[[VAL_SPLIT_LIST]], %[[VAL_AXIS]] : !torch.vtensor<[2,6],f32>, !torch.list, !torch.int -> !torch.list> +// CHECK: %[[VAL_VARIADIC_RETURN_VALUE:.*]]:2 = torch.prim.ListUnpack %[[VAL_RESULT_LIST]] : !torch.list> -> !torch.vtensor<[2,2],f32>, !torch.vtensor<[2,4],f32> +// CHECK: return %[[VAL_VARIADIC_RETURN_VALUE]]#0, %[[VAL_VARIADIC_RETURN_VALUE]]#1 : !torch.vtensor<[2,2],f32>, !torch.vtensor<[2,4],f32> +func.func @test_split_variable_parts_2d_opset18(%arg0: !torch.vtensor<[2,6],f32>, %arg1: !torch.vtensor<[2],si64>) -> (!torch.vtensor<[2,2],f32>, !torch.vtensor<[2,4],f32>) attributes {torch.onnx_meta.ir_version = 8 : si64, torch.onnx_meta.opset_version = 18 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} { + %0:2 = torch.operator "onnx.Split"(%arg0, %arg1) {torch.onnx.axis = 1 : si64} : (!torch.vtensor<[2,6],f32>, !torch.vtensor<[2],si64>) -> (!torch.vtensor<[2,2],f32>, !torch.vtensor<[2,4],f32>) + return %0#0, %0#1 : !torch.vtensor<[2,2],f32>, !torch.vtensor<[2,4],f32> +} + +// ----- + +// CHECK-LABEL: func.func @test_split_2d_uneven_split_opset18( +// CHECK-SAME: %[[INPUT_TENSOR:.*]]: !torch.vtensor<[2,8],f32>) -> (!torch.vtensor<[2,3],f32>, !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,2],f32>) attributes {torch.onnx_meta.ir_version = 8 : si64, torch.onnx_meta.opset_version = 18 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} { +// CHECK: %[[AXIS:.*]] = torch.constant.int 1 +// CHECK: %[[SPLIT_SIZE:.*]] = torch.constant.int 3 +// CHECK: %[[SPLIT_RESULT:.*]] = torch.aten.split.Tensor %[[INPUT_TENSOR]], %[[SPLIT_SIZE]], %[[AXIS]] : !torch.vtensor<[2,8],f32>, !torch.int, !torch.int -> !torch.list> +// CHECK: %[[UNPACKED_TENSORS:.*]]:3 = torch.prim.ListUnpack %[[SPLIT_RESULT]] : !torch.list> -> !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,2],f32> +// CHECK: return %[[UNPACKED_TENSORS]]#0, %[[UNPACKED_TENSORS]]#1, %[[UNPACKED_TENSORS]]#2 : !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,2],f32> +// CHECK: } +func.func @test_split_2d_uneven_split_opset18(%arg0: !torch.vtensor<[2,8],f32>) -> (!torch.vtensor<[2,3],f32>, !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,2],f32>) attributes {torch.onnx_meta.ir_version = 8 : si64, torch.onnx_meta.opset_version = 18 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} { + %0:3 = torch.operator "onnx.Split"(%arg0) {torch.onnx.axis = 1 : si64, torch.onnx.num_outputs = 3 : si64} : (!torch.vtensor<[2,8],f32>) -> (!torch.vtensor<[2,3],f32>, !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,2],f32>) + return %0#0, %0#1, %0#2 : !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,3],f32>, !torch.vtensor<[2,2],f32> +} + +// ----- + // CHECK-LABEL: func.func @test_tan func.func @test_tan(%arg0: !torch.vtensor<[3,4,5],f32>) -> !torch.vtensor<[3,4,5],f32> attributes {torch.onnx_meta.ir_version = 3 : si64, torch.onnx_meta.opset_version = 7 : si64, torch.onnx_meta.producer_name = "backend-test", torch.onnx_meta.producer_version = ""} { // CHECK: %[[TAN:.+]] = torch.aten.tan %arg0