Skip to content

Conversation

jessicacotturone21
Copy link
Collaborator

Look at "Draft of Func to Tessera Pass" commit.
Not yet tested.

Copy link
Contributor

@github-actions github-actions bot left a comment

Choose a reason for hiding this comment

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

EnzymeJAX Benchmarks

Benchmark suite Current: c4d10eb Previous: 7f88655 Ratio
scatter_sum / JaX / cpu / Primal 0.000004314978991169483 s 0.000004144166014157235 s 1.04
scatter_sum / JaXPipe / cpu / Primal 0.000004397982999216765 s 0.000004226277908310294 s 1.04
scatter_sum / JaX / tpu / Primal 0.0001386552981915 s 0.0001389866188168 s 1.00
scatter_sum / JaXPipe / tpu / Primal 0.0001493743569939 s 0.0001376257929019 s 1.09

This comment was automatically generated by workflow using github-action-benchmark.

PatternRewriter &rewriter) const override {
FunctionType fnType = funcOp.getFunctionType();

if (fnType.getNumResults() > 1)
Copy link
Member

Choose a reason for hiding this comment

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

not sure this is required?

LogicalResult
matchAndRewrite(func::ReturnOp returnOp,
PatternRewriter &rewriter) const override {

Copy link
Member

Choose a reason for hiding this comment

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

we should first check whether our parent op is a tessera defineop and only then do the replacement

matchAndRewrite(func::FuncOp funcOp,
PatternRewriter &rewriter) const override {
FunctionType fnType = funcOp.getFunctionType();

Copy link
Member

Choose a reason for hiding this comment

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

eventually we will want to do some check if funcOp.hasAttr("tessera.convert") then do this conversion, otherwise leave the funcop alone

matchAndRewrite(func::CallOp callOp,
PatternRewriter &rewriter) const override {

rewriter.replaceOpWithNewOp<tessera::CallOp>(callOp, callOp.getResultTypes(),
Copy link
Member

Choose a reason for hiding this comment

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

we need to get the symbol we're calling into, and only if that symbol is a tessera.define do we do the call conversion

Copy link
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.

add some tests (and the checks) otherwise looks good!

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.

3 participants