scatter reduce decomposition #3008

apbose · 2024-07-15T21:10:34Z

#2740 should be using this. Will change it once this PR is finalized

github-actions

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

--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_decompositions.py	2024-07-15 21:13:39.692683+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_decompositions.py	2024-07-15 21:15:36.206907+00:00
@@ -1163,11 +1163,13 @@
            (
                "scatter_reduce_amax_zero_dim_indexOne_constant",
                0,
                torch.tensor([[0, 1, 2, 0]]).cuda(),
                torch.tensor([[1, 2, 3, 4]], dtype=torch.int32).cuda(),
-                {torch.ops.aten.amax.default,},
+                {
+                    torch.ops.aten.amax.default,
+                },
                torch.zeros(3, 5, dtype=torch.int32).cuda(),
                "amax",
            ),
            (
                "scatter_reduce_amax_zero_dim_indexTwo_constant",

peri044 · 2024-07-24T21:19:54Z

tests/py/dynamo/lowering/test_decompositions.py

+            inputs,
+            expected_ops=expected_ops,
+            unexpected_ops=unexpected_ops,
+            min_block_size=2,


Why is the min_block size = 2 here ? What does lower_graph_testing do ? Does it use our partitioning ?

Yes lower_graph_testing uses our partition. It is used for returning the expected ops unseen and the seen unexpected ops. The default is 3, but sometimes the graph in the test case is too small, and the block size is lesser than 3 and it errors out. Thats why I set it to 1 or 2, since that is not something which we are testing explicitly in the test.
Let me try with the default 3. If it passes, I will remove the min_block_size =2 then.

peri044 · 2024-07-24T21:21:48Z

py/torch_tensorrt/dynamo/lowering/_decompositions.py

@@ -243,6 +244,99 @@ def empty_strided_decomposition(*args, **kwargs) -> torch.Tensor:
    )


+# enum class for reduce operation of scatter_reduce
+class reduceOperation(Enum):


minor - consider renaming it to ReduceOperation

github-actions

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

--- /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_decompositions.py	2024-07-30 02:56:59.084675+00:00
+++ /home/runner/work/TensorRT/TensorRT/tests/py/dynamo/lowering/test_decompositions.py	2024-07-30 02:58:55.960957+00:00
@@ -1020,11 +1020,10 @@
            0,
            DECIMALS_OF_AGREEMENT,
            f"Scatter_add TRT outputs don't match with the original model.",
        )

-
    @parameterized.expand(
        [
            ############################sum###########################
            (
                "scatter_reduce_add_zero_dim_indexOne_constant",

py/torch_tensorrt/dynamo/lowering/_decompositions.py

peri044 · 2024-08-15T23:54:48Z

py/torch_tensorrt/dynamo/lowering/_decompositions.py

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


let's use the device where the input_tensor exists

peri044 · 2024-09-03T14:49:12Z

@apbose CI is failing on scatter tests

peri044 · 2024-09-05T18:37:50Z

py/torch_tensorrt/dynamo/lowering/_decompositions.py

+            print("Invalid Operation for Reduce op!!")
+
+        operation_rhs = torch.scatter(scatter_tensor, dim, index_tensor, src_tensor)
+        device = to_torch_device(default_device())


use the device of initial_tensor here instead of default

github-actions

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

--- /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/utils.py	2024-09-10 16:33:48.731288+00:00
+++ /home/runner/work/TensorRT/TensorRT/py/torch_tensorrt/dynamo/utils.py	2024-09-10 16:34:09.607993+00:00
@@ -186,11 +186,11 @@
    """
    device = None
    for parameter in list(module.parameters()):
        if isinstance(parameter, (torch.nn.parameter.Parameter, torch.Tensor)):
            return parameter.device
-    
+
    for buffer in list(module.buffers()):
        if isinstance(buffer, (torch.Tensor)):
            return buffer.device

    if device is None:

HolyWu · 2024-09-10T17:04:03Z

py/torch_tensorrt/dynamo/lowering/_decompositions.py

+    index: torch.Tensor,
+    src_tensor: torch.Tensor,
+    reduce: str,
+) -> torch.Tensor:


There is a kwarg include_self in https://github.com/pytorch/pytorch/blob/bc1b8f094d24de27432f4c29f0729e85a6b5ba63/aten/src/ATen/native/native_functions.yaml#L8237. Is it intentionally not handled in our decomposition?

Thanks for the review! Most of the cases which I have seen is with include_self = True. Here we have the implementation with the default case. No particular reason, I could add cases with include_self = False

Add include_self=True in the function arguments. And raise an error saying we don't support the case when user sets it False

peri044 · 2024-09-10T18:10:31Z

py/torch_tensorrt/dynamo/utils.py

+            return parameter.device
+
+    for buffer in list(module.buffers()):


The buffer device overrides the parameter device here which shouldn't be the case. Check device of parameters first, if not found, use buffers.
Also consider adding break once the device is found.

peri044 · 2024-09-10T18:14:01Z

py/torch_tensorrt/dynamo/lowering/_decompositions.py

+    index: torch.Tensor,
+    src_tensor: torch.Tensor,
+    reduce: str,
+) -> torch.Tensor:


Add include_self=True in the function arguments. And raise an error saying we don't support the case when user sets it False

apbose requested a review from peri044 July 15, 2024 21:10

facebook-github-bot added the cla signed label Jul 15, 2024

github-actions bot added component: tests Issues re: Tests component: lowering Issues re: The lowering / preprocessing passes component: api [Python] Issues re: Python API component: dynamo Issues relating to the `torch.compile` or `torch._dynamo.export` paths labels Jul 15, 2024

apbose force-pushed the scatter_reduce_decomposition branch from ef97199 to 8e5151f Compare July 15, 2024 21:13

apbose marked this pull request as draft July 15, 2024 21:14

github-actions bot requested changes Jul 15, 2024

View reviewed changes

apbose force-pushed the scatter_reduce_decomposition branch 3 times, most recently from 5dbc520 to 6c77c44 Compare July 17, 2024 17:11

github-actions bot added component: conversion Issues re: Conversion stage component: converters Issues re: Specific op converters labels Jul 17, 2024

apbose force-pushed the scatter_reduce_decomposition branch from 6c77c44 to 33e76dc Compare July 17, 2024 17:13

apbose marked this pull request as ready for review July 17, 2024 17:14

apbose force-pushed the scatter_reduce_decomposition branch from 33e76dc to d137714 Compare July 17, 2024 17:23

peri044 reviewed Jul 24, 2024

View reviewed changes

apbose force-pushed the scatter_reduce_decomposition branch 3 times, most recently from 9aea7dd to f0ccb92 Compare July 30, 2024 02:58

github-actions bot requested changes Jul 30, 2024

View reviewed changes

apbose force-pushed the scatter_reduce_decomposition branch 2 times, most recently from 4bf82d5 to b6aa19d Compare August 6, 2024 00:00

peri044 reviewed Aug 15, 2024

View reviewed changes

apbose force-pushed the scatter_reduce_decomposition branch 2 times, most recently from 64442d6 to 2ce4933 Compare August 22, 2024 23:49

apbose requested a review from peri044 August 27, 2024 15:44

apbose force-pushed the scatter_reduce_decomposition branch 2 times, most recently from b689e76 to 020d32c Compare August 30, 2024 07:54

peri044 reviewed Sep 5, 2024

View reviewed changes

apbose added 3 commits September 9, 2024 12:25

scatter reduce decomposition

f1be9fe

addressing review comments-move all tensors to input tensor device

cd7d682

removing full_like decomposition op after PR-3077

485adf9

apbose force-pushed the scatter_reduce_decomposition branch 3 times, most recently from 648fa95 to f124297 Compare September 10, 2024 16:33

github-actions bot requested changes Sep 10, 2024

View reviewed changes

changing the device setting in conversion.py

35f2b00

apbose force-pushed the scatter_reduce_decomposition branch from f124297 to 35f2b00 Compare September 10, 2024 16:36

HolyWu reviewed Sep 10, 2024

View reviewed changes

peri044 reviewed Sep 10, 2024

View reviewed changes

Assertion error for include_self=False

e0eda18

apbose merged commit 501a1e1 into main Sep 11, 2024
13 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

scatter reduce decomposition #3008

scatter reduce decomposition #3008

apbose commented Jul 15, 2024 •

edited

Loading

github-actions bot left a comment

peri044 Jul 24, 2024

apbose Jul 24, 2024 •

edited

Loading

peri044 Jul 24, 2024

github-actions bot left a comment

peri044 Aug 15, 2024

peri044 commented Sep 3, 2024

peri044 Sep 5, 2024

github-actions bot left a comment

HolyWu Sep 10, 2024

apbose Sep 10, 2024 •

edited

Loading

peri044 Sep 10, 2024

peri044 Sep 10, 2024

peri044 Sep 10, 2024

peri044 Sep 10, 2024

		return parameter.device

		for buffer in list(module.buffers()):

scatter reduce decomposition #3008

scatter reduce decomposition #3008

Conversation

apbose commented Jul 15, 2024 • edited Loading

github-actions bot left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apbose Jul 24, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

peri044 commented Sep 3, 2024

Choose a reason for hiding this comment

github-actions bot left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apbose Sep 10, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apbose commented Jul 15, 2024 •

edited

Loading

apbose Jul 24, 2024 •

edited

Loading

apbose Sep 10, 2024 •

edited

Loading