Convert `aten.slice.Tensor` to `ttnn.slice` #179

jdh8 · 2024-09-11T08:11:44Z

Ticket

None

Problem description

PyTorch lowers input[:, :, start:end] to 3 aten.slice.Tensor ops for each dimensions. I'd like to eliminate the outer two no-op slicing ops, but I need help from @kevinwuTT.
Note that ending indices of ttnn.slice are inclusive

What's changed

Convert aten.slice.Tensor to ttnn.slice in supported cases

jdh8 · 2024-09-11T08:13:47Z

I'm getting this error message. I test slicing like tensor[:, :, start:end, :]. What are the specs/restrictions of ttnn.slice?

E       RuntimeError: TT_FATAL @ /home/runner/work/tt-metal/tt-metal/ttnn/cpp/ttnn/operations/data_movement/slice/device/slice_op.cpp:91: (output_tensor_shape[-2] % TILE_HEIGHT == 0) && (this->slice_start[-2] % TILE_HEIGHT == 0)
E       info:
E       Can only unpad tilized tensor with full tiles
E       backtrace:
E        --- ttnn::operations::data_movement::SliceDeviceOperation::validate_with_output_tensors(std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>> const&, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor>>> const&) const

tests/lowering/tensor_manipulation/test_slice.py

ayerofieiev-tt · 2024-09-11T17:32:51Z

I'm getting this error message. I test slicing like tensor[:, :, start:end, :]. What are the specs/restrictions of ttnn.slice?

E       RuntimeError: TT_FATAL @ /home/runner/work/tt-metal/tt-metal/ttnn/cpp/ttnn/operations/data_movement/slice/device/slice_op.cpp:91: (output_tensor_shape[-2] % TILE_HEIGHT == 0) && (this->slice_start[-2] % TILE_HEIGHT == 0)
E       info:
E       Can only unpad tilized tensor with full tiles
E       backtrace:
E        --- ttnn::operations::data_movement::SliceDeviceOperation::validate_with_output_tensors(std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>> const&, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor>>> const&) const

I think if the tensor is tilized, you can't slice to non-tile aligned chunks

sjameelTT · 2024-09-11T17:39:28Z

I'm getting this error message. I test slicing like tensor[:, :, start:end, :]. What are the specs/restrictions of ttnn.slice?

E       RuntimeError: TT_FATAL @ /home/runner/work/tt-metal/tt-metal/ttnn/cpp/ttnn/operations/data_movement/slice/device/slice_op.cpp:91: (output_tensor_shape[-2] % TILE_HEIGHT == 0) && (this->slice_start[-2] % TILE_HEIGHT == 0)
E       info:
E       Can only unpad tilized tensor with full tiles
E       backtrace:
E        --- ttnn::operations::data_movement::SliceDeviceOperation::validate_with_output_tensors(std::__1::vector<tt::tt_metal::Tensor, std::__1::allocator<tt::tt_metal::Tensor>> const&, std::__1::vector<std::__1::optional<tt::tt_metal::Tensor>, std::__1::allocator<std::__1::optional<tt::tt_metal::Tensor>>> const&) const

I think if the tensor is tilized, you can't slice to non-tile aligned chunks

Yeah try ROW major. You're slicing the outer collapsed dim of 1416=64 to 8, which is not tile-aligned. Try row major for this.

ayerofieiev-tt · 2024-09-11T17:45:47Z

torch_ttnn/passes/lowering/add_data_move_pass.py

@@ -118,7 +118,8 @@ def is_function_call(node) -> bool:
    ttnn.permute,
    #  ttnn.repeat,  in target_wrapper
    ttnn.concat,
-    # ttnn.split,  # ttnn has no split, remote the comment in the future when it has
+    # ttnn.split,  # ttnn has no split, remove the comment in the future when it has


@jdh8 please see https://github.com/tenstorrent/tt-metal/tree/main/ttnn/cpp/ttnn/operations/data_movement/split

I made #184 because the patch is useful on its own. After that, ttnn.split becomes a valid conversion target.

ayerofieiev-tt · 2024-09-11T17:48:07Z

torch_ttnn/passes/lowering/to_tt_pass.py

+                    if step != 1 or dim >= len(input_size):
+                        return None
+
+                    # NOTE(jdh8): is there a way to skip a no-op?


@kevinwuTT is there a way to skip an op if its a no-op?

Returning tensor seems fine here. The following code will automatically replace the uses of the aten.split op to point to the arg[0] tensor, which effectively removes this op.

pytorch2.0_ttnn/torch_ttnn/passes/lowering/to_tt_pass.py

Lines 592 to 596 in a5ce2a6

if new_nodes:

node.replace_all_uses_with(

new_nodes[-1],

delete_user_cb=lambda node: node != new_nodes[-1],

)

ayerofieiev-tt · 2024-09-11T17:48:51Z

torch_ttnn/passes/lowering/to_tt_pass.py

+
+                    slice_start[dim] = start
+                    slice_end[dim] = min(end, input_size[dim])
+                    slice_end -= 1


thats because torch is exclusive and ttnn is inclusive, right?

Yes! This is exactly my intention, to convert exclusive ends to inclusive ends.

if you use python style splices (eg) tensor[..., begin:end:step]) then it's exclusive. If you use ttnn.slice() it's inclusive.

I am going to refactor ttnn.slice to be exclusive

…ent pass This patch originates from #179 but is good on its own

…ent pass (#184) This patch originates from #179 but is good on its own Co-authored-by: Artem Yerofieiev <[email protected]>

ayerofieiev-tt · 2024-09-12T15:23:20Z

torch_ttnn/passes/lowering/to_tt_pass.py

+                    return None
+
+                if start == 0 and end >= input_size[dim]:
+                    return tensor


what is the difference in returning None vs tensor?

None marks failed conversion and leaves the graph as is. Returning a tenor lets the tenor replace the original op.

ayerofieiev-tt · 2024-09-12T15:25:54Z

Running manually here
https://github.com/tenstorrent/pytorch2.0_ttnn/actions/runs/10833885090

Tests failed

ayerofieiev-tt · 2024-09-12T15:48:27Z

@jdh8 negative values are now supported with this PR
tenstorrent/tt-metal#12469

So you can improve when this gets merged
#187

ayerofieiev-tt · 2024-09-12T16:59:23Z

@jdh8 that PR is now merged.

jdh8 · 2024-09-11T08:16:36Z

torch_ttnn/passes/lowering/to_tt_pass.py

+
+                    # NOTE(jdh8): is there a way to skip a no-op?
+                    # if start == 0 and end >= input_size[dim]:
+                    #    return tensor


@kevinwuTT, I'd like to eliminate no-op slices, but this doesn't seem to be the right way.

From:

aten.slice.Tensor(tensor, dim, 0, INT64_MAX)

To:

tensor

I guess we can just return the tensor as the new node after merging #180. RAUW of the return value is collectively handled.

pytorch2.0_ttnn/torch_ttnn/passes/lowering/to_tt_pass.py

Lines 595 to 601 in 9f80cdb

with g.inserting_before(node):

new_node = rewrite_node(node)

if new_node is not None:

node.replace_all_uses_with(

new_node,

delete_user_cb=lambda node: node != new_node,

)

jdh8 · 2024-09-11T17:55:21Z

torch_ttnn/passes/lowering/to_tt_pass.py

+
+                    slice_start[dim] = start
+                    slice_end[dim] = min(end, input_size[dim])
+                    slice_end -= 1


Yes! This is exactly my intention, to convert exclusive ends to inclusive ends.

jdh8 · 2024-09-11T18:56:05Z

torch_ttnn/passes/lowering/add_data_move_pass.py

@@ -118,7 +118,8 @@ def is_function_call(node) -> bool:
    ttnn.permute,
    #  ttnn.repeat,  in target_wrapper
    ttnn.concat,
-    # ttnn.split,  # ttnn has no split, remote the comment in the future when it has
+    # ttnn.split,  # ttnn has no split, remove the comment in the future when it has


I made #184 because the patch is useful on its own. After that, ttnn.split becomes a valid conversion target.

jdh8 · 2024-09-12T15:28:40Z

torch_ttnn/passes/lowering/to_tt_pass.py

+                    return None
+
+                if start == 0 and end >= input_size[dim]:
+                    return tensor


None marks failed conversion and leaves the graph as is. Returning a tenor lets the tenor replace the original op.

jdh8 requested review from kevinwuTT and ayerofieiev-tt September 11, 2024 08:11

jdh8 self-assigned this Sep 11, 2024

jdh8 added the conversion label Sep 11, 2024

jdh8 mentioned this pull request Sep 11, 2024

Refactor manual lowering #180

Merged

ayerofieiev-tt reviewed Sep 11, 2024

View reviewed changes

tests/lowering/tensor_manipulation/test_slice.py Show resolved Hide resolved

ayerofieiev-tt reviewed Sep 11, 2024

View reviewed changes

jdh8 added a commit that referenced this pull request Sep 11, 2024

Hook ttnn.split (since 0.51.0) and ttnn.slice into the data movem…

f7f093c

…ent pass This patch originates from #179 but is good on its own

jdh8 mentioned this pull request Sep 11, 2024

Hook ttnn.split (since 0.51.0) and ttnn.slice into the data movement pass #184

Merged

jdh8 added 2 commits September 12, 2024 05:58

Convert aten.slice.Tensor to ttnn.slice

f9d0459

Fix formatting

6dc7eb4

jdh8 force-pushed the feature/slice branch from f1df110 to 6dc7eb4 Compare September 12, 2024 05:58

jdh8 added 2 commits September 12, 2024 06:03

Skip no-op slicing op

29411f8

Reduce one level of nesting

e11a33a

jdh8 force-pushed the feature/slice branch from 684aae4 to e11a33a Compare September 12, 2024 06:15

ayerofieiev-tt added a commit that referenced this pull request Sep 12, 2024

Hook ttnn.split (since 0.51.0) and ttnn.slice into the data movem…

57226ac

…ent pass (#184) This patch originates from #179 but is good on its own Co-authored-by: Artem Yerofieiev <[email protected]>

Merge branch 'main' into feature/slice

0f0c8bd

ayerofieiev-tt reviewed Sep 12, 2024

View reviewed changes

jdh8 commented Sep 13, 2024

View reviewed changes

jdh8 added the blocked label Sep 13, 2024

ayerofieiev-tt added the bert label Sep 13, 2024

Merge branch 'main' into feature/slice

cb05df2

jdh8 and others added 5 commits September 13, 2024 18:20

Fix formatting

5251a21

Merge branch 'main' into feature/slice

9c3d444

Handle non-4d slice inputs; slice in row major

14d2071

Merge branch 'main' into feature/slice

80ff576

Simplify handling of non-4D tensors to slice

6f78bd3

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Convert `aten.slice.Tensor` to `ttnn.slice` #179

Convert `aten.slice.Tensor` to `ttnn.slice` #179

jdh8 commented Sep 11, 2024

jdh8 commented Sep 11, 2024

ayerofieiev-tt commented Sep 11, 2024

sjameelTT commented Sep 11, 2024

ayerofieiev-tt Sep 11, 2024

jdh8 Sep 11, 2024

ayerofieiev-tt Sep 11, 2024

kevinwuTT Sep 11, 2024

ayerofieiev-tt Sep 11, 2024

jdh8 Sep 11, 2024

sjameelTT Sep 11, 2024

sjameelTT Sep 12, 2024

ayerofieiev-tt Sep 12, 2024

jdh8 Sep 12, 2024

ayerofieiev-tt commented Sep 12, 2024 •

edited

Loading

ayerofieiev-tt commented Sep 12, 2024

ayerofieiev-tt commented Sep 12, 2024

jdh8 Sep 11, 2024

jdh8 Sep 11, 2024

jdh8 Sep 11, 2024

jdh8 Sep 11, 2024

jdh8 Sep 12, 2024

	if new_nodes:
	node.replace_all_uses_with(
	new_nodes[-1],
	delete_user_cb=lambda node: node != new_nodes[-1],
	)

	with g.inserting_before(node):
	new_node = rewrite_node(node)
	if new_node is not None:
	node.replace_all_uses_with(
	new_node,
	delete_user_cb=lambda node: node != new_node,
	)

Convert aten.slice.Tensor to ttnn.slice #179

Are you sure you want to change the base?

Convert aten.slice.Tensor to ttnn.slice #179

Conversation

jdh8 commented Sep 11, 2024

Ticket

Problem description

What's changed

jdh8 commented Sep 11, 2024

ayerofieiev-tt commented Sep 11, 2024

sjameelTT commented Sep 11, 2024

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

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

ayerofieiev-tt commented Sep 12, 2024 • edited Loading

ayerofieiev-tt commented Sep 12, 2024

ayerofieiev-tt commented Sep 12, 2024

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

Convert `aten.slice.Tensor` to `ttnn.slice` #179

Convert `aten.slice.Tensor` to `ttnn.slice` #179

ayerofieiev-tt commented Sep 12, 2024 •

edited

Loading