From d7d006ec39b45fcb2e32baf1645723d73480cbc5 Mon Sep 17 00:00:00 2001
From: Meenakshi Ramanathan <mramanathan@tenstorrent.com>
Date: Thu, 9 Jan 2025 11:15:23 +0530
Subject: [PATCH] Add test for grid_sample and floor op (#953)

Test cases for grid sample op and floor op has been added.
Assertion from tm.py has been removed to support greater than 1 batch
size inputs for grid_sample op.
---
 forge/forge/op/eval/forge/tm.py |  2 --
 forge/test/mlir/test_ops.py     | 60 +++++++++++++++++++++++++++++++++
 third_party/tvm                 |  2 +-
 3 files changed, 61 insertions(+), 3 deletions(-)

diff --git a/forge/forge/op/eval/forge/tm.py b/forge/forge/op/eval/forge/tm.py
index 6fc35a1fb..0e580d90a 100644
--- a/forge/forge/op/eval/forge/tm.py
+++ b/forge/forge/op/eval/forge/tm.py
@@ -1028,8 +1028,6 @@ def decompose(type, attr, dc, inputs):
             # by adding the size on that dimension
             stop = act.shape[dim] + stop
 
-        assert dim != -4, "No support for indexing on dimension -4 (w)"
-
         is_one_dim = len(act.shape) == 1
         if is_one_dim:
             # If input is a one-dimensional tensor, reshape it to a 2D tensor with one dimension equal to 1
diff --git a/forge/test/mlir/test_ops.py b/forge/test/mlir/test_ops.py
index a50f84068..c222447c9 100644
--- a/forge/test/mlir/test_ops.py
+++ b/forge/test/mlir/test_ops.py
@@ -7,6 +7,7 @@
 import pytest
 import torch
 from torch import nn
+import torch.nn.functional as F
 
 import forge
 from forge.tensor import to_forge_tensors
@@ -2017,3 +2018,62 @@ def forward(self, x):
     compiled_model = forge.compile(framework_model, sample_inputs=inputs)
 
     verify(inputs, framework_model, compiled_model)
+
+
+@pytest.mark.xfail(
+    reason="RuntimeError: Found Unsupported operations while lowering from TTForge to TTIR in forward graph - adv_index"
+)
+@pytest.mark.parametrize(
+    "img, grid",
+    [
+        ((1, 2, 4, 4), (1, 6, 2, 2)),
+        ((1, 32, 50, 50), (1, 2500, 4, 2)),
+        ((1, 3, 8, 8), (1, 3, 3, 2)),
+        ((1, 3, 16, 16), (1, 8, 8, 2)),
+        ((5, 2, 10, 10), (5, 12, 3, 2)),
+        ((3, 8, 32, 32), (3, 25, 4, 2)),
+    ],
+)
+@pytest.mark.parametrize("align_corners", [True, False])
+def test_grid_sample(img, grid, align_corners, test_device):
+    class GridSampleModule(nn.Module):
+        def __init__(self, interpolation="bilinear", align_corners=align_corners):
+            super(GridSampleModule, self).__init__()
+            self.interpolation = interpolation
+            self.align_corners = align_corners
+
+        def forward(self, img, grid):
+            output = F.grid_sample(img, grid, mode=self.interpolation, align_corners=align_corners)
+            return output
+
+    # TO-DO: Support for nearest interpolation mode is yet to be added
+    model = GridSampleModule(interpolation="bilinear", align_corners=align_corners)
+    model.eval()
+    img = torch.randn(img)
+    grid = torch.randn(grid)
+    output = model(img, grid)
+    compiled_model = forge.compile(model, sample_inputs=[img, grid], module_name="grid_sample")
+
+
+@pytest.mark.xfail(reason="RuntimeError: BinaryOpType cannot be mapped to BcastOpMath")
+@pytest.mark.parametrize(
+    "input_data",
+    [
+        torch.tensor([-0.8166, 1.5308, -0.2530, -0.2091]),
+        torch.tensor([-3.7, -1.2, 0.0, 1.5, 3.9]),
+        torch.tensor([1.0, 2.0, -1.0, -2.0]),
+        torch.tensor([-12345.678, 12345.678, -0.999, 0.999, 3.14159, -3.14159]),
+    ],
+)
+def test_floor(input_data):
+    class Floor(nn.Module):
+        def __init__(self):
+            super().__init__()
+
+        def forward(self, a):
+            return torch.floor(a)
+
+    framework_model = Floor()
+    compiled_model = forge.compile(framework_model, sample_inputs=[input_data], module_name="floor")
+
+    verify([input_data], framework_model, compiled_model)
diff --git a/third_party/tvm b/third_party/tvm
index 0a429cfcb..8aed232d1 160000
--- a/third_party/tvm
+++ b/third_party/tvm
@@ -1 +1 @@
-Subproject commit 0a429cfcbe016178709f62ce653437df4f0e7a4f
+Subproject commit 8aed232d15fae56f2e8946cb738d2ec01e59b866