Complete register op from python

[xqdan]

Thanks for contributing to TVM! Please refer to guideline https://tvm.apache.org/docs/contribute/ for useful information and tips. After the pull request is submitted, please request code reviews from Reviewers by @ them in the pull request thread.

This is a following up PR based on #8002

1， Support add_type_rel() and others in python
2， Make register op utilities exposed under relay.op

@zackcquic @tqchen

[xqdan]

also cc @FrozenGene please review

[zackcquic]

@xqdan Thanks a lot for this PR.

I'd really like to complete registering ops from python.
Some discussions are in RFC 9903.

The major problem (why I didn't complete it at once) is that the TypeRelation currently is pure C++ with very limited FFI.
Also, need to think about how to use TypeReporter while InferType from python?

[xqdan]

@xqdan Thanks a lot for this PR.

I'd really like to complete registering ops from python.
Some discussions are in RFC 9903.

The major problem (why I didn't complete it at once) is that the TypeRelation currently is pure C++ with very limited FFI.
Also, need to think about how to use TypeReporter while InferType from python?

1, We may need to expose some basic TypeRelation function without TypeReporter .
2, We apply infer result by TypeReporter in ir.OpAddTypeRel, which is c++ side. we rewrite the relation function here.

[xqdan]

@zackcquic please take another look

[altanh]

thanks for the PR, being able to register ops fully in python is a great tool for prototyping. I have a few concerns for this approach though:

1. IMO, we should emphasize this python API strictly for user prototyping and not for any checked in code, since the type function exposed through FFI is not the C++ type relation (it's strictly weaker since we don't have access to the type reporter, and cannot propagate constraints to the inputs, only to the output). We should probably not use "type relation" to describe the python API but I'm not sure what other name to use.
2. Related to the previous point, the python API (at least currently) does not have access to diagnostics for reporting type errors, which is not good for user experience (indeed we're trying to port ICHECKs to diagnostics where possible).

Ideally we should expose all the machinery needed but for prototyping it might not be necessary. cc @jroesch for thoughts

[xqdan]

thanks for the PR, being able to register ops fully in python is a great tool for prototyping. I have a few concerns for this approach though:

1. IMO, we should emphasize this python API strictly for user prototyping and not for any checked in code, since the type function exposed through FFI is not the C++ type relation (it's strictly weaker since we don't have access to the type reporter, and cannot propagate constraints to the inputs, only to the output). We should probably not use "type relation" to describe the python API but I'm not sure what other name to use.

Since we still call C++ type relation funcs underhood（such as test_custom_op_infer）, I keep the name but add more comments in python api, is it ok?

2. Related to the previous point, the python API (at least currently) does not have access to diagnostics for reporting type errors, which is not good for user experience (indeed we're trying to port ICHECKs to diagnostics where possible).

Let me try add a case for python exceptions.

Ideally we should expose all the machinery needed but for prototyping it might not be necessary. cc @jroesch for thoughts

[xqdan]

@altanh @zackcquic @ZihengJiang please take a look

[ZihengJiang]

The PR overall looks good. One question is, is there any support or have you thought about to register TVMCompute via such API?

[xqdan]

The PR overall looks good. One question is, is there any support or have you thought about to register TVMCompute via such API?

Thanks for your review, @ZihengJiang

Do you mean this?

tvm/python/tvm/relay/op/op.py

Line 171 in 8131364

def register_compute(op_name, compute=None, level=10):

Since we can define custom op thru tvm compute in python, and we can register it in python, I think it's ok for now.

[ZihengJiang]

Hi @xqdan , I am thinking to use relay IR and TE in a hybrid way but that might beyond this PR. I will approve it and wait to see if @altanh and @zackcquic have any opinion for your updated changes.

[xqdan]

Hi @xqdan , I am thinking to use relay IR and TE in a hybrid way but that might beyond this PR. I will approve it and wait to see if @altanh and @zackcquic have any opinion for your updated changes.

Thanks @ZihengJiang
Do you mean define custom op thru relay ir like relay.add, relay.sub ..., and then register it?
We are trying to do this, this is an example:

def layernorm_register_func():
    data = relay.var("input", dtype="float16")    
    gamma = relay.var("gamma",  dtype="float16")
    beta = relay.var("beta", dtype="float16")    
    mean = relay.mean(data, -1, True, False)
    sub_mean = relay.subtract(data, mean)
    mean_square = relay.multiply(sub_mean, sub_mean)
    variance = relay.mean(mean_square, -1, True, False)    
    z = rel.npu.scalaradd(variance, 1e-5)
    zz = relay.sqrt(z)
    u = relay.divide(sub_mean, zz)
    v = relay.multiply(u, gamma)
    w = relay.add(v, beta)    
    return relay.Function([data, gamma, beta], w).with_attr("Primitive", 1)

[zackcquic]

Hi @xqdan,

Sorry for the late reply.
This PR is very useful.
(Actually I already use this PR to do some experiments, thanks a lot 👍)

I am wondering if you have plan to make it more like C++ API (eg export TypeReporter, etc).
I agree with @altanh, should clarify what is done and what is not done, so users could quickly catch up.
I can only think about adding comments/documents, maybe you and others may come up with better idea?

[xqdan]

Hi @xqdan,

Sorry for the late reply.
This PR is very useful.
(Actually I already use this PR to do some experiments, thanks a lot 👍)

I am wondering if you have plan to make it more like C++ API (eg export TypeReporter, etc).
I agree with @altanh, should clarify what is done and what is not done, so users could quickly catch up.
I can only think about adding comments/documents, maybe you and others may come up with better idea?

Thanks for review. @zackcquic
I think it's possible to expose TypeReporter to python side, in this PR, as your suggestion, I put more comments for now.

[zackcquic]

Thanks a lot @xqdan.

[ZihengJiang]

Merged. Thanks @xqdan