Grid test changes

Author: apbose

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -245,20 +245,37 @@ def aten_ops_fmod(
     return impl.elementwise.fmod(ctx, target, SourceIR.ATEN, name, args[0], args[1])
 
 
-@dynamo_tensorrt_converter(torch.ops.aten.grid_sampler.out)
-@dynamo_tensorrt_converter(torch.ops.aten.grid_sampler_backward.out)
+@dynamo_tensorrt_converter(torch.ops.aten.grid_sampler)
 @dynamo_tensorrt_converter(torch.ops.aten.grid_sampler_2d.out)
 @dynamo_tensorrt_converter(torch.ops.aten.grid_sampler_2d_backward.out)
 @dynamo_tensorrt_converter(torch.ops.aten.grid_sampler_3d.out)
 @dynamo_tensorrt_converter(torch.ops.aten.grid_sampler_3d_backward.out)
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+        1: (TRTTensor,),
+    }
+)  # type: ignore[misc]
 def aten_ops_grid(
     ctx: ConversionContext,
     target: Target,
     args: Tuple[Argument, ...],
     kwargs: Dict[str, Argument],
     name: str,
 ) -> Union[TRTTensor, Sequence[TRTTensor]]:
-    return impl.grid.grid(ctx, target, SourceIR.ATEN, name, args[0], args[1], args[2], args[3], args[4])
+    return impl.grid.grid(
+        ctx, 
+        target, 
+        SourceIR.ATEN, 
+        name, 
+        input=args[0], 
+        grid=args[1], 
+        interpolation_mode=args[2], 
+        padding_mode=args[3], 
+        align_corners=args_bounds_check(args, 4, True),
+        output_mask=args_bounds_check(args, 5, None),
+
+    )
 
 
 @dynamo_tensorrt_converter(torch.ops.aten.relu.default)

py/torch_tensorrt/dynamo/conversion/impl/grid.py

@@ -1,14 +1,16 @@
-from typing import Optional
+from typing import Optional, Sequence
 
 import torch
+import tensorrt as trt
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
-from torch_tensorrt.dynamo.conversion.converter_utils import GridSamplerInterpolation, GridSamplerPadding
+from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
+from torch_tensorrt.dynamo.conversion.converter_utils import cast_trt_tensor, GridSamplerInterpolation, GridSamplerSampling
 from torch_tensorrt.fx.converters.converter_utils import set_layer_name
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 
 def grid(
-    network: TRTNetwork,
+    ctx: ConversionContext,
     target: Target,
     source_ir: Optional[SourceIR],
     name: str,
@@ -17,10 +19,21 @@ def grid(
     interpolation_mode: int,
     padding_mode: int,
     align_corners: bool,
+    output_mask: Optional[Sequence[bool]] = None,
 ) -> TRTTensor:
-    grid_layer = network.add_grid_sample(input, grid)
-    grid_layer.interpolation_mode = GridSamplerInterpolation(interpolation_mode)
-    grid_layer.padding_mode = GridSamplerPadding(padding_mode)
+    grid_layer = ctx.net.add_grid_sample(input, grid)
+    interpolation_mode_trt = GridSamplerInterpolation()    
+    grid_layer.interpolation_mode = interpolation_mode_trt(interpolation_mode)
+    sample_mode_trt = GridSamplerSampling()
+    grid_layer.sample_mode = sample_mode_trt(padding_mode) 
     grid_layer.align_corners = align_corners
     set_layer_name(grid_layer, target, name + "_grid_layer", source_ir)
-    return grid_layer.get_output(0)
+    if(output_mask is None):
+        return grid_layer.get_output(0)
+    else:
+        if(output_mask[0] and output_mask[1]):
+            return (grid_layer.get_output(0), None)
+        elif(output_mask[0]):
+            return grid_layer.get_output(0)
+        else:
+            return None

tests/py/dynamo/conversion/test_grid_aten.py

@@ -4,32 +4,32 @@
 from torch.testing._internal.common_utils import run_tests
 from torch_tensorrt import Input
 from parameterized import parameterized
-from .harness import DispatchTestCase
+from harness import DispatchTestCase
 
 class TestGridConverter(DispatchTestCase):
     @parameterized.expand(
         [
-            ("input_grid_interpolation_nearest_sample_fill", [5,5], [5,2], 0, 0),
-            ("input_grid_interpolation_nearest_sample_clamp", [5,5], [5,2], 0, 1),
-            ("input_grid_interpolation_nearest_sample_reflect", [5,5], [5,2], 0, 2),
-            ("input_grid_interpolation_linear_sample_fill", [5,5], [5,2], 1, 0),
-            ("input_grid_interpolation_linear_sample_clamp", [5,5], [5,2], 1, 1),
-            ("input_grid_interpolation_linear_sample_reflect", [5,5], [5,2], 1, 2),
-            ("input_grid_interpolation_cubic_sample_fill", [5,5], [5,2], 2, 0),
-            ("input_grid_interpolation_cubic_sample_clamp", [5,5], [5,2], 2, 1),
-            ("input_grid_interpolation_cubic_sample_reflect", [5,5], [5,2], 2, 2),
+            ("input_grid_interpolation_nearest_sample_fill", [1,1,5,5], [1,5,2,2], 0, 0),
+            ("input_grid_interpolation_nearest_sample_clamp", [1,1,5,5], [1,5,2,2], 0, 1),
+            ("input_grid_interpolation_nearest_sample_reflect", [1,1,5,5], [1,5,2,2], 0, 2),
+            ("input_grid_interpolation_linear_sample_fill", [1,1,5,5], [1,5,2,2], 1, 0),
+            ("input_grid_interpolation_linear_sample_clamp", [1,1,5,5], [1,5,2,2], 1, 1),
+            ("input_grid_interpolation_linear_sample_reflect", [1,1,5,5], [1,5,2,2], 1, 2),
+            ("input_grid_interpolation_cubic_sample_fill", [1,1,5,5], [1,5,2,2], 2, 0),
+            ("input_grid_interpolation_cubic_sample_clamp", [1,1,5,5], [1,5,2,2], 2, 1),
+            ("input_grid_interpolation_cubic_sample_reflect", [1,1,5,5], [1,5,2,2], 2, 2),
         ]
     )
-    def test_grid(self,_, input_shape, dim_shape, interpolation, sample):
+    def test_grid(self, _, input_shape, dim_shape, interpolation, sample):
         class TestModule(nn.Module):
-            def forward(self, x):
-                input = torch.randn(10).reshape(input_shape)
-                grid = torch.randint(-1, 1, dim_shape)
-                return nn.functional.grid(input, grid, interpolation, sample)
-
-        inputs = [torch.randn(1, 10)]
-        self.run_test(TestModule(), inputs, expected_ops={torch.ops.aten.grid_sampler.out})
+            def forward(self, x): 
+                grid = torch.randint(-1, 1, dim_shape, dtype=torch.float32)
+                return torch.ops.aten.grid_sampler(x, grid, interpolation, sample, True)
+        inputs = [torch.randn(input_shape, dtype = torch.float32)]
+        self.run_test(TestModule(), inputs)
 
+if __name__ == "__main__":
+    run_tests()