diff --git a/lib/Conversion/TorchToLinalg/Linear.cpp b/lib/Conversion/TorchToLinalg/Linear.cpp index 6c04dd12f..44ac95ce0 100644 --- a/lib/Conversion/TorchToLinalg/Linear.cpp +++ b/lib/Conversion/TorchToLinalg/Linear.cpp @@ -882,14 +882,14 @@ public: if (bias.getType().isa()) { Value c0; if (resultDTy.isa()) { - c0 = rewriter.create( - loc, FloatAttr::get(resultDTy, 0.0)); + c0 = rewriter.create(loc, + FloatAttr::get(resultDTy, 0.0)); } else if (resultDTy.isa()) { - c0 = rewriter.create( - loc, IntegerAttr::get(resultDTy, 0)); + c0 = rewriter.create(loc, + IntegerAttr::get(resultDTy, 0)); } - outputTensor = rewriter.create(loc, c0, initTensor) - .getResult(0); + outputTensor = + rewriter.create(loc, c0, initTensor).getResult(0); } else { auto biasType = bias.getType().cast(); @@ -1058,11 +1058,11 @@ public: loc, collapsedType, weight, collapsedDims); conv = rewriter - .create( - loc, outputTensor.getType(), - ValueRange{paddedInput, collapsedWeight}, outputTensor, - stridesAttr, dilationAttr) - .getResult(0); + .create( + loc, outputTensor.getType(), + ValueRange{paddedInput, collapsedWeight}, outputTensor, + stridesAttr, dilationAttr) + .getResult(0); Type newResultType = getTypeConverter()->convertType(op.getType()); rewriter.replaceOpWithNewOp(op, newResultType, conv); diff --git a/lib/Conversion/TorchToLinalg/Uncategorized.cpp b/lib/Conversion/TorchToLinalg/Uncategorized.cpp index 657ea460f..d28369cc5 100644 --- a/lib/Conversion/TorchToLinalg/Uncategorized.cpp +++ b/lib/Conversion/TorchToLinalg/Uncategorized.cpp @@ -1274,6 +1274,29 @@ static Value createLinalgPayloadCalculationForElementwiseOp( return result; } + if (auto fmod = dyn_cast(op)) { + Type newResultType = converter->convertType(fmod.getType()) + .cast() + .getElementType(); + + Value self = convertScalarToDtype(b, loc, payloadArgs[0], newResultType); + Value other = convertScalarToDtype(b, loc, payloadArgs[1], newResultType); + Value result; + + if (newResultType.isa()) { + Value n = b.create(loc, self, other); + n = b.create(loc, n); + Value n_y = b.create(loc, n, other); + result = b.create(loc, self, n_y); + } else if (newResultType.isa()) { + Value n = b.create(loc, self, other); + Value n_y = b.create(loc, n, other); + result = b.create(loc, self, n_y); + } else { + fmod.emitError("Unsupported type encountered for AtenFmodTensorOp."); + } + return result; + } if (auto reciprocal = dyn_cast(op)) { Type dtype = converter->convertType(reciprocal.getType()) .cast() @@ -1541,22 +1564,22 @@ public: AtenMulScalarOp, AtenLogOp, AtenErfOp, AtenSqrtOp, AtenFloorOp, AtenPowScalarOp, AtenPowTensorScalarOp, AtenPowTensorTensorOp, AtenLog2Op, AtenLog10Op, AtenLog1pOp, AtenRsqrtOp, AtenDivScalarOp, - AtenRemainderScalarOp, AtenRemainderTensorOp, AtenAbsOp, - AtenReciprocalOp, AtenBitwiseAndTensorOp, AtenBitwiseAndScalarOp, - AtenBitwiseOrTensorOp, AtenBitwiseXorTensorOp, - AtenBitwiseLeftShiftTensorOp, AtenBitwiseRightShiftTensorOp, - AtenGtScalarOp, AtenGeScalarOp, AtenEqScalarOp, AtenLtScalarOp, - AtenLeScalarOp, AtenWhereSelfOp, AtenCeilOp, AtenGtTensorOp, - AtenGeTensorOp, AtenEqTensorOp, AtenNeTensorOp, AtenLtTensorOp, - AtenLeTensorOp, AtenSubScalarOp, AtenAddScalarOp, AtenThresholdOp, - AtenThresholdBackwardOp, AtenHardtanhBackwardOp, AtenCloneOp, - AtenSinOp, AtenCosOp, AtenNeScalarOp, AtenNegOp, - AtenMaskedFillTensorOp, AtenLogicalOrOp, AtenLogicalAndOp, - AtenLogicalXorOp, AtenLogicalNotOp, AtenIsinfOp, AtenTriuOp, - AtenTrilOp, AtenBitwiseNotOp, AtenRoundOp, AtenFillScalarOp, - AtenFillTensorOp, AtenAtanOp, AtenAcosOp, AtenAtanhOp, AtenAcoshOp, - AtenAsinOp, AtenAsinhOp, AtenRealOp, AtenImagOp, - AtenDequantizeSelfOp, AtenDequantizeTensorOp, + AtenRemainderScalarOp, AtenRemainderTensorOp, AtenFmodTensorOp, + AtenAbsOp, AtenReciprocalOp, AtenBitwiseAndTensorOp, + AtenBitwiseAndScalarOp, AtenBitwiseOrTensorOp, + AtenBitwiseXorTensorOp, AtenBitwiseLeftShiftTensorOp, + AtenBitwiseRightShiftTensorOp, AtenGtScalarOp, AtenGeScalarOp, + AtenEqScalarOp, AtenLtScalarOp, AtenLeScalarOp, AtenWhereSelfOp, + AtenCeilOp, AtenGtTensorOp, AtenGeTensorOp, AtenEqTensorOp, + AtenNeTensorOp, AtenLtTensorOp, AtenLeTensorOp, AtenSubScalarOp, + AtenAddScalarOp, AtenThresholdOp, AtenThresholdBackwardOp, + AtenHardtanhBackwardOp, AtenCloneOp, AtenSinOp, AtenCosOp, + AtenNeScalarOp, AtenNegOp, AtenMaskedFillTensorOp, AtenLogicalOrOp, + AtenLogicalAndOp, AtenLogicalXorOp, AtenLogicalNotOp, AtenIsinfOp, + AtenTriuOp, AtenTrilOp, AtenBitwiseNotOp, AtenRoundOp, + AtenFillScalarOp, AtenFillTensorOp, AtenAtanOp, AtenAcosOp, + AtenAtanhOp, AtenAcoshOp, AtenAsinOp, AtenAsinhOp, AtenRealOp, + AtenImagOp, AtenDequantizeSelfOp, AtenDequantizeTensorOp, AtenQuantizePerTensorOp>(op)) return rewriter.notifyMatchFailure(op, "not a supported elementwise op"); @@ -2584,9 +2607,10 @@ void mlir::torch::torch_to_linalg::populateUncategorizedPatternsAndLegality( AtenCloneOp, AtenSinOp, AtenCosOp, AtenNeScalarOp, AtenMaskedFillTensorOp, AtenLogicalOrOp, AtenLogicalAndOp, AtenAtanOp, AtenAcosOp, AtenLogicalXorOp, AtenLogicalNotOp, AtenIsinfOp, AtenTriuOp, AtenTrilOp, - AtenRemainderScalarOp, AtenRemainderTensorOp, AtenBitwiseNotOp, - AtenRoundOp, AtenFillScalarOp, AtenFillTensorOp, AtenRealOp, AtenImagOp, - AtenDequantizeSelfOp, AtenDequantizeTensorOp, AtenQuantizePerTensorOp>(); + AtenRemainderScalarOp, AtenFmodTensorOp, AtenRemainderTensorOp, + AtenBitwiseNotOp, AtenRoundOp, AtenFillScalarOp, AtenFillTensorOp, + AtenRealOp, AtenImagOp, AtenDequantizeSelfOp, AtenDequantizeTensorOp, + AtenQuantizePerTensorOp>(); patterns.add(typeConverter, context); target.addIllegalOp(); patterns.add(typeConverter, context); diff --git a/lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp b/lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp index b40ca2cde..f9ad383d1 100644 --- a/lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp +++ b/lib/Dialect/Torch/Transforms/AbstractInterpLibrary.cpp @@ -6865,6 +6865,10 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() { " %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list) -> !torch.list\n" " return %0 : !torch.list\n" " }\n" +" func.func @\"__torch_mlir_shape_fn.aten.fmod.Tensor\"(%arg0: !torch.list, %arg1: !torch.list) -> !torch.list {\n" +" %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list) -> !torch.list\n" +" return %0 : !torch.list\n" +" }\n" " func.func @\"__torch_mlir_shape_fn.aten.floor_divide.Scalar\"(%arg0: !torch.list, %arg1: !torch.float) -> !torch.list {\n" " %0 = call @__torch__.torch.jit._shape_functions.unary(%arg0) : (!torch.list) -> !torch.list\n" " return %0 : !torch.list\n" @@ -11395,6 +11399,14 @@ StringRef mlir::torch::Torch::getAbstractInterpLibrary() { " %4 = call @__torch__.torch_mlir.jit_ir_importer.build_tools.library_generator.promote_dtypes(%1, %3) : (!torch.list>, !torch.list) -> !torch.int\n" " return %4 : !torch.int\n" " }\n" +" func.func @\"__torch_mlir_dtype_fn.aten.fmod.Tensor\"(%arg0: !torch.tuple, %arg1: !torch.tuple) -> !torch.int {\n" +" %0:2 = torch.prim.TupleUnpack %arg1 : !torch.tuple -> !torch.int, !torch.int\n" +" %1:2 = torch.prim.TupleUnpack %arg0 : !torch.tuple -> !torch.int, !torch.int\n" +" %2 = torch.prim.ListConstruct %1#0, %0#0 : (!torch.int, !torch.int) -> !torch.list>\n" +" %3 = torch.prim.ListConstruct %1#1, %0#1 : (!torch.int, !torch.int) -> !torch.list\n" +" %4 = call @__torch__.torch_mlir.jit_ir_importer.build_tools.library_generator.promote_dtypes(%2, %3) : (!torch.list>, !torch.list) -> !torch.int\n" +" return %4 : !torch.int\n" +" }\n" " func.func @\"__torch_mlir_dtype_fn.aten.floor_divide.Scalar\"(%arg0: !torch.tuple, %arg1: !torch.number) -> !torch.int {\n" " %none = torch.constant.none\n" " %str = torch.constant.str \"AssertionError: \"\n" diff --git a/projects/pt1/e2e_testing/xfail_sets.py b/projects/pt1/e2e_testing/xfail_sets.py index adb406f60..e2b198839 100644 --- a/projects/pt1/e2e_testing/xfail_sets.py +++ b/projects/pt1/e2e_testing/xfail_sets.py @@ -1641,6 +1641,9 @@ ONNX_XFAIL_SET = { "ElementwiseOrTensorStaticShapeModule_basic", "ElementwiseQuantizePerTensorModule_basic", "ElementwiseRemainderTensorModule_Int_basic", + "ElementwiseFmodTensor_Float_basic", + "ElementwiseFmodTensor_Int_Float_basic", + "ElementwiseFmodTensor_Int_basic", "EmptyStridedModule_basic", "EmptyStridedSizeIntStrideModule_basic", "EqIntModule_basic", diff --git a/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py b/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py index 99f4f2200..a9bf5640d 100644 --- a/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py +++ b/projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py @@ -438,6 +438,9 @@ def aten〇remainder〇Scalar〡shape(self: List[int], other: float) -> List[int def aten〇remainder〇Tensor〡shape(self: List[int], other: List[int]) -> List[int]: return upstream_shape_functions.unary(self) +def aten〇fmod〇Tensor〡shape(self: List[int], other: List[int]) -> List[int]: + return upstream_shape_functions.unary(self) + def aten〇floor_divide〇Scalar〡shape(self: List[int], other: float) -> List[int]: return upstream_shape_functions.unary(self) @@ -3491,6 +3494,14 @@ def aten〇fmod〇Scalar〡dtype(self_rank_dtype: Tuple[int, int], other: Union[ dtypes = [self_dtype, get_dtype_of_scalar(other)] return promote_dtypes(ranks, dtypes) +@check_dtype_function(_check_two_tensor_op()) +def aten〇fmod〇Tensor〡dtype(self_rank_dtype: Tuple[int, int], other_rank_dtype: Tuple[int, int]) -> int: + other_rank, other_dtype = other_rank_dtype + self_rank, self_dtype = self_rank_dtype + ranks: List[Optional[int]] = [self_rank, other_rank] + dtypes = [self_dtype, other_dtype] + return promote_dtypes(ranks, dtypes) + @check_dtype_function( _check_tensors_with_the_same_dtype(num_of_tensors=1, error_types={torch.complex64, torch.complex128}, other=1) + _check_tensors_with_the_same_dtype(num_of_tensors=1, error_types={torch.complex64, torch.complex128}, other=1.0)) diff --git a/projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise.py b/projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise.py index 24bbe2919..d9921d23d 100644 --- a/projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise.py +++ b/projects/pt1/python/torch_mlir_e2e_test/test_suite/elementwise.py @@ -2526,6 +2526,68 @@ def ElementwiseRemainderScalarModule_Bool_basic(module, tu: TestUtils): # ============================================================================== + +class ElementwiseFmodTensor_Float(torch.nn.Module): + + def __init__(self): + super().__init__() + + @export + @annotate_args([ + None, + ([-1], torch.float32, True), + ([-1], torch.float32, True) + ]) + def forward(self, x, y): + return torch.fmod(x, y) + + +@register_test_case(module_factory=lambda: ElementwiseFmodTensor_Float()) +def ElementwiseFmodTensor_Float_basic(module, tu: TestUtils): + module.forward(tu.rand(100, low=-10, high=10), tu.rand(100, low=-10, high=10)) + +# ============================================================================== + +class ElementwiseFmodTensor_Int_Float(torch.nn.Module): + + def __init__(self): + super().__init__() + + @export + @annotate_args([ + None, + ([-1], torch.int32, True), + ([-1], torch.float32, True) + ]) + def forward(self, x, y): + return torch.fmod(x, y) + + +@register_test_case(module_factory=lambda: ElementwiseFmodTensor_Int_Float()) +def ElementwiseFmodTensor_Int_Float_basic(module, tu: TestUtils): + module.forward(tu.randint(100, low=-10, high=10).to(torch.int32), tu.rand(100, low=-10, high=10)) + +# ============================================================================== + +class ElementwiseFmodTensor_Int(torch.nn.Module): + + def __init__(self): + super().__init__() + + @export + @annotate_args([ + None, + ([-1], torch.int32, True), + ([-1], torch.int32, True), + ]) + def forward(self, x, y): + return torch.fmod(x, y) + + +@register_test_case(module_factory=lambda: ElementwiseFmodTensor_Int()) +def ElementwiseFmodTensor_Int_basic(module, tu: TestUtils): + module.forward(tu.randint(100, low=0, high=1000).to(torch.int32), tu.randint(100, low=1, high=1000).to(torch.int32)) + # ============================================================================== class ElementwiseRemainderTensorModule_Int_Float(torch.nn.Module):