How to build PyTorch+Caffe2 with BLAS library, instead of Eigen


TLDR: Is there a flag or some configuration option to make FULL_CAFFE2=1 python3 install link to a custom BLAS library for optimizations, instead of using Eigen?

For background, we have built a novel data-parallel accelerator, and have compiled an optimized BLAS library targeting this architecture. We would like to link PyTorch+Caffe2 with this library, such that as many NN routines as possible fall back to calling our optimized BLAS routines.

As I understand it, the default CPU-only Pytorch build leverages Eigen for generating optimized BLAS kernels. Is there some way to tear out Eigen, and replace those kernel calls with my own BLAS routines? In Caffe1 this is simple, by just changing the value of BLAS_DIR in the Makefile. In Pytorch+Caffe2 I have looked at the script, but I do not see a build flag which might be relevant.