Support non-aligned shape

tests/lowering/eltwise/unary/test_fill.py

@@ -13,21 +13,24 @@ def forward(self, input, value):
 
 
 @pytest.mark.parametrize(
-    "input_shape",
+    "input_shape, value, converted",
     [
-        (64,),
-        (32, 32),
-        (1, 64, 32),
-        (2, 32, 64),
-        (32, 1),
-        (1, 32),
-        (16, 64),
+        ((32, 32), 0.5, True),
+        ((32, 1), 114, True),
+        ((1, 32), 120, True),
+        ((4, 4), 32, True),
+        ((1, 2, 100, 100), 16, True),
+        ((1, 1, 1, 1), 61, True),
+        ((2, 1, 1, 1), 64, True),
+        ((1, 1, 2, 1), 8, True),
+        # Not supported: dims > 4 or < 2
+        ((1, 1, 1, 2, 2), 1.0, False),
+        ((64,), 1.0, False),
     ],
 )
-def test_fill_scalar(device, input_shape):
+def test_fill_scalar(device, input_shape, value, converted):
     m = FillScalarModule()
     input = torch.rand(input_shape, dtype=torch.bfloat16)
-    value = 0.125
     result_before = m.forward(input, value)
 
     option = torch_ttnn.TorchTtnnOption(device=device, gen_graphviz=True)
@@ -38,10 +41,6 @@ def test_fill_scalar(device, input_shape):
 
     # Check the graph has be rewritten and contains ttnn ops
     nodes = list(option._out_fx_graphs[0].nodes)
-    # Currently only tensor sizes divisible by 32x32 are converted.
-    if len(input_shape) > 1 and input_shape[-1] % ttnn.TILE_SIZE == 0 and input_shape[-2] % ttnn.TILE_SIZE == 0:
-        assert [node.target for node in nodes].count(ttnn.full) == 1
-    else:
-        assert [node.target for node in nodes].count(ttnn.full) == 0
+    assert [node.target for node in nodes].count(ttnn.full) == (1 if converted else 0)
     # Check inference result
     assert torch.allclose(result_before, result_after)

torch_ttnn/passes/lowering/add_data_move_pass.py

@@ -123,7 +123,7 @@ def is_function_call(node) -> bool:
     ttnn.to_layout,
 ]
 
-TTNN_TARGET_WRAPPERS = [target_wrappers.clone, target_wrappers.repeat]
+TTNN_TARGET_WRAPPERS = [target_wrappers.clone, target_wrappers.repeat, target_wrappers.getitem]
 
 TTNN_NORM_OPS = [
     ttnn.group_norm,

torch_ttnn/passes/lowering/target_wrappers.py

@@ -10,3 +10,8 @@ def clone(t):
 @torch.fx.wrap
 def repeat(t, sizes):
     return ttnn.repeat(t, ttnn.Shape(sizes))
+
+
+@torch.fx.wrap
+def getitem(t, bounds):
+    return t[[slice(lower, upper) for lower, upper in bounds]]

torch_ttnn/passes/lowering/to_tt_pass.py

@@ -443,18 +443,19 @@ def rewrite_node(node):
                 return None
 
             if node.target == torch.ops.aten.fill.Scalar:
-                shape = tuple(node.meta["val"].size())
-                if has_valid_page_size(shape, strict=True):
-                    new_kwargs = {
-                        "fill_value": args[1],
-                        "device": TtnnDevice(),
-                        "layout": TtnnTileLayout(),
-                    }
-                    return g.call_function(
-                        ttnn.full,
-                        args=(shape,),
-                        kwargs=new_kwargs,
-                    )
+                shape = list(node.meta["val"].size())
+                if len(shape) < 2 or len(shape) > 4:
+                    return None
+                aligned_shape = shape[:-2] + [
+                    (dim + ttnn.TILE_SIZE - 1) // ttnn.TILE_SIZE * ttnn.TILE_SIZE for dim in shape[-2:]
+                ]
+                new_node = g.call_function(
+                    ttnn.full,
+                    args=(aligned_shape,),
+                    kwargs={"fill_value": args[1], "device": TtnnDevice(), "layout": TtnnTileLayout()},
+                )
+                bounds = [(0, dim) for dim in shape]
+                return g.call_function(target_wrappers.getitem, args=(new_node, bounds))
 
             if node.target == torch.ops.aten.baddbmm.default:
                 # out = beta * input + alpha * (batch1 @ batch2)