Fix cumsum dtype

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"),
     )
 
 

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)

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()