gemini-code-assist[bot] commented on code in PR #19837:
URL: https://github.com/apache/tvm/pull/19837#discussion_r3448293228
##########
python/tvm/relax/frontend/torch/exported_program_translator.py:
##########
@@ -918,6 +918,189 @@ def _gru(self, node: fx.Node) -> relax.Var:
return output
+ def _rnn_tanh_cell_unroll(
+ self,
+ input_reshaped,
+ weight_ih,
+ weight_hh,
+ bias_ih,
+ bias_hh,
+ h_prev,
+ seq_len,
+ reverse=False,
+ ):
+ """Unroll vanilla tanh-RNN cells for a single direction."""
+ # Transpose weights for matmul: (hidden_size, in) -> (in, hidden_size)
+ weight_ih_t = self.block_builder.emit(relax.op.permute_dims(weight_ih,
axes=[1, 0]))
+ weight_hh_t = self.block_builder.emit(relax.op.permute_dims(weight_hh,
axes=[1, 0]))
+
+ outputs = []
+ time_steps = range(seq_len - 1, -1, -1) if reverse else range(seq_len)
+
+ for t in time_steps:
+ # Input at time t: (batch_size, input_size)
+ x_t = self.block_builder.emit(
+ relax.op.take(input_reshaped, relax.const(t, "int64"), axis=0,
mode="clip")
+ )
+
+ # h_t = tanh(W_ih @ x_t + b_ih + W_hh @ h_{t-1} + b_hh)
+ ih = self.block_builder.emit(relax.op.linear_algebra.matmul(x_t,
weight_ih_t))
+ hh =
self.block_builder.emit(relax.op.linear_algebra.matmul(h_prev, weight_hh_t))
+ if bias_ih is not None and bias_hh is not None:
+ ih = self.block_builder.emit(relax.op.add(ih, bias_ih))
+ hh = self.block_builder.emit(relax.op.add(hh, bias_hh))
+ h_t = self.block_builder.emit(relax.op.tanh(relax.op.add(ih, hh)))
+
+ outputs.append(h_t)
+ h_prev = h_t
Review Comment:
Since `bias_ih` and `bias_hh` are constant across all time steps, we can
pre-compute their sum `bias = bias_ih + bias_hh` outside the loop. This avoids
performing `seq_len` redundant additions and emitting unnecessary intermediate
variables in the Relax block for each time step.
```python
bias = None
if bias_ih is not None and bias_hh is not None:
bias = self.block_builder.emit(relax.op.add(bias_ih, bias_hh))
outputs = []
time_steps = range(seq_len - 1, -1, -1) if reverse else
range(seq_len)
for t in time_steps:
# Input at time t: (batch_size, input_size)
x_t = self.block_builder.emit(
relax.op.take(input_reshaped, relax.const(t, "int64"),
axis=0, mode="clip")
)
# h_t = tanh(W_ih @ x_t + W_hh @ h_{t-1} + bias)
ih = self.block_builder.emit(relax.op.linear_algebra.matmul(x_t,
weight_ih_t))
hh =
self.block_builder.emit(relax.op.linear_algebra.matmul(h_prev, weight_hh_t))
ih_hh = self.block_builder.emit(relax.op.add(ih, hh))
if bias is not None:
ih_hh = self.block_builder.emit(relax.op.add(ih_hh, bias))
h_t = self.block_builder.emit(relax.op.tanh(ih_hh))
outputs.append(h_t)
h_prev = h_t
```
--
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]
