gemini-code-assist[bot] commented on code in PR #19660:
URL: https://github.com/apache/tvm/pull/19660#discussion_r3348352476
##########
python/tvm/relax/frontend/torch/base_fx_graph_translator.py:
##########
@@ -523,6 +524,27 @@ def call_binary_op(op, lhs, rhs):
return convert
+ def _pow(self, node: fx.Node) -> relax.Var:
+ lhs, rhs = self.retrieve_args(node)
+ # torch integer pow returns an integer tensor, but relax.op.power
legalizes to
+ # TOPI power which requires floating-point inputs. Decompose an
integer base with
+ # a constant non-negative integer exponent into repeated
multiplication instead.
+ if (
+ isinstance(lhs, relax.Expr)
+ and isinstance(lhs.struct_info, relax.TensorStructInfo)
+ and "int" in lhs.struct_info.dtype
+ and isinstance(rhs, int)
+ and not isinstance(rhs, bool)
+ and rhs >= 0
+ ):
+ if rhs == 0:
+ return self.block_builder.emit(relax.op.ones_like(lhs))
+ result = lhs
+ for _ in range(rhs - 1):
+ result = self.block_builder.emit(relax.op.multiply(result,
lhs))
+ return result
Review Comment:
Decomposing integer power using linear repeated multiplication takes $O(N)$
operations, which can lead to extremely bloated IR or compiler stack overflows
for larger exponents (e.g., `x.pow(100)`).
Using **exponentiation by squaring** (binary exponentiation) reduces the
complexity to $O(\log N)$ multiplications, which is much more efficient and
robust.
```suggestion
if rhs == 0:
return self.block_builder.emit(relax.op.ones_like(lhs))
result = None
base = lhs
exponent = rhs
while exponent > 0:
if exponent % 2 == 1:
if result is None:
result = base
else:
result =
self.block_builder.emit(relax.op.multiply(result, base))
if exponent > 1:
base = self.block_builder.emit(relax.op.multiply(base,
base))
exponent //= 2
return result
```
--
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]
