Lunderberg commented on code in PR #15842:
URL: https://github.com/apache/tvm/pull/15842#discussion_r1358513235
##########
tests/python/relax/test_transform_legalize_ops.py:
##########
@@ -282,5 +284,77 @@ def main(A: R.Tensor([16, 32]), B: R.Tensor([32, 8])) ->
R.Tensor([16, 8]):
assert err_message.startswith("To legalize R.matmul")
+emit_legalization_through_builder = tvm.testing.parameter(
+ by_dict={
+ "return_relax_expr": False,
+ "return_relax_var": True,
+ }
+)
+
+
[email protected]
+def custom_op(emit_legalization_through_builder):
+ op_name = "custom_op.matmul_bias_add"
+
+ def infer_struct_info(call: relax.Call, context):
+ activations, weight, bias = call.args
+
+ matmul_call = relax.op.matmul(activations, weight)
+ matmul_sinfo =
tvm.ir.Op.get("relax.matmul").get_attr("FInferStructInfo")(
+ matmul_call, context
+ )
+
+ matmul_var = relax.Var("dummy_var", matmul_sinfo)
+ add_call = matmul_var + bias
+ add_sinfo =
tvm.ir.Op.get("relax.add").get_attr("FInferStructInfo")(add_call, context)
+
+ return add_sinfo
+
+ def legalize(bb: relax.BlockBuilder, call: relax.Call):
Review Comment:
There's a couple of reasons I'd been thinking of, most of which fall
somewhere between future-planning and user-friendliness. (Bit of a brain dump
as follows.)
1. User-friendliness to make it easier to write legalization steps. For
example, `R.nn.rms_norm` could be written in terms of `R.std` instead of
requiring a direct lowering to a TIR implementation.
2. Future-planning for user-defined custom intrinsics. If the legalization
of these custom operators is defined in terms of standard relax operators,
`LegalizeOps` would need to recursively expand them to allow
`AnnotateTIROpPattern` to recognize the results.
3. Future-planning for partial legalization. If each operator has a
"composite_level", then we could selectively lower operators that are above
some level of complexity. This would be a generalization of the
`OpDecomposer`, to decompose any
4. Future-planning for defining the requirements of graph-level optimization
passes. If an optimization pass handles all relax operators up to some
`composite_level`, new operators could be added without impacting that
optimization pass, so long as those operators define a partial legalization
that decomposes it.
5. Centralizing the definition of each operator. With composite operators
defined in terms of lower-complexity operators, the `OpDecomposer` could be
identical to the rules used by `LegalizeOps`, avoiding duplicate operator
definitions.
6. Future-planning to minimize the need for TIR pattern recognition. For
example, `R.nn.attention` is implemented in terms of `topi.transpose` and
`topi.reshape`, and would require pattern-matching similar to
`RewriteDataflowReshape` to un-lower these back to Relax operations. If
`R.nn.attention` were instead decomposed into `R.permute_dims` and `R.reshape`,
we'd get this for free.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
