forked from ccf-ai-infra/TileOPs-Metax
dfebdd432d
Closes #1246 - Extend `MaskedFillFwdKernel.SUPPORTED_DTYPES` to cover the int/bool dtypes the manifest declares for `MaskedFillScalarFwdOp` (bool, uint8, int8, int16, int32, int64, float16, bfloat16, float32). - Generalize `_clamp_to_dtype_range` and `_validate_scalar_param_repr` to match PyTorch's `Tensor.masked_fill` scalar coercion: - **Signed int dtypes**: range-check the real float against `iinfo.min/max`, then truncate toward zero (`int32.fill(1.5) -> 1`, `int8.fill(127.5)` raises because `127.5 > 127`). - **`torch.uint8`**: additionally accepts Python ints in `[-255, 0)` and wraps via `value & 0xFF` (`fill(-1) -> 255`). Float negatives stay rejected (matches PyTorch). - **fp16 / bf16 / fp32**: `+/-Inf` and `NaN` pass through untouched. `fp8_e5m2` preserves `Inf`; only `fp8_e4m3fn` (no Inf representation) saturates to `finfo` extrema. - **bool**: any truthy/falsy int or float reduces to `{0, 1}`. - Validator gains an `allow_nonfinite_float` opt-in so masked_fill accepts `Inf`/`NaN` while `elu` / `softplus` / `clamp` keep their finite-only contract. - Dispatch `bool` at the Op layer by viewing input as `uint8` and re-viewing the result as `bool`; TileLang does not vectorize bool storage. Kernel + Op changes only. Manifest entry is **byte-identical** in this PR. - [x] AC-1: `pytest tests/ops/test_special_elementwise.py tests/ops/test_elementwise_independent_fp8.py tests/ops/test_elementwise_caching_autotune.py` -> 131 passed - [x] AC-2: All 9 manifest dtypes (bool, uint8, int8, int16, int32, int64, float16, bfloat16, float32) construct, compile, run, and match `torch.Tensor.masked_fill` reference - [x] AC-3: Scalar coercion validator matches PyTorch exactly: rejects out-of-range integers, `Inf`/`NaN` into int dtypes, and uint8 float negatives; accepts uint8 int wraparound and float `Inf`/`NaN` for floating-point dtypes - [x] AC-4: Per-dtype correctness tests added to `tests/ops/test_special_elementwise.py`: - signed int min/max + uint8 negative wrap (`test_masked_fill_int_dtypes`) - bool truthy/falsy fill values via uint8 storage view (`test_masked_fill_bool`) - `+/-Inf` and `NaN` through fp16/bf16/fp32 (`test_masked_fill_float_nonfinite`) - validator rejection set (`test_masked_fill_rejects_out_of_range`) - [x] pre-commit passed `python scripts/test_node_delta.py --base upstream/testbed` ``` File Base HEAD Delta -------------------------------------------------------------- tests/ops/test_special_elementwise.py 63 81 +18 -------------------------------------------------------------- TOTAL 63 81 +18 Growth: +28.6% ``` --------- Co-authored-by: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com> |
||
|---|---|---|
| .. | ||
| __init__.py | ||
| _base.py | ||
| activations.py | ||
| alibi.py | ||
| arithmetic.py | ||
| bitwise.py | ||
| clamp.py | ||
| comparison.py | ||
| fused_gated.py | ||
| logical.py | ||
| masked_fill.py | ||
| math_unary.py | ||
| nan_to_num.py | ||
| predicates.py | ||
| prelu.py | ||
| sinusoidal.py | ||
| where.py | ||