Add epilogue subtiling #948

PaulZhang12 · 2025-10-15T20:59:34Z

Stacked PRs:

->Add epilogue subtiling #948

Add epilogue subtiling

stack-info: PR: #948, branch: PaulZhang12/stack/14

oulgen · 2025-10-15T21:57:27Z

examples/matmul.py

+    config=helion.Config(
+        block_sizes=[64, 64, 64],
+        loop_orders=[[0, 1]],
+        l2_groupings=[4],
+        range_unroll_factors=[0, 1],
+        range_num_stages=[0, 3],
+        range_multi_buffers=[None, False],
+        range_flattens=[None, None],
+        num_warps=8,
+        num_stages=6,
+        indexing='tensor_descriptor',
+        pid_type='flat'
+    )


you probably dont want to check this in since the best config will depend on the machine

stack-info: PR: #948, branch: PaulZhang12/stack/14

jansel

Does this help with matmul perf?

jansel · 2025-10-16T00:10:05Z

helion/_compiler/device_function.py

+        import re
        host_function = HostFunction.current()
        block_size_expr = ", ".join(map(self.literal_expr, block_size))
+        pattern = r'triton_helpers\.div_floor_integer\(([^,]+),\s*(\d+)\)'


@yf225 didn't you add something to fix this somewhere else?

yes - we have sanitization pass for triton_helpers.* right now at

helion/helion/_compiler/device_function.py

Lines 519 to 532 in 1aaba3f

if isinstance(value, sympy.Expr):

sanitized = value.replace( # pyright: ignore[reportAttributeAccessIssue]

lambda node: isinstance(node, sympy.Function)

and getattr(node.func, "__name__", "")

== "triton_helpers.div_floor_integer",

lambda node: sympy.floor(node.args[0] / node.args[1]), # pyright: ignore[reportAttributeAccessIssue]

).replace( # pyright: ignore[reportAttributeAccessIssue]

lambda node: isinstance(node, sympy.Function)

and getattr(node.func, "__name__", "")

== "triton_helpers.remainder_integer",

lambda node: sympy.Mod(node.args[0], node.args[1]), # pyright: ignore[reportAttributeAccessIssue]

)

expr = cast("sympy.Expr", sanitized)

return HostFunction.current().sympy_expr(expr)

for constexpr arg, maybe we can extract a common util function to be used in both sites

jansel · 2025-10-16T00:10:16Z

