Prepack conv weights

[pinzhenx]

No description provided.

[jgong5]

I have a concern related but not specific to this PR w.r.t. the in-place changes of input data/weight tensors. PyTorch originally assumes the input tensors as constant and it is safe for re-entrance, e.g. called by multiple threads from JIT with fork and join. The in-place changes of these input tensors might break this assumption. Do we have the design to protect these in-place changes from multi-threaded access?

[pinzhenx]

As of right now, no. Modifying constant tensors on the fly, especially constant parameters, do have some implications we need to take care of.

[pinzhenx]

BTW, we are also implementing a standalone prepack jit pass. This PR is kind of a workaround to optimize the imperative path, so that we could achieve the same performance as the to_mkldnn(model) approach.

We may revert this design if customers are more inclined to use JIT, or if something goes wrong with it.

[jgong5]

Use cases for your consideration:

1. Shared weight by multiple in-process inference models in both imperative and JIT mode.
2. Shared data by multiple ops.
   For 1), pre-packing during model initialization should work fine. But pre-packing on the first model run needs extra care on re-entrance.
   For 2), there won't be re-entrance in the imperative path since there is just one model but maybe there is a problem if some downstream ops need blocked layout while others need plain layout. Also, we need to take care of the re-entrance in the JIT mode.

[EikanWang]

Thanks, Jiong. I discussed with Pinzhen, for case 1, we can capture module.to and then prepack the weight tensors.

[pinzhenx]

We currently adopt prepacking conv weights in module.to

Pros:

• thread-safe
• more explicit, compared to implicitly prepack at runtime

Cons:

• no input info. meaning that queried format might not be optimal. This is not manifest in conv but this could be a problem for linear weights.
• unable to re-pack conv weights if it's been reordered back to plain (maybe even not a con)

[jgong5]

no input info. meaning that queried format might not be optimal. This is not manifest in conv but this could be a problem for linear weights.

It will be a problem for conv too if we use winograd. Just FYI.

[EikanWang]

We need to check auto_dnnl here. If the user disables auto_dnnl, it should go through the original path.