TensorRT icon indicating copy to clipboard operation
TensorRT copied to clipboard

scatter_add_decomposition

Open apbose opened this issue 1 year ago • 1 comments

The issue in this PR is for cases where there is index collision. Example of such cases- scatter_add(input_tensor = torch.zeros(3,5), dim=1, index_tensor = torch.tensor([[0, 1, 2, 0]]).cuda(), src_tensor = torch.tensor([[1, 2, 3, 1]], dtype=torch.int32).cuda()) Looks like straighforward decomposition by stacking the src_tensor would not work, since it would again lead to collision for some cases. The better way to implement is to do converter implementation and allot the values using advanced indexing. Working on this now.

apbose avatar Apr 09 '24 21:04 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?

apbose avatar May 16 '24 22:05 apbose