tlc-pack/relax

[Tracking Issue] High-level operator

MasterJH5574 opened this issue · 3 comments

As StructInfo stably landed in recent refactoring, it is a good time now to introducing high-level tensor operators to our codebase, under the infrastructure of StructInfo.

The naming and categorization of operators refers to DataAPI specification. For operators that are not included in the DataAPI specification, we mainly refer to NumPy op API. On C++ side, we choose use lowercase to name the operators, to keep consistency with DataAPI op names and operator naming in PyTorch C++ API.

This introduction of operators mainly has two phases: the utility functions that are shared by operators, and the interface, attribute, implementation, StructInfo inference and test cases for each operator. The second phase will be split to multiple parts, according to the op categorization.

The temporary plan is as the following. Note that in O2, there is no dependency among categories of operators, so the actual upstreaming order may not completely follow with the order below.

  • O1. Key op utilities for op registration and StructInfo inference (op_common.h/.cc). #333
  • O2. Operator interface, attribute, implementation, StructInfo inference and test cases
    • O2a. Tensor arithmetic and comparison operators #333
    • O2b. Tensor statistical operators #334
    • O2c. Tensor creation operators #336
    • O2d. Tensor manipulation operators #337
    • O2e. Tensor indexing operators, datatype operators #338
    • O2f. Tensor set operators, searching operators #339
    • O2g. Neural network operators (convolution, pooling), image operators #341
    • O2h. Linear algebra operator and other neural network operators #343

Are you also planning to attach compute to each high-level op, so that they can actually execute?

Are you also planning to attach compute to each high-level op, so that they can actually execute?

@masahi We don’t touch the computation part at this moment while only focus on the ops themselves. Our plan is to have a pass which legalizes all high-level op calls in a Relax func to corresponding CallTIRs with PrimFuncs properly generated. That part can be introduced separately as followup works.

Great that all tasks are finished now. Thanks @sunggg @YuchenJin and @tqchen for reviewing!