[Relax][PyTorch] Add support for sparse matrix multiplication by Dayuxiaoshui · Pull Request #18499 · apache/tvm

Dayuxiaoshui · 2025-11-25T04:36:22Z

This commit adds support for sparse matrix multiplication and random number generation in PyTorch frontend.

Changes:

Add _sparse_mm() method to handle sparse matrix multiplication
Add _sparse_addmm() method to handle sparse addmm operations
Register these operations in the convert_map

The fix ensures that PyTorch models containing sparse matrix operations
and random number generation can be successfully converted to TVM Relax modules.

Fixes #18476

gemini-code-assist · 2025-11-25T04:36:48Z

Summary of Changes

Hello @Dayuxiaoshui, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the PyTorch frontend for TVM Relax by introducing support for several key operations. It enables the conversion of PyTorch models that utilize sparse matrix multiplications and random number generation, as well as as_strided operations, into TVM Relax modules. This broadens the range of PyTorch models that can be successfully compiled and optimized by TVM.

Highlights

Sparse Matrix Operations: Added support for PyTorch sparse matrix multiplication (_sparse_mm) and sparse addmm operations (_sparse_addmm) within the Relax frontend, converting sparse inputs to dense for processing.
Random Number Generation: Implemented support for torch.randn by introducing a _randn method, which currently acts as a placeholder by emitting relax.op.zeros until native random number generation is available in TVM.
Tensor Conversion Utility: Introduced a new static method _convert_pytorch_tensor_to_tvm to handle the robust conversion of PyTorch tensors to TVM tensors, including the necessary conversion of sparse PyTorch tensors to dense format and a fallback mechanism for DLPack conversion failures.
Strided Tensor Support: Added support for torch.as_strided operations, specifically for view-like cases where the provided strides align with a contiguous layout, translating them to relax.op.reshape.
Operation Registration: Registered the newly added _sparse_mm.default, _sparse_addmm.default, mul, as_strided.default, and randn.default operations in the create_convert_map to enable their translation during the conversion process.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces support for sparse matrix multiplication and random number generation from PyTorch's ExportedProgram to Relax. The implementation for sparse operations (_sparse_mm, _sparse_addmm) cleverly works by converting sparse tensors to dense, which is a practical approach for enabling model conversion. For torch.randn, a placeholder of zeros is used, with clear comments about this limitation. The PR also adds a handler for as_strided for contiguous memory layouts and refactors tensor conversion logic into a new _convert_pytorch_tensor_to_tvm helper, which cleans up the code and centralizes tensor handling. Overall, the changes are well-implemented and improve the frontend's capabilities. I have one minor suggestion to simplify the code.

python/tvm/relax/frontend/torch/exported_program_translator.py

…dom number generation This commit adds support for sparse matrix multiplication and random number generation in PyTorch frontend. Changes: - Add _sparse_mm() method to handle sparse matrix multiplication - Add _sparse_addmm() method to handle sparse addmm operations - Add _randn() method to handle torch.randn random number generation - Register these operations in the convert_map The fix ensures that PyTorch models containing sparse matrix operations and random number generation can be successfully converted to TVM Relax modules. Fixes #18476

Resolved merge conflict by keeping our sparse matrix operations support: - _sparse_mm() method - _sparse_addmm() method - _randn() method - Registration in convert_map Fixes #18476

Dayuxiaoshui · 2025-11-25T14:56:33Z

@tlopex Request Review

tlopex · 2025-11-26T05:15:42Z

@Dayuxiaoshui Could you add tests for those new ops?

Dayuxiaoshui · 2025-11-26T06:07:09Z

@tlopex No problem, I'll add it.

…number generation This commit adds support for sparse matrix multiplication, sparse addmm, and random number generation in PyTorch frontend. Changes: - Add _sparse_mm() method to handle sparse matrix multiplication - Add _sparse_addmm() method to handle sparse addmm operations - Add _randn() method to handle torch.randn random number generation - Register these operations in the convert_map - Add comprehensive tests for all three new operations The implementation converts sparse tensors to dense format before matrix operations, which enables model conversion for PyTorch models containing sparse computations. Fixes #18476

Dayuxiaoshui · 2025-11-26T06:32:37Z

@tlopex By the way, checked and fixed the format of the test file

…mm, and _randn

…n output PyTorch's run_decompositions() decomposes _sparse_mm.default into full.default + _sparse_addmm.default (with beta=0). This commit updates the expected IR in test_sparse_mm to include the R.full operation that is generated by the decomposition.

Dayuxiaoshui · 2025-11-27T08:33:29Z

@tlopex

tlopex · 2025-11-27T18:32:36Z

