gemini-code-assist[bot] commented on code in PR #19897:
URL: https://github.com/apache/tvm/pull/19897#discussion_r3487183182
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -546,7 +546,15 @@ def _impl_v7(cls, bb, inputs, attr, params):
if isinstance(inputs[1], relax.Constant) and
bool(_np.any(inputs[1].data.numpy() == 0)):
raise ValueError("ONNX Div with integer inputs encountered divisor
value 0.")
- return cls.base_impl(bb, inputs, attr, params)
+ if isinstance(inputs[1], relax.Constant):
+ return cls.base_impl(bb, inputs, attr, params)
+
+ rhs_nonzero = bb.normalize(relax.op.not_equal(inputs[1],
relax.op.zeros_like(inputs[1])))
+ safe_rhs = bb.normalize(
+ relax.op.where(rhs_nonzero, inputs[1],
relax.op.ones_like(inputs[1]))
+ )
+ quotient = bb.normalize(relax.op.divide(inputs[0], safe_rhs))
+ return relax.op.where(rhs_nonzero, quotient,
relax.op.zeros_like(quotient))
Review Comment:
### Efficiency & Cleanliness Improvement
Instead of allocating and filling intermediate tensors of the same shape as
`inputs[1]` or `quotient` using `zeros_like` and `ones_like`, we can use scalar
constants `relax.const(0, dtype)` and `relax.const(1, dtype)`. Since Relax
operators support broadcasting with scalar constants, this avoids unnecessary
memory allocations and kernel launches for filling these intermediate tensors,
which is especially beneficial for large or dynamic tensors.
This also simplifies the generated Relax IR by removing three intermediate
variables (`zeros_like` and `ones_like` bindings).
Note that if you apply this suggestion, you will need to update the expected
IR in `test_div_integer_dynamic_zero_divisor_ir_guard` and
`test_div_integer_dynamic_zero_divisor_broadcast_ir_guard` to use `R.const(0,
"int32")` and `R.const(1, "int32")` instead of `R.zeros_like` and `R.ones_like`.
```suggestion
dtype = inputs[1].ty.dtype
rhs_nonzero = bb.normalize(relax.op.not_equal(inputs[1],
relax.const(0, dtype)))
safe_rhs = bb.normalize(
relax.op.where(rhs_nonzero, inputs[1], relax.const(1, dtype))
)
quotient = bb.normalize(relax.op.divide(inputs[0], safe_rhs))
return relax.op.where(rhs_nonzero, quotient, relax.const(0, dtype))
```
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
