mirage-project / mirage

A multi-level tensor algebra superoptimizer
https://mirage-project.readthedocs.io/
Apache License 2.0
318 stars 18 forks source link

[Transpiler] Threadblock Matmul does not perform epilogue #82

Open jiazhihao opened 4 days ago

jiazhihao commented 4 days ago

Threadblock Matmul operators currently do not perform epilogue: https://github.com/mirage-project/mirage/blob/main/src/transpiler/transpiler_tb.cc#L715-L725

This issue can be reproduced by running the demo/demo_gqa.py example, where the exp operator after matmul is not shown in the transpiled code.

xinhaoc commented 1 day ago

it should be done at here https://github.com/mirage-project/mirage/blob/main/include/mirage/transpiler/runtime/threadblock/matmul.h#L353 I can add other element unary ops