Skip to content

Tessera Pattern Matching with Struct Arguments#2437

Draft
jessicacotturone21 wants to merge 7 commits into
mainfrom
tessera-calls
Draft

Tessera Pattern Matching with Struct Arguments#2437
jessicacotturone21 wants to merge 7 commits into
mainfrom
tessera-calls

Conversation

@jessicacotturone21
Copy link
Copy Markdown
Collaborator

@jessicacotturone21 jessicacotturone21 commented Apr 17, 2026

The goal of this pull request is to allow for more pattern matching to occur by eliminating intermediate stores when we have struct arguments and returns. We also allow users to specify whether arguments are byref in the source code (as well as the size of the arguments), and any arguments specified as byref will be loaded and converted to byval in our tessera passes.

The user can now specify that an argument is byref like this:
tessera_op = eigen.inv(x: byref, sizeof(Mat))
where Mat is a struct and x is the matrix we want to take the inverse of.

Previously, after running the llvm-to-tessera pass and pulling out the sret argument, we would end up with something that looks like this:

%3 = tessera.call @eigen.inv(%2)
llvm.store %3, %1
%4 = tessera.call @eigen.inv(%1)

By loading the struct value from the pointer if it is byref and checking that the tessera.call op does not capture the pointer, we can perform store-to-load forwarding with the polygeist-mem2reg pass and end up with:

%3 = tessera.call @eigen.inv(%2)
%4 = tessera.call @eigen.inv(%3)

which our PDL pattern can match. This pull request updates the polygeist-mem2reg pass with additional type conversion logic and adds a canonicalization on tessera calls to allow this to happen.

Any operands that were changed from pointers to loaded values are kept track of and new llvm.alloca instances are created and pointer operands are reinserted when lowering back to LLVM.

The MemoryEffectOpInterface has also been implemented and pure_tessera_ops added, so that a user can declare a function as side effect free. This will allow for the elimination of lingering tessera.call ops that remain after the optimization rewrite has been performed.

@jessicacotturone21 jessicacotturone21 changed the title Tessera Pattern Matching with Struct Arguments Tessera Pattern Matching with Struct Operands Apr 17, 2026
@jessicacotturone21 jessicacotturone21 changed the title Tessera Pattern Matching with Struct Operands Tessera Pattern Matching with Struct Arguments Apr 17, 2026
if (getOperand(i - startIdx).getType() == fnType.getInput(i))
continue;
if (isa<LLVM::LLVMPointerType>(fnType.getInput(i)) &&
fn.getArgAttr(i, LLVM::LLVMDialect::getReadonlyAttrName()))
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.

you also need to check nocapture, and it can be readonly or readnone

Comment thread src/enzyme_ad/jax/Passes/Tessera/LLVMToTessera.cpp Outdated
Copy link
Copy Markdown
Member

@wsmoses wsmoses left a comment

Choose a reason for hiding this comment

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

consider a fn

double* pointeradd(double* x, int y) {
   return &x[y];
}

we could define an opt pointeradd(pointeradd(x, y), z) -> pointeradd(x, y+z)

double x[10] -> double* [ptr]

@jessicacotturone21 jessicacotturone21 marked this pull request as draft April 19, 2026 01:18
@jessicacotturone21 jessicacotturone21 force-pushed the tessera-calls branch 2 times, most recently from 64e632d to 9edf2eb Compare April 24, 2026 13:49
@codecov
Copy link
Copy Markdown

codecov Bot commented Apr 24, 2026

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 25.47%. Comparing base (aeecbf9) to head (cb0df8e).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files
@@           Coverage Diff           @@
##             main    #2437   +/-   ##
=======================================
  Coverage   25.47%   25.47%           
=======================================
  Files         220      220           
  Lines       44600    44600           
=======================================
  Hits        11364    11364           
  Misses      33236    33236           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@jessicacotturone21 jessicacotturone21 force-pushed the tessera-calls branch 2 times, most recently from 02a826d to d0393b4 Compare May 1, 2026 00:32
@jessicacotturone21 jessicacotturone21 marked this pull request as ready for review May 11, 2026 18:46
@jessicacotturone21 jessicacotturone21 marked this pull request as draft May 16, 2026 19:15
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.

2 participants