google / google/heir

MNIST with rolled kernels takes 7 minutes per inference

Open
#2,975 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.