Support 1D convolution in HEIR
This ticket is to support `linalg.conv_1d` in the HEIR compiler. This came up because Marc from Belfort was trying to lower a torch model that used 1d convolution. Together we did a deep dive on the layout system and he will contribute the implementation.
The basic idea is to reduce conv1d to Halevi-Shoup matvec with sparse diagonals, similar to the conv2d work.
The outline of supporting it:
- Add a filter-to-undiagonalized-matrix layout relation in lib/Utils/Layout, and then compose it with the matrix diagonalization relation
- Add a new function in lib/Kernel/KernelImplementation.h that constructs the arithmetic dag for the conv1d kernel, and uses `implementHaleviShoup` as a subroutine. This can be used to run the resulting kernel on test data by instantiating the template type for the ArithmeticDag with a concrete type (std::vector<int>) and using EvalVisitor in a unit test.
- Add a new RewritePattern in ConvertToCiphertextSemantics.cpp (similar to the one for [conv2d](https://github.com/google/heir/blob/d0c64c24bff045d4182103dc103a0335dad0aab8/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp#L811)) that lowers conv1d using the arithmetic dag, and guards this by the kernel attribute (KernelName::MatvecDiagonal).
- Add support for conv1d in `layout-propagation` (I think this should be easy because the ciphertext input does not change its layout after going through this kernel).
- [optional]: Implement LayoutConversionHoistableOpInterface for linalg.conv_1d, which allows layout conversions to be hoistable through this op in `layout-optimization`.
3 条评论