diff --git a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py index 56ad2512e0..0b4848c537 100644 --- a/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py +++ b/py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py @@ -1427,6 +1427,7 @@ def aten_ops_cumsum( name, args[0], args[1], + kwargs.get("dtype"), ) diff --git a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py index 90153e24b6..2a41367491 100644 --- a/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py +++ b/py/torch_tensorrt/dynamo/conversion/impl/slice/ops.py @@ -4,13 +4,16 @@ import numpy as np import tensorrt as trt +import torch from tensorrt import ITensor as TRTTensor from torch.fx.node import Target +from torch_tensorrt import _enums from torch_tensorrt.dynamo._SourceIR import SourceIR from torch_tensorrt.dynamo.conversion import impl from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext from torch_tensorrt.dynamo.conversion.converter_utils import ( calculate_strides, + cast_trt_tensor, flatten_dims, get_positive_dim, get_trt_tensor, @@ -366,7 +369,21 @@ def cumsum( name: str, input: TRTTensor, dim: int, + dtype: Optional[torch.dtype] = None, ) -> TRTTensor: + # aten.cumsum accumulates bool and integer inputs in int64 and floats in + # their own dtype; an explicit dtype wins over both + input_dtype = _enums.dtype._from(input.dtype).to(torch.dtype) + if dtype is not None: + acc_dtype = dtype + elif not input_dtype.is_floating_point: + acc_dtype = torch.int64 + else: + acc_dtype = input_dtype + + if input_dtype != acc_dtype: + input = cast_trt_tensor(ctx, input, acc_dtype, f"{name}_input_cast") + input_shape = input.shape dim = get_positive_dim(dim, len(input_shape)) if input_shape[dim] < 0: @@ -399,13 +416,13 @@ def cumsum( ) else: data_shape.append(input_shape[i]) - zero_trttensor = impl.full.full( - ctx, target, source_ir, name + "_full", data_shape, 0.0 - ) else: - new_dims = tuple(data.shape) - zeros = np.zeros(new_dims, dtype=np.float32) - zero_trttensor = get_trt_tensor(ctx, zeros, f"{name}_initial_value") + data_shape = list(data.shape) + + # full rather than np.zeros: numpy has no bf16 + zero_trttensor = impl.full.full( + ctx, target, source_ir, f"{name}_initial_value", data_shape, 0, dtype=acc_dtype + ) running_sum = loop.add_recurrence(zero_trttensor) set_layer_name(running_sum, target, f"{name}_running_sum", source_ir) diff --git a/tests/py/dynamo/conversion/test_cumsum_aten.py b/tests/py/dynamo/conversion/test_cumsum_aten.py index 1a83c9c45f..4cee79febd 100644 --- a/tests/py/dynamo/conversion/test_cumsum_aten.py +++ b/tests/py/dynamo/conversion/test_cumsum_aten.py @@ -98,6 +98,80 @@ def forward(self, x): immutable_weights=False, ) + @parameterized.expand( + [ + (torch.int32, torch.int32), # explicit dtype keeps int32 + (torch.int32, None), # integral promotes to int64 + (torch.int64, None), + (torch.float16, None), + (torch.bfloat16, None), + (torch.float32, torch.float16), + (torch.float32, torch.bfloat16), + ] + ) + def test_cumsum_dtype(self, input_dtype, out_dtype): + class Cumsum(nn.Module): + def forward(self, x): + if out_dtype is None: + return torch.ops.aten.cumsum.default(x, 0) + return torch.ops.aten.cumsum.default(x, 0, dtype=out_dtype) + + # 1,2,3,4 accumulate exactly in every dtype under test + inputs = [torch.tensor([1, 2, 3, 4], dtype=input_dtype)] + self.run_test( + Cumsum(), + inputs, + use_dynamo_tracer=True, + immutable_weights=False, + ) + + @parameterized.expand( + [ + (torch.int32, None), + (torch.int64, None), + (torch.float32, torch.int64), + ] + ) + def test_cumsum_accumulator_is_exact(self, input_dtype, out_dtype): + class Cumsum(nn.Module): + def forward(self, x): + if out_dtype is None: + return torch.ops.aten.cumsum.default(x, 0) + return torch.ops.aten.cumsum.default(x, 0, dtype=out_dtype) + + # 2**24+1 is unrepresentable in float32, so a float accumulator stalls + # here while an integral one keeps counting; the sums must be exact + inputs = [torch.tensor([2**24, 1, 1, 1], dtype=input_dtype)] + self.run_test( + Cumsum(), + inputs, + rtol=0, + atol=0, + use_dynamo_tracer=True, + immutable_weights=False, + ) + + @parameterized.expand([(torch.float16,), (torch.bfloat16,)]) + def test_cumsum_dynamic_shape_dtype(self, input_dtype): + class Cumsum(nn.Module): + def forward(self, x): + return torch.ops.aten.cumsum.default(x, 0) + + # a dynamic non-cumsum dim sends the seed down full's shape-tensor path + inputs = [ + torch_tensorrt.Input( + min_shape=(1, 2), + opt_shape=(2, 3), + max_shape=(3, 4), + dtype=input_dtype, + ), + ] + self.run_test_with_dynamic_shape( + Cumsum(), + inputs, + immutable_weights=False, + ) + if __name__ == "__main__": run_tests()