MNIST with rolled kernels takes 7 minutes per inference
- Dominant language
- MLIR
- Stars
- 906
- Forks
- 171
- Avg merge
- 4d 12h
- Merged PRs (30d)
- 32
Description
In https://github.com/google/heir/pull/2971 I'm trying to support MNIST with rolled kernels. The code compiles, but runs slowly (7 minutes per inference).
In profiling, I found two main issues:
1. One of [the generated packing loops](https://gist.github.com/j2kun/7aa21ed702ae391e73e761cc561beb0f) uses inefficient slice copying in the middle of the loop. While this could be factored out into split-preprocessing, it can also be dramatically improved by not copying a slice and modifying the array in-place. I suspect this (and other issues with generated code being superficially slow) is partly due to our decision to keep the program in tensor (value) semantics even at code generation time. The codegen for OpenFHE and Lattigo have some awkward places where, despite tensor dialect ops having value semantics, we instead mutate tensors in-place for performance reasons. I'd like to explore potentially adding a one-shot-bufferize pipeline just before codegen, so that we can more naturally codegen memref ops with pointer semantics that avoid copying more trivially.
2. RotateNew takes up 3 minutes of runtime alone, which seems partly due to https://github.com/google/heir/issues/2635 being disabled, but I also suspect that alloc-to-inplace would not work with loops to begin with. I think this pass may also benefit from one-shot-bufferize, but I'm not entirely sure how the two would interact.
Contributor guide
Assessment
This issue has not been assessed yet.