tlc-pack / relax

Apache License 2.0
193 stars 58 forks source link

[BYOC] Add matmul to relax CUTLASS BYOC #418

Closed yelite closed 1 year ago

yelite commented 1 year ago

This PR adds support of matmul, including the variants with bias and with relu/gelu, to relax cutlass BYOC. This PR also makes cutlass profiler and codegen support column major RHS matrix, which is the default layout for R.matmul.

There will be a followup PR to add batched matmul.

cc: @vinx13 @masahi @mbaret

masahi commented 1 year ago

@yelite Continuing dense vs matmul terminology issue, so far CUTLASS BYOC has only supported dense, so dense terminology is used everywhere. For example, in the code below we should rename DenseOp to MatmulOp and replace the condition with func_name.find("dense") != std::string::npos || func_name.find("matmul") != std::string::npos.

https://github.com/apache/tvm/blob/7bc41ecca26204112f266d146371d0e867ec1240/src/relay/backend/contrib/cutlass/codegen.cc#L769-L770