Overall LGTM. Please remove the randn func. Besides, the added trailing commas in the TVMScript blocks introduce formatting inconsistencies with existing tests. Please reformat and revert them, keep only the logical changes. Thanks!

Add test_sparse_addmm and test_sparse_mm to verify sparse tensor operations. Also remove randn support as requested.

tlopex

LGTM. Thank you!

tlopex · 2025-11-27T18:20:47Z

python/tvm/relax/frontend/torch/exported_program_translator.py

+
+        # Create a tensor filled with zeros (as placeholder)
+        # In practice, this should use a random number generator
+        # For now, we use zeros as a workaround since TVM doesn't have built-in randn


I think we shouldn't support a valid function until it is available. So temporarily I think we don't need it.

No problem, I agree with this point of view.

gemini-code-assist bot reviewed Nov 25, 2025

View reviewed changes

python/tvm/relax/frontend/torch/exported_program_translator.py Outdated Show resolved Hide resolved

Dayuxiaoshui force-pushed the main branch from 485b56d to 8745c0e Compare November 25, 2025 04:40

Dayuxiaoshui force-pushed the main branch from 8745c0e to 0f667ad Compare November 25, 2025 04:50

Merge commit 4be951d and resolve conflicts

5182b0f

Resolved merge conflict by keeping our sparse matrix operations support: - _sparse_mm() method - _sparse_addmm() method - _randn() method - Registration in convert_map Fixes #18476

Dayuxiaoshui added 2 commits November 26, 2025 06:28

Merge branch 'main' of github.com:Dayuxiaoshui/tvm

e8d4f05

Dayuxiaoshui added 4 commits November 26, 2025 06:54

[CI] Fix black formatting in test_frontend_from_exported_program.py

1f931ca

[CI] Fix duplicate key warnings in create_convert_map

9e10254

[Fix] Remove duplicate method definitions for _sparse_mm, _sparse_add…

631519d

…mm, and _randn

[Relax][PyTorch] Add tests for sparse matrix operations

804718d

Add test_sparse_addmm and test_sparse_mm to verify sparse tensor operations. Also remove randn support as requested.

Dayuxiaoshui force-pushed the main branch from 2f5dfa0 to 804718d Compare November 28, 2025 04:20

[CI] Fix missing newline at end of test file

d47f1f6

tlopex approved these changes Nov 28, 2025

View reviewed changes

tlopex merged commit 25a37e7 into apache:main Nov 28, 2025
10 checks passed

tlopex changed the title ~~[Relax][PyTorch] Add support for sparse matrix multiplication and random number generation~~ [Relax][PyTorch] Add support for sparse matrix multiplication Nov 28, 2025

ysh329 mentioned this pull request Feb 1, 2026

[Release] v0.23.0 Release Candidate Notes #18701

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Relax][PyTorch] Add support for sparse matrix multiplication#18499

[Relax][PyTorch] Add support for sparse matrix multiplication#18499
tlopex merged 10 commits intoapache:mainfrom
Dayuxiaoshui:main

Dayuxiaoshui commented Nov 25, 2025 •

edited by tlopex

Loading

Uh oh!

gemini-code-assist bot commented Nov 25, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Dayuxiaoshui commented Nov 25, 2025 •

edited

Loading

Uh oh!

tlopex commented Nov 26, 2025

Uh oh!

Dayuxiaoshui commented Nov 26, 2025

Uh oh!

Dayuxiaoshui commented Nov 26, 2025

Uh oh!

Dayuxiaoshui commented Nov 27, 2025

Uh oh!

tlopex commented Nov 27, 2025

Uh oh!

tlopex left a comment

Uh oh!

tlopex Nov 27, 2025

Uh oh!

Dayuxiaoshui Nov 29, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Dayuxiaoshui commented Nov 25, 2025 • edited by tlopex Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist bot commented Nov 25, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Dayuxiaoshui commented Nov 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tlopex commented Nov 26, 2025

Uh oh!

Dayuxiaoshui commented Nov 26, 2025

Uh oh!

Dayuxiaoshui commented Nov 26, 2025

Uh oh!

Dayuxiaoshui commented Nov 27, 2025

Uh oh!

tlopex commented Nov 27, 2025

Uh oh!

tlopex left a comment

Choose a reason for hiding this comment

Uh oh!

tlopex Nov 27, 2025

Choose a reason for hiding this comment

Uh oh!

Dayuxiaoshui Nov 29, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Dayuxiaoshui commented Nov 25, 2025 •

edited by tlopex

Loading

Dayuxiaoshui commented Nov 25, 2025 •

edited

Loading