apache / tvm

Open deep learning compiler stack for cpu, gpu and specialized accelerators
https://tvm.apache.org/
Apache License 2.0
11.81k stars 3.48k forks source link

[Relax][MetaSchedule] Support CPU weight prepack #17445

Closed Hzfengsy closed 1 month ago

Hzfengsy commented 1 month ago

This PR adds support for CPU weight prepacking. To be specific, this PR adds a new pass AttachAttrLayoutFreeBuffers to attach layout free buffers to the weight parameters, so that we can leverage MetaSchedule to optimize the prepacking process.

After the pass and tuning, we introduce a new pass SplitLayoutRewritePreproc to split the layout rewrite pass into multiple functions, so that we can lift the parameters transform pass function with existing pass.