google / google/heir

Improve torch operator coverage (as exported by torch-mlir)

Open
#2,923 2 comments 0 reactions 0 assignees View on GitHub
torch support
Dominant language
MLIR
Stars
906
Forks
171
Avg merge
4d 12h
Merged PRs (30d)
32

Description

In a separate repository https://github.com/j2kun/heir_torch_coverage, I created a set of torch-mlir export scripts, one for each torch operator.

For Conv2d for example, the [operators/conv2d](https://github.com/j2kun/heir_torch_coverage/tree/main/operators/conv2d) contains a single-op torch model with garbage weights, exported to MLIR and annotated with secret before being run through `heir-opt` (2026-05-01 release). If it succeeds, the resulting openfhe MLIR is saved, otherwise an erorr log is saved.

I let AI decide on what "representative" set of torch operators should be included in the list (there are over 2k torch operators), so if new ones are needed I can add them.

So this issue exists to track progress on supporting these operators in HEIR. Specifically, if the relevant torch operator compiles to linalg that HEIR can handle _at all_ (even if inefficiently) then it is considered done for this issue.

I also created a new label [torch-support](https://github.com/google/heir/issues?q=is%3Aissue%20state%3Aopen%20label%3A%22torch%20support%22) for tickets related to this.

Comment https://github.com/google/heir/issues/2923#issuecomment-4391907478 has a regularly-updated table containing results of a CI that tests all the torch operators in the test suite.

You can find the failure logs by navigating to the appropriate GH action run of https://github.com/google/heir/actions/workflows/torch_mlir_coverage.yml and downloading the "failure-logs" artifact.

The torch-ci models tested are checked in at https://github.com/google/heir/tree/main/scripts/torch_ci

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.