So the main goal of this thread is to gather thoughts on where we place C2 lowering(in form of target dependent library dispatch).
From an API perspective, I think C0, C1, and C2 only differ by the replacement used for each operator. Since they each produce a
relax::Expr as output, the only difference is in the functional a
- C0: The relax operator is replaced by a call to a newly-generated
- C1: The relax operator may be replaced by a call to a newly-generated
PrimFunc, or may be replaced by a call to
- C2: The relax operator is replaced by a call to a
C0 and C1 generally will have legalization. C2 usually do not.
Wouldn’t C2 cases still have a legalization step in which
relax.op.some_operator(*args) gets replaced with
R.ExternFunc("implementation_of_some_operator")(*args)? That is, every operator that isn’t a low-level builtin gets replaced at some point. C2 operators just have a 1:1 correspondence between the relax representation and the external function that implements it.
Users can think about introducing an optimized version later for a new target when the performance of that op becomes important enough for them.
I like this approach and this categorization quite a bit. The categorization also defines what can be done externally to the operator, If I were to summarize the key aspects of it:
- Every replacement, whether “legalization” or “optimization”, replaces the call to the operator with a
- Every operator must have a legalization that is valid for every target.
- Some operators may have replacements that are applied prior to legalization.
The main difference is that in such cases these TensorIR might still be specialized (aka for sort we will need to have gpu radix sort while on cpu maybe we we can do quick sort)
In the view of “legalization” and “optimization”, I’d consider this case as two possible optimizations (radix sort if gpu, quicksort if cpu) that are independently checked prior to legalization. That is, rather than having specialized forms produced during legalization, all specialized forms are produced prior to legalization. If that happens to cover all possible cases, great, but no there’s no requirement for them to do so.