Skip to content
This repository was archived by the owner on Aug 1, 2025. It is now read-only.

[inductor] Revise #956 to avoid a perf regression #981

Merged
merged 2 commits into from
Aug 24, 2022
Merged

[inductor] Revise #956 to avoid a perf regression #981

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b4c38bb
[inductor] Rework #956 to avoid a perf regression
desertfire Aug 24, 2022
4689b59
Change how index 0 load is implemented
desertfire Aug 24, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 7 additions & 2 deletions torchinductor/codegen/triton.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -652,7 +652,7 @@ def indexing(
mask.append(f"{tree.prefix}mask")
dense_mask.append(f"{tree.prefix}mask")

if (need_dense and not have_dense) or index == 0:
if need_dense and not have_dense:
mask = dense_mask
index_str = f"{index_str} + tl.zeros({self.dense_size_str()}, tl.int32)"
Comment on lines +655 to 657
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
if need_dense and not have_dense:
mask = dense_mask
index_str = f"{index_str} + tl.zeros({self.dense_size_str()}, tl.int32)"
if need_dense and not have_dense or index == 0:
index_str = f"{index_str} + tl.zeros({self.dense_size_str()}, tl.int32)"
if index == 0:
mask = ["None"]
else:
mask = dense_mask

and remove changes to load. If index is 0, you don't need mask at all, a single element load is always valid.

elif not have_loop_vars and copy_shape:
Expand Down Expand Up @@ -704,7 +704,12 @@ def mask_loads(self, mask):
def load(self, name: str, index: sympy.Expr, upcast: bool = False):
var = self.args.input(name)
indirect_indexing = self.is_indirect_indexing(index)
index, mask = self.indexing(index)
if index == 0:
# No need to use mask when loading a single element from index 0
index, mask = "0", "None"
else:
index, mask = self.indexing(index)

if "rmask" in mask:
# This eviction policy heuristic is untested.
# ptillet suggested we should try only doing this for
Expand Down