examples/matmul.py

    Returns:
        Tensor: Resulting matrix of shape [m, n].
    """
+


revert unrelated change

jansel · 2025-10-16T00:10:49Z

helion/_compiler/indexing_strategy.py


    return None

+def _supports_epilogue_subtiling():


Should this be the same as the supports_tensor_descriptor helper we already have?

jansel · 2025-10-16T00:11:56Z

helion/autotuner/config_spec.py

        config.setdefault(
            "load_eviction_policies", self.load_eviction_policies.default()
        )
+        config.setdefault("epilogue_subtiling", False)


Should this be a list since we can have multiple stores in the program?

helion/_compiler/indexing_strategy.py

stack-info: PR: #948, branch: PaulZhang12/stack/14

yf225 · 2025-10-20T23:00:38Z

helion/_compiler/device_function.py

        block_size_expr = ", ".join(map(self.literal_expr, block_size))
+        pattern = r"triton_helpers\.div_floor_integer\(([^,]+),\s*(\d+)\)"
+        replacement = r"\1 // \2"
+        block_size_expr = re.sub(pattern, replacement, block_size_expr)


I believe this should unify the triton_helpers handling: https://gist.github.com/yf225/dbb4045c44df97c902906290e0f6affa

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 added a commit that referenced this pull request Oct 15, 2025

Add epilogue subtiling

fcc7492

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from cf439ac to fcc7492 Compare October 15, 2025 20:59

meta-cla bot added the CLA Signed This label is managed by the Meta Open Source bot. label Oct 15, 2025

PaulZhang12 added a commit that referenced this pull request Oct 15, 2025

Add epilogue subtiling

cdbedf6

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from fcc7492 to cdbedf6 Compare October 15, 2025 21:05

PaulZhang12 added a commit that referenced this pull request Oct 15, 2025

Add epilogue subtiling

58496fb

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from cdbedf6 to 58496fb Compare October 15, 2025 21:05

oulgen reviewed Oct 15, 2025

View reviewed changes

PaulZhang12 added a commit that referenced this pull request Oct 15, 2025

Add epilogue subtiling

965b193

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from 58496fb to 965b193 Compare October 15, 2025 22:00

jansel requested changes Oct 16, 2025

View reviewed changes

PaulZhang12 added a commit that referenced this pull request Oct 16, 2025

Add epilogue subtiling

2bc36d0

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from 965b193 to 2bc36d0 Compare October 16, 2025 01:53

PaulZhang12 added a commit that referenced this pull request Oct 17, 2025

Add epilogue subtiling

1c1e282

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from 2bc36d0 to 1c1e282 Compare October 17, 2025 22:11

PaulZhang12 added a commit that referenced this pull request Oct 20, 2025

Add epilogue subtiling

cccb0af

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from 1c1e282 to cccb0af Compare October 20, 2025 19:19

PaulZhang12 added a commit that referenced this pull request Oct 20, 2025

Add epilogue subtiling

a6dd082

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from cccb0af to a6dd082 Compare October 20, 2025 19:19

PaulZhang12 mentioned this pull request Oct 20, 2025

Remove unrolling with tma + pipelining #994

Merged

PaulZhang12 changed the base branch from main to PaulZhang12/stack/16 October 20, 2025 19:20

PaulZhang12 changed the base branch from PaulZhang12/stack/16 to main October 20, 2025 19:22

PaulZhang12 added a commit that referenced this pull request Oct 20, 2025

Add epilogue subtiling

88d46a8

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from a6dd082 to 88d46a8 Compare October 20, 2025 19:22

PaulZhang12 changed the base branch from main to PaulZhang12/stack/16 October 20, 2025 19:22

PaulZhang12 changed the base branch from PaulZhang12/stack/16 to main October 20, 2025 19:24

PaulZhang12 added a commit that referenced this pull request Oct 20, 2025

Add epilogue subtiling

48eed82

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from 88d46a8 to 48eed82 Compare October 20, 2025 19:24

PaulZhang12 changed the base branch from main to PaulZhang12/stack/16 October 20, 2025 19:25

PaulZhang12 changed the base branch from PaulZhang12/stack/16 to main October 20, 2025 19:28

PaulZhang12 changed the base branch from PaulZhang12/stack/16 to main October 20, 2025 19:48

PaulZhang12 added a commit that referenced this pull request Oct 20, 2025

Add epilogue subtiling

ec85d55

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from b5692b3 to ec85d55 Compare October 20, 2025 19:48

PaulZhang12 changed the base branch from main to PaulZhang12/stack/16 October 20, 2025 19:48

PaulZhang12 changed the base branch from PaulZhang12/stack/16 to main October 20, 2025 20:00

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from ec85d55 to c315444 Compare October 20, 2025 20:00

PaulZhang12 added a commit that referenced this pull request Oct 20, 2025

Add epilogue subtiling

c315444

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 changed the base branch from main to PaulZhang12/stack/16 October 20, 2025 20:00

PaulZhang12 changed the base branch from PaulZhang12/stack/16 to main October 20, 2025 21:34

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from c315444 to c0ed00f Compare October 20, 2025 21:34

PaulZhang12 mentioned this pull request Oct 20, 2025

Prioritize outermost loop for warp spec #995

Merged

yf225 reviewed Oct 20, 2025

View reviewed changes

PaulZhang12 added a commit that referenced this pull request Oct 21, 2025

Add epilogue subtiling

589fc42

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from c0ed00f to 589fc42 Compare October 21, 2025 20:56

PaulZhang12 added a commit that referenced this pull request Oct 21, 2025

Add epilogue subtiling

2da3497

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from 589fc42 to 2da3497 Compare October 21, 2025 22:14

PaulZhang12 added a commit that referenced this pull request Oct 21, 2025

Add epilogue subtiling

3e336b1

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from 2da3497 to 3e336b1 Compare October 21, 2025 22:24

PaulZhang12 requested review from jansel, njriasan, oulgen and yf225 October 21, 2025 22:37

PaulZhang12 added a commit that referenced this pull request Oct 21, 2025

Add epilogue subtiling

ded2521

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from 3e336b1 to ded2521 Compare October 21, 2025 22:49

PaulZhang12 added a commit that referenced this pull request Oct 21, 2025

Add epilogue subtiling

ad6ca82

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from ded2521 to ad6ca82 Compare October 21, 2025 22:51

PaulZhang12 added a commit that referenced this pull request Oct 21, 2025

Add epilogue subtiling

0c3d607

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from 0826b24 to 0c3d607 Compare October 21, 2025 23:59

Add epilogue subtiling

bdf0793

stack-info: PR: #948, branch: PaulZhang12/stack/14

PaulZhang12 force-pushed the PaulZhang12/stack/14 branch from 0c3d607 to bdf0793 Compare October 22, 2025 00:08

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add epilogue subtiling #948

Add epilogue subtiling #948

Uh oh!

PaulZhang12 commented Oct 15, 2025 •

edited

Loading

Uh oh!

oulgen Oct 15, 2025

Uh oh!

jansel left a comment

Uh oh!

jansel Oct 16, 2025

Uh oh!

yf225 Oct 16, 2025

Uh oh!

jansel Oct 16, 2025

Uh oh!

jansel Oct 16, 2025

Uh oh!

jansel Oct 16, 2025

Uh oh!

Uh oh!

yf225 Oct 20, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

	if isinstance(value, sympy.Expr):
	sanitized = value.replace( # pyright: ignore[reportAttributeAccessIssue]
	lambda node: isinstance(node, sympy.Function)
	and getattr(node.func, "__name__", "")
	== "triton_helpers.div_floor_integer",
	lambda node: sympy.floor(node.args[0] / node.args[1]), # pyright: ignore[reportAttributeAccessIssue]
	).replace( # pyright: ignore[reportAttributeAccessIssue]
	lambda node: isinstance(node, sympy.Function)
	and getattr(node.func, "__name__", "")
	== "triton_helpers.remainder_integer",
	lambda node: sympy.Mod(node.args[0], node.args[1]), # pyright: ignore[reportAttributeAccessIssue]
	)
	expr = cast("sympy.Expr", sanitized)
	return HostFunction.current().sympy_expr(expr)

Add epilogue subtiling #948

Are you sure you want to change the base?

Add epilogue subtiling #948

Uh oh!

Conversation

PaulZhang12 commented Oct 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!