Skip to content

Conversation

@tlopex
Copy link
Member

@tlopex tlopex commented Sep 13, 2025

This PR allows calling Python functions directly from Relax IR, where integration between Relax computations and Python/PyTorch operations can be supported.

Usage Example

@I.ir_module
class MyModule(BasePyModule):
    @I.pyfunc
    def pytorch_add(self, x, y):
        return x + y
    
    @R.function
    def compute(x: R.Tensor((5,), "float32"), y: R.Tensor((5,), "float32")) -> R.Tensor((5,), "float32"):
        result = R.call_py_func("pytorch_add", (x, y), out_sinfo=R.Tensor((5,), "float32"))
        return result

THIS PR represents the final milestone in the Relax Python integration design, and completes the full feature set for Python/Relax interoperability. The complete Relax-Python integration ecosystem has been built upon the following key design principles:

Cross-level Calls

  • Two-way interoperability: Python functions can invoke Relax/TIR/packed functions, and Relax functions can invoke Python functions via R.call_py_func
  • Seamless data conversion: DLPack enables efficient conversion between TVM Tensors and PyTorch Tensors with minimal overhead

Just-in-time (JIT) Compilation

  • Delayed compilation: TIR and Relax functions are compiled only when the IRModule is instantiated
  • Flexible integration: Allows late-stage modifications and seamless integration with the Python runtime
  • Relax VM execution: Compiled Relax functions are executed using a Relax VM created at instantiation time

Conversion between Relax and Python Functions

  • IRModule printer: Converts Relax functions into executable Python code for debugging and deployment
  • Operator mapping: High-level Relax operators (e.g., R.nn.relu) are mapped to corresponding PyTorch APIs (e.g., F.relu)
  • Multi-stage conversion: Can happen at any stage of compilation, from early Relax functions to fully lowered TIR modules
  • DLPack integration: Handles call_tir and call_dps_packed by converting PyTorch tensors to/from DLPack format

@tlopex
Copy link
Member Author

tlopex commented Sep 13, 2025

cc @tqchen @MasterJH5574

@MasterJH5574 MasterJH5574 self-assigned this Sep 17, 2025
@tlopex
Copy link
Member Author

tlopex commented Sep 19, 2025

cc @MasterJH5574

@MasterJH5574 MasterJH5574 merged commit 4041f89 into apache:main Sep 19, 2025
13 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants