Skip to content

[WIP] Starter FWD mode rules for triton.#2549

Draft
spikerheado1234 wants to merge 2 commits into
EnzymeAD:mainfrom
spikerheado1234:main-rebuilt
Draft

[WIP] Starter FWD mode rules for triton.#2549
spikerheado1234 wants to merge 2 commits into
EnzymeAD:mainfrom
spikerheado1234:main-rebuilt

Conversation

@spikerheado1234

Copy link
Copy Markdown
Collaborator

Adds some expressive fwd mode rules for ttir ops as well as simple test cases to demonstrate their correctness.

Comment thread test/lit_tests/diffrules/triton/err.txt Outdated

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

can we delete this

Comment thread test/lit_tests/diffrules/triton/out.txt Outdated

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

same here

@vimarsh6739 vimarsh6739 requested a review from dcw02 June 10, 2026 18:28
@wsmoses

wsmoses commented Jun 10, 2026

Copy link
Copy Markdown
Member

cc @Pangoraw @avik-pal

@@ -0,0 +1,64 @@
// RUN: enzymexlamlir-opt %s --enzyme-wrap="infn=main outfn= argTys=enzyme_dup,enzyme_dup,enzyme_const,enzyme_const,enzyme_const,enzyme_const retTys=enzyme_dup,enzyme_dup mode=ForwardMode" --canonicalize | FileCheck %s

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

forgot to add file checks patterns here?

@@ -0,0 +1,37 @@
// RUN: enzymexlamlir-opt %s -canonicalize | FileCheck %s

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

forgot to add file checks patterns here?

@@ -0,0 +1,37 @@
// RUN: enzymexlamlir-opt %s -canonicalize | FileCheck %s

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

same

/// shadow = dc (or zero if c is constant)
/// if a is active: shadow = dot(da, b, shadow)
/// if b is active: shadow = dot(a, db, shadow)
class AutoDiffTritonDotFwd

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

this pattern could be implemented in tablegen, it may need to specify an unimplemented rule for the reverse

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants