tlc-pack / relax

Apache License 2.0
193 stars 58 forks source link

[CUTLASS] Move matmul host codegen to python #442

Closed masahi closed 1 year ago

masahi commented 1 year ago

To make extending CUTLASS BYOC easier, I'm moving C++-side codegen that's responsible for cutlass host code generation to python. This is the first step, only modifying matmul codegen.

I've add a packed func contrib.cutlass.instantiate_template, that gets an attribute map from C++, instantiates a host code template, and return the instantiated code back to C++. Op-specific attribute parsing and template instantiation are both done in python, so adding a new op no longer needs modification to the C++ code.

cc @vinx13 @yelite @mbaret

masahi commented 1 year ago

Now batch matmul host codegen is also moved to python, which lets us remove big C++ boilerplate.