scatter reduce decomposition

Author: apbose

py/torch_tensorrt/dynamo/lowering/_decompositions.py

@@ -1,4 +1,5 @@
 import logging
+from enum import Enum, auto
 from typing import Any, Callable, Dict, List, Optional
 
 import torch
@@ -243,6 +244,84 @@ def empty_strided_decomposition(*args, **kwargs) -> torch.Tensor:
     )
 
 
+# enum class for reduce operation of scatter_reduce
+class reduceOperation(Enum):
+    SUM = ("Sum reduce operation", lambda x, y: torch.add(x, y))
+    PROD = ("Product reduce operation", lambda x, y: torch.mul(x, y))
+    MEAN = ("Mean reduce operation", lambda x, y: torch.div(torch.add(x, y), 2))
+    AMAX = ("Amax reduce operation", lambda x, y: torch.amax(x, y))
+    AMIN = ("Amin reduce operation", lambda x, y: torch.amin(x, y))
+
+    def __new__(cls, description, func):
+        obj = object.__new__(cls)
+        obj._value_ = auto()  # Assign a unique value based on the number of members
+        obj.description = description
+        obj.func = func
+        return obj
+
+    def reduce_operation_with_scatter(
+        self, operation_lhs, initial_tensor, dim, index_tensor, src_tensor
+    ):
+        scatter_tensor = None
+        if self == reduceOperation.SUM or self == reduceOperation.MEAN:
+            scatter_tensor = torch.zeros_like(initial_tensor)
+        elif self == reduceOperation.PROD:
+            scatter_tensor = torch.ones_like(initial_tensor)
+        elif self.name == reduceOperation.AMIN or self.name == reduceOperation.AMAX:
+            scatter_tensor = initial_tensor
+        else:
+            # This case would not be encountered from torch itself
+            print("Invalid Operation for Reduce op!!")
+
+        operation_rhs = torch.scatter(scatter_tensor, dim, index_tensor, src_tensor)
+        device = to_torch_device(default_device())
+        operation_lhs = operation_lhs.to(device)
+        operation_rhs = operation_rhs.to(device)
+        return self.func(operation_lhs, operation_rhs)
+
+
+@register_torch_trt_decomposition(
+    torch.ops.aten.scatter_reduce.two, registry=TORCH_TRT_DECOMPOSITIONS
+)
+def scatter_reduce_decomposition(
+    input_tensor: torch.Tensor,
+    dim: int,
+    index: torch.Tensor,
+    src_tensor: torch.Tensor,
+    reduce: str,
+) -> torch.Tensor:
+    scatter_loop_tensor = input_tensor
+    src_shape = list(src_tensor.shape)
+    src_dim = src_shape[dim]
+
+    for i in range(0, src_dim):
+        src_slice = torch.select(src_tensor, dim, i)
+        index_slice = torch.select(index, dim, i)
+        # unsqueeze src and index in dim
+        src_slice = torch.unsqueeze(src_slice, dim)
+        index_slice = torch.unsqueeze(index_slice, dim)
+        device = to_torch_device(default_device())
+
+        # moving tensor to default device
+        scatter_loop_tensor = scatter_loop_tensor.to(device)
+        index_slice = index_slice.to(device)
+        src_slice = src_slice.to(device)
+        if reduce == "sum":
+            reduceOp = reduceOperation.SUM
+        elif reduce == "prod":
+            reduceOp = reduceOperation.PROD
+        elif reduce == "mean":
+            reduceOp = reduceOperation.MEAN
+        elif reduce == "amax":
+            reduceOp = reduceOperation.AMAX
+        elif reduce == "amin":
+            reduceOp = reduceOperation.AMIN
+        scatter_loop_tensor = reduceOp.reduce_operation_with_scatter(
+            scatter_loop_tensor, input_tensor, dim, index_slice, src_slice
+        )
+    return scatter_loop_tensor
+
+
 def get_decompositions(
     enable_experimental_decompositions: bool = False,
 ) -> Dict[OpOverload, Callable[[Any], Any]]: