Skip to content

Commit

Permalink
Remove foreach triton workaround that is no longer needed (pytorch#10…
Browse files Browse the repository at this point in the history
…4016)

Fixes #ISSUE_NUMBER

Pull Request resolved: pytorch#104016
Approved by: https://github.com/eellison
  • Loading branch information
mlazos authored and pytorchmergebot committed Jun 22, 2023
1 parent 1f1fb58 commit 8d9581a
Showing 1 changed file with 0 additions and 7 deletions.
7 changes: 0 additions & 7 deletions torch/_inductor/codegen/triton_foreach.py
Original file line number Diff line number Diff line change
Expand Up @@ -123,13 +123,6 @@ def codegen_kernel(self, name=None):
code.splice("pid = tl.program_id(0)")
code.splice(f"XBLOCK: tl.constexpr = {self.block_size}")

# Initialize all range variables to avoid a triton bug
# with defining vars in if/else blocks
for i in range(next(self.iter_vars_count)):
code.splice("# Note: initialize vars to work around triton bug")
code.splice(f"x{i} = tl.arange(0, XBLOCK)")
code.splice("xmask = tl.arange(0, XBLOCK) > XBLOCK")

for sub_kernel in self.sub_kernels:
num_elems = int(sympy_product(sub_kernel.numels))
self.codegen_pid_range(code, num_elems)
Expand Down

0 comments on commit 8d9581a

Please sign in to comment.