gemini-code-assist[bot] commented on code in PR #19626:
URL: https://github.com/apache/tvm/pull/19626#discussion_r3350321796
##########
python/tvm/relax/frontend/onnx/onnx_frontend.py:
##########
@@ -1104,6 +1104,63 @@ def _impl_v13(cls, bb, inputs, attr, params):
return relax.const(output, to_type)
if isinstance(inputs[0], relax.PrimValue):
return relax.PrimValue(inputs[0].value.astype(to_type))
+
+ try:
+ np_dst = _np.dtype(str(to_type))
+ except Exception:
+ return relax.op.astype(inputs[0], to_type)
+
+ if np_dst.kind in ("i", "u"):
+ src = inputs[0]
+ src_dtype = getattr(getattr(src, "struct_info", None), "dtype",
None) or getattr(
+ src, "dtype", None
+ )
+ if src_dtype is not None and
_relax_dtype_is_floating_point(src_dtype):
+ x_sanitized = bb.emit(
+ relax.op.where(
+ relax.op.logical_not(relax.op.isfinite(src)),
+ relax.const(0.0, src_dtype),
+ src,
+ )
+ )
+ dst_str = str(to_type)
+ if dst_str.startswith("uint"):
+ signed = False
+ bits = int(dst_str[4:])
+ elif dst_str.startswith("int"):
+ signed = True
+ bits = int(dst_str[3:])
+ else:
+ return relax.op.astype(x_sanitized, to_type)
+
+ if bits == 64:
+ return relax.op.astype(x_sanitized, to_type)
+
+ temp_dtype = "int64" if bits >= 32 else "int32"
+ t = relax.op.astype(x_sanitized, temp_dtype)
+ if bits == 32:
+ two_pow = relax.const(1 << bits, temp_dtype)
+ uw = relax.op.floor_mod(t, two_pow)
+ else:
+ mask_val = (1 << bits) - 1
+ mask = relax.const(mask_val, temp_dtype)
+ uw = relax.op.bitwise_and(t, mask)
+ if signed:
+ half = 1 << (bits - 1)
+ half_c = relax.const(half, temp_dtype)
+ if bits == 32:
+ two_pow = relax.const(1 << bits, temp_dtype)
+ else:
+ two_pow = relax.op.add(mask, relax.const(1,
temp_dtype))
+ wrapped = relax.op.where(
+ relax.op.greater_equal(uw, half_c),
+ relax.op.subtract(uw, two_pow),
+ uw,
+ )
+ else:
+ wrapped = uw
+ return relax.op.astype(wrapped, to_type)
Review Comment:
The logic for wrapping integer values can be significantly simplified and
unified.
By using `bitwise_and` with a mask of `(1 << bits) - 1` for all bit widths
less than 64 (including 32-bit), we can completely eliminate the special casing
for `bits == 32`, avoid the expensive `floor_mod` operator, and remove the
runtime addition `relax.op.add(mask, 1)`.
Since `temp_dtype` is strictly larger than `bits` (i.e., `int64` for 32-bit,
and `int32` for 8/16-bit), `bitwise_and` with the mask correctly handles both
positive and negative values in two's complement representation.
```python
mask_val = (1 << bits) - 1
mask = relax.const(mask_val, temp_dtype)
uw = relax.op.bitwise_and(t, mask)
if signed:
half = 1 << (bits - 1)
half_c = relax.const(half, temp_dtype)
two_pow = relax.const(1 << bits, temp_dtype)
wrapped = relax.op.where(
relax.op.greater_equal(uw, half_c),
relax.op.subtract(uw, two_pow),
uw,
)
else:
wrapped = uw
return relax.op.astype(wrapped, to_type)
```
--
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]
