gemini-code-assist[bot] commented on code in PR #18426:
URL: https://github.com/apache/tvm/pull/18426#discussion_r2502618527
##########
python/tvm/topi/math.py:
##########
@@ -481,7 +481,9 @@ def log2(x):
y : tvm.te.Tensor
The result.
"""
- return te.compute(x.shape, lambda *i: te.log2(x(*i)))
+ if x.dtype.startswith("int"):
+ x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
Review Comment:
Similar to the `log` function, the intermediate `te.compute` for casting
here should be tagged as `elemwise` to enable operator fusion. This ensures
that the cast can be inlined by the scheduler.
```suggestion
x = te.compute(x.shape, lambda *i: x(*i).astype("float32"),
tag=tag.ELEMWISE)
```
##########
python/tvm/topi/math.py:
##########
@@ -464,10 +463,11 @@ def log(x):
y : tvm.te.Tensor
The result.
"""
- return te.compute(x.shape, lambda *i: te.log(x(*i)))
+ if x.dtype.startswith("int"):
+ x = te.compute(x.shape, lambda *i: x(*i).astype("float32"))
Review Comment:
The intermediate `te.compute` for casting should be tagged as `elemwise` to
allow for operator fusion. Without this tag, the cast operation might not be
inlined, potentially leading to suboptimal performance.
```suggestion
x = te.compute(x.shape, lambda *i: x(*i).astype("float32"),
tag=tag.ELEMWISE)
```
--
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]
