Skip to content

properly codegen hl.triton_kernel#1797

Merged
shunting314 merged 1 commit intomainfrom
shunting314/stack/18
Mar 26, 2026
Merged

properly codegen hl.triton_kernel#1797
shunting314 merged 1 commit intomainfrom
shunting314/stack/18

Conversation

@shunting314
Copy link
Copy Markdown
Contributor

@shunting314 shunting314 commented Mar 24, 2026

Stacked PRs:


properly codegen hl.triton_kernel

fix #1642 . When we roll reductions, we need consider side-effect of a node. We can not lift the node out of a loop if it has side effect. Instead, we should close the current loop, codegen the side-effect node and start a new loop.

shunting314 added a commit that referenced this pull request Mar 24, 2026
stack-info: PR: #1797, branch: shunting314/stack/18
@shunting314 shunting314 force-pushed the shunting314/stack/18 branch from 2afcd24 to 9160dce Compare March 24, 2026 04:19
@shunting314 shunting314 requested review from jansel and yf225 March 24, 2026 04:19
@meta-cla meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Mar 24, 2026
@shunting314 shunting314 marked this pull request as draft March 24, 2026 04:44
@shunting314 shunting314 changed the base branch from shunting314/stack/14 to main March 24, 2026 04:44
@shunting314 shunting314 changed the base branch from main to shunting314/stack/14 March 24, 2026 04:44
@shunting314 shunting314 marked this pull request as ready for review March 24, 2026 04:44
@shunting314 shunting314 marked this pull request as draft March 24, 2026 04:57
@shunting314 shunting314 changed the base branch from shunting314/stack/14 to main March 24, 2026 04:57
@shunting314 shunting314 changed the base branch from main to shunting314/stack/14 March 24, 2026 04:58
@shunting314 shunting314 marked this pull request as ready for review March 24, 2026 04:58
@shunting314 shunting314 marked this pull request as draft March 24, 2026 17:51
@shunting314 shunting314 changed the base branch from shunting314/stack/14 to main March 24, 2026 17:51
@shunting314 shunting314 force-pushed the shunting314/stack/18 branch from 9160dce to f0cc4bd Compare March 24, 2026 17:52
@shunting314 shunting314 changed the base branch from main to shunting314/stack/14 March 24, 2026 17:52
@shunting314 shunting314 marked this pull request as ready for review March 24, 2026 17:52
@shunting314 shunting314 marked this pull request as draft March 24, 2026 19:03
@shunting314 shunting314 changed the base branch from shunting314/stack/14 to main March 24, 2026 19:03
@shunting314 shunting314 changed the base branch from shunting314/stack/14 to main March 25, 2026 17:24
@shunting314 shunting314 force-pushed the shunting314/stack/18 branch from da612eb to bfdb55a Compare March 25, 2026 17:25
@shunting314 shunting314 changed the base branch from main to shunting314/stack/14 March 25, 2026 17:25
@shunting314 shunting314 marked this pull request as ready for review March 25, 2026 17:25
@shunting314 shunting314 marked this pull request as draft March 25, 2026 17:35
@shunting314 shunting314 changed the base branch from shunting314/stack/14 to main March 25, 2026 17:35
@shunting314 shunting314 force-pushed the shunting314/stack/18 branch from bfdb55a to 4d2178a Compare March 25, 2026 17:35
@shunting314 shunting314 changed the base branch from main to shunting314/stack/14 March 25, 2026 17:35
@shunting314 shunting314 marked this pull request as ready for review March 25, 2026 17:35
if (
not all((n in self.available) for n in node.all_input_nodes)
or node.op == "output"
or (node.is_impure() and self.inner_count > 0)
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Do we have a test for the issue this is fixing?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

test_distributed.py has a test covering this case

fixed_config = helion.Config(
block_sizes=[8],
num_warps=8,
block_sizes=[8], num_warps=8, reduction_loops=[1024]
Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

@jansel this test covered the issue fixed.

@shunting314 shunting314 marked this pull request as draft March 26, 2026 06:31
@shunting314 shunting314 changed the base branch from shunting314/stack/14 to main March 26, 2026 06:31
@shunting314 shunting314 force-pushed the shunting314/stack/18 branch from 4d2178a to 3cecf1e Compare March 26, 2026 06:31
@shunting314 shunting314 changed the base branch from main to shunting314/stack/14 March 26, 2026 06:31
@shunting314 shunting314 marked this pull request as ready for review March 26, 2026 06:31
@shunting314 shunting314 marked this pull request as draft March 26, 2026 06:37
@shunting314 shunting314 changed the base branch from shunting314/stack/14 to main March 26, 2026 06:37
@shunting314 shunting314 force-pushed the shunting314/stack/18 branch from 3cecf1e to 1724043 Compare March 26, 2026 06:37
@shunting314 shunting314 changed the base branch from main to shunting314/stack/14 March 26, 2026 06:37
@shunting314 shunting314 marked this pull request as ready for review March 26, 2026 06:37
stack-info: PR: #1797, branch: shunting314/stack/18
@shunting314 shunting314 marked this pull request as draft March 26, 2026 06:49
@shunting314 shunting314 changed the base branch from shunting314/stack/14 to main March 26, 2026 06:49
@shunting314 shunting314 force-pushed the shunting314/stack/18 branch from 1724043 to c5a7440 Compare March 26, 2026 06:49
@shunting314 shunting314 marked this pull request as ready for review March 26, 2026 06:49
@shunting314 shunting314 merged commit 6eaf5ee into main Mar 26, 2026
8 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

hl.triton_kernel codegen at the wrong places

3 participants