scatter_add_decomposition

[apbose]

No description provided.

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/_decompositions.py	2024-05-14 21:04:24.249027+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/_decompositions.py	2024-05-14 21:06:14.833958+00:00
@@ -181,13 +181,16 @@
    input_tensor: torch.Tensor,
    src_tensor: torch.Tensor,
    dim: int,
    index: torch.Tensor,
) -> torch.Tensor:
-    input_tensor_to_add = torch.scatter(torch.empty_like(input_tensor), dim, index, src_tensor)
+    input_tensor_to_add = torch.scatter(
+        torch.empty_like(input_tensor), dim, index, src_tensor
+    )
    scatter_add_tensor = torch.add(input_tensor, input_tensor_to_add.cuda())
    return scatter_add_tensor
+

def get_decompositions(
    enable_experimental_decompositions: bool = False,
) -> Dict[OpOverload, Callable[[Any], Any]]:
    if enable_experimental_decompositions:

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/_decompositions.py	2024-05-14 21:04:33.408462+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/_decompositions.py	2024-05-14 21:06:24.734368+00:00
@@ -181,13 +181,16 @@
    input_tensor: torch.Tensor,
    src_tensor: torch.Tensor,
    dim: int,
    index: torch.Tensor,
) -> torch.Tensor:
-    input_tensor_to_add = torch.scatter(torch.empty_like(input_tensor), dim, index, src_tensor)
+    input_tensor_to_add = torch.scatter(
+        torch.empty_like(input_tensor), dim, index, src_tensor
+    )
    scatter_add_tensor = torch.add(input_tensor, input_tensor_to_add.cuda())
    return scatter_add_tensor
+

def get_decompositions(
    enable_experimental_decompositions: bool = False,
) -> Dict[OpOverload, Callable[[Any], Any]]:
    if enable_experimental_decompositions:

[github-actions]

There are some changes that do not conform to Python style guidelines:

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/_decompositions.py	2024-05-14 21:01:28.751182+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/lowering/_decompositions.py	2024-05-14 21:09:51.626381+00:00
@@ -181,13 +181,16 @@
    input_tensor: torch.Tensor,
    src_tensor: torch.Tensor,
    dim: int,
    index: torch.Tensor,
) -> torch.Tensor:
-    input_tensor_to_add = torch.scatter(torch.empty_like(input_tensor), dim, index, src_tensor)
+    input_tensor_to_add = torch.scatter(
+        torch.empty_like(input_tensor), dim, index, src_tensor
+    )
    scatter_add_tensor = torch.add(input_tensor, input_tensor_to_add.cuda())
    return scatter_add_tensor
+

def get_decompositions(
    enable_experimental_decompositions: bool = False,
) -> Dict[OpOverload, Callable[[Any], Any]]:
    if enable_experimental_decompositions:

[gs-olive]

Could torch.empty_like(input_tensor) instead be just input_tensor, or is it required to be empty?

[apbose]

Yes it would be required. The logic is like we would apply scatter with the indices to an empty tensor and then add it to the input tensor.

[gs-olive]

Without .cuda() does the decomposition fail?

[apbose]

Yes in this case it complains that the two tensors are on different devices.

[apbose]

@gs-olive though the test cases pass, I don't think that the decomposition is taking place.
It fails in the backend at aot_export_joint_simple. Shows me something like

CRITICAL:torch_tensorrt.dynamo.backend.backends:Halting compilation on build failure since pass_through_build_failures was specified as True. To
return the default Torch implementation and avoid halting compilation on engine build failures, specify pass_through_build_failures=False.
W0516 15:51:10.818862 139669023303488 torch/_dynamo/exc.py:201] [0/0] Backend compiler failed with a fake tensor exception at
W0516 15:51:10.818862 139669023303488 torch/_dynamo/exc.py:201] [0/0]   File "<eval_with_key>.5 from /home/abose/Documents/work/torchTRT_TRT10_scatter_5_2/TensorRT/tests/py/dynamo/lowering/test_decompositions.py:523 in forward", line 9, in forward
W0516 15:51:10.818862 139669023303488 torch/_dynamo/exc.py:201] [0/0]     return scatter_add_default
W0516 15:51:10.818862 139669023303488 torch/_dynamo/exc.py:201] [0/0] Adding a graph break.

Is this due to torch.unique?