forked from ccf-ai-infra/TileOPs-Metax
15 Commits
| Author | SHA1 | Message | Date |
|---|---|---|---|
Cao Ying
|
9bda1ac537
|
[Refactor][ELEMENTWISE] move strategy into kernel config; drop ctor kwarg (#1778)
Closes #1768 ## Summary - Move elementwise `strategy` selection into the kernel `config` dict; delete the ctor kwarg and its pass-through plumbing from all elementwise kernels and Ops. - fp8/bool coercions and the `register_copy` broadcast downgrade stay inside the kernel; same kernel body selected for every (op, dtype, shape) as before. - Migrate strategy tests to config-based construction; add guards asserting no elementwise Op/kernel exposes a `strategy` kwarg. - Fused-gated strategy bench now records a torch baseline and the measured kernel object. - Validator: `"strategy"` removed from `_CTOR_INFRA_PARAMS` — reintroducing the kwarg on any op now fails validation (validator suite 142 passed). ## Test plan - [x] pre-commit passed; 335 passed across the six modified test files - [x] AC-1: no elementwise Op/kernel exposes `strategy` (signature guards, +2 test nodes) - [x] AC-2: register_copy broadcast-downgrade regression preserved under config form - [x] AC-3: elementwise GPU smoke tier green (219 passed, H200) ## Benchmark NVIDIA H200, CUDA 12.8, PyTorch 2.9.1+cu128, TileLang 0.1.11. Fused-gated explicit_parallel (4096, 4096) fp16: | Op | TileOPs (ms) | torch (ms) | Speedup | BW (TB/s) | | --- | ---: | ---: | ---: | ---: | | SiluAndMul | 0.0292 | 0.1130 | 3.87× | 3.45 | | GeluAndMul | 0.0332 | 0.1213 | 3.65× | 3.03 | | GeluTanhAndMul | 0.0294 | 0.1147 | 3.90× | 3.43 | All rows meet or exceed the documented bandwidth basis (3.04/2.72/3.38 TB/s) — perf-neutral or better. ## Regression `test_register_copy_downgrades_on_broadcast` PASSED — config-form downgrade matches PyTorch under broadcast strides. --------- Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com> |
|
|
Cao Ying
|
6f21f92d8e
|
[Test][TestInfra] collapse test_pool per-dim blocks; merge signed-zero mirror pairs (#1773)
Round-2 test-structure cleanup (from #1763's deferred list). Test-only;
no behavior change.
## What
- **test_pool.py 2587 -> 1468 lines**: the six near-identical per-dim
max/avg blocks collapse into dim-parametrized harnesses
(`MaxPoolTest`/`AvgPoolTest` + one runner each). Curated PARAMS tables
kept value-identical (AST-verified vs previous revision). Genuine
per-dim deltas preserved: avg kernel-dispatch expectations,
`divisor_override` 2d/3d-only, max special-values config asymmetry, 2d's
exhaustive invalid-param set. All nine `register_compile_contract`
registrations intact.
- **test_binary_arith.py**: mirrored maximum/minimum signed-zero
(+with-nan) pairs parametrized over `(op_cls, torch_ref[, inputs])`;
per-op input orderings unchanged.
## Test node delta
```
tests/ops/test_binary_arith.py 147 147 0
tests/ops/test_pool.py 194 194 0
TOTAL 341 341 0
```
Zero growth; per-tier counts identical (test_pool: smoke 87 / full 107 /
packaging 6).
## Verification
GPU smoke tier on test_pool: 87 passed. Collection identical to base per
tier. ruff + pre-commit clean.
---------
Co-authored-by: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com>
|
|
|
Cao Ying
|
0a9bf1e40a
|
[Chore][Cleanup] repo slimming: dead code, duplicated tests/benches, compat shims, file defragmentation (#1764)
Closes #1763 ## Summary - Remove dead helpers, redundant elementwise benchmark drivers, decorative banners, and stale compatibility aliases. - Consolidate duplicated tests and private benchmark wrappers onto shared parametrized fixtures and `ManifestBenchmark`. - Merge fragmented workload, reduction-op, MHC, and normalization benchmark modules; retarget imports and manifest source paths. - Preserve canonical runtime behavior while reducing the repository by roughly 4.3k net lines. ## Test plan - [x] pre-commit passed - [x] `python scripts/validate_manifest.py --strict` passed - [x] Repository-wide pytest collection completed: 5,178 tests, 0 collection errors - [x] Touched test and benchmark modules collect cleanly --------- Co-authored-by: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com> |
|
|
Cao Ying
|
55ae09dd3e |
[Feat][Kernels] int / bool kernel coverage for elementwise_binary (#1408)
Closes #1243 ## Summary - Widen `_BINARY_FULL_DTYPES` (bool/uint8/int8/int16/int32/int64/fp16/bf16/fp32) and `_BINARY_NO_BOOL_DTYPES` (Sub) so 11 elementwise_binary ops accept every manifest-declared dtype. Public op layer matches the manifest dtype union exactly; fp8 remains kernel-internal. - Force `BinaryKernel` to the `direct` strategy when `dtype == torch.bool` (TileLang cannot lower vectorised `bool×N`). `RuntimeWarning` when a non-default caller-provided strategy is overridden. - Per-dtype correctness tests on three decoupled axes (dtype-axis × representative op, op-axis × int32, bool-axis × every bool-supporting op) per `.claude/domain-rules/testing-budget.md`. References are `torch.*`; bool ops mapped to `torch.logical_or` / `torch.logical_and`. Sentinel `test_add_bool_is_or_not_xor` pins TileLang's bool `+` lowering. - `tests/ops/test_elementwise_fp8.py` rerouted: fp8 acceptance / forward-dtype onto `DivFwdKernel`, saturation / Inf overflow onto `ExpFwdOp` — `AddFwdOp` is no longer reachable with fp8 at the public layer. - Manifest byte-identical. `status: spec-only → implemented` flip is a follow-up manifest-only PR per the trust model. ## Test plan - [x] `pytest tests/ops/test_binary_arith.py tests/ops/test_comparison.py tests/ops/test_elementwise_fp8.py` green (193 nodes) - [x] `scripts/validate_manifest.py` clean for all 11 ops - [x] pre-commit, gitleaks, actionlint, ci-gate green --------- Co-authored-by: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com> |
|
|
Cao Ying
|
731744d0d5 |
[Feat][Kernels] rounding-mode div kernel path for trunc/floor (#1410)
Closes #1245 ## Summary Adds `rounding_mode='trunc' | 'floor'` support to `DivFwdOp`. Previously only `rounding_mode=None` (true division) worked; non-None raised `NotImplementedError`. - `tileops/kernels/elementwise.py`: new `DivTruncFwdKernel` (fp32 promotion → divide → `T.trunc` → cast back). `'floor'` reuses existing `FloorDivideFwdKernel` (byte-identical to what an early draft `DivFloorFwdKernel` would have been; that draft class was removed during review). - `tileops/ops/elementwise/arithmetic.py`: `_DIV_KERNEL_BY_ROUNDING_MODE` dispatch (`None → DivFwdKernel`, `'trunc' → DivTruncFwdKernel`, `'floor' → FloorDivideFwdKernel`). Forward-time `NotImplementedError` guard removed; invalid modes raise `ValueError` at construction. - `tests/ops/test_binary_arith.py`: `test_div_rounding_mode_eager` (6 cells: trunc/floor × fp16/bf16/fp32, mixed-sign quotients vs `torch.div`) + `test_div_rounding_mode_dispatch` (CPU-only dispatch table + invalid-mode rejection). - `tests/ops/test_elementwise_compile.py`: `test_div_rounding_mode_compile` (6 cells under `torch.compile(fullgraph=True)`, same oracle). ## Test plan - [x] `pytest tests/ops/test_binary_arith.py tests/ops/test_elementwise_compile.py -m smoke` — 89 passed. - [x] Manual parity vs `torch.div(..., rounding_mode=...)` on fp16/bf16/fp32; tolerances follow the existing `FloorDivideFwdKernel` convention (`atol=1.0, rtol=0` for half precision rounding-mode kernels, `1e-5` for fp32). - [x] `scripts/validate_manifest.py` clean on `DivFwdOp`. --------- Co-authored-by: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com> |
|
|
Cao Ying
|
0451c058a9
|
[Chore][Release] merge testbed into main: op-family alignment + new ops (#1380)
## Summary Merge 27 commits accumulated on `testbed` into `main`. Covers manifest-spec alignment for the elementwise (unary/binary/multi-input), reduction, and normalization families; new ops (LerpTensorFwd, Reciprocal, fused tanh-GELU, InstanceNormFwdOpNoAffine, GroupNorm/InstanceNorm no-affine variants); shared dispatch/refactor cleanup; and supporting manifest DSL (`promote_int_to_float`), tests, and process docs. ## Highlights - **Family alignment to manifest spec**: elementwise_unary_activation (#1211), elementwise_binary (#1222), elementwise_multi_input (#1229), reduction (#1235), normalization (#1240). - **New ops / variants**: LerpTensorFwd (#1264), Reciprocal + `promote_int_to_float` DSL (#1252), fused tanh-GELU (#1262), GroupNorm/InstanceNorm no-affine (#1274, #1284), InstanceNormFwdOpNoAffine running-stats (#1375). - **Refactors**: activation shared bases (#1230), dispatch_kernel routing (#1260, #1368), normalization hotfix cleanup (#1256), elementwise_binary torch-fallback removal (#1248). - **Tests / docs**: manifest-driven L1 parity for elementwise_binary (#1370), drop manifest-mirror tests (#1374), FLIP_STATUS carve-out (#1258), spec-only -> implemented promotion rule (#1271). ## Test plan - [ ] CI green on the PR branch - [ ] \`make test\` passes locally - [ ] Manifest validator and strict-parity gates pass --------- Co-authored-by: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com> |
|
|
Cao Ying
|
c39d65d31e
|
[Refactor][Naming] Add Fwd direction suffix to all elementwise ops and kernels (#978)
## Summary
Rename all elementwise op and kernel classes to include the `Fwd`
direction suffix, aligning with the naming convention used by
non-elementwise ops (e.g., `RMSNormFwdOp`, `GemmFwdKernel`). This makes
the codebase consistent and prepares for future backward-pass
implementations.
- 66 Op classes renamed: `{Name}Op` -> `{Name}FwdOp`
- 66 Kernel classes renamed: `{Name}Kernel` -> `{Name}FwdKernel`
- All imports updated across tests, benchmarks, and internal references
Closes #977
## Test plan
- [x] **AC-1**: All elementwise ops follow `{Name}FwdOp` pattern — AST
scan of `tileops/ops/elementwise.py` found 66 `*FwdOp` classes and 0
legacy elementwise concrete `*Op` classes.
- [x] **AC-2**: All elementwise kernels follow `{Name}FwdKernel` pattern
— AST scan of `tileops/kernels/elementwise.py` found 66 `*FwdKernel`
classes and 0 legacy elementwise concrete `*Kernel` classes.
- [x] **AC-3**: Manifest keys match class names — `ops_manifest.yaml`
contains no legacy elementwise op/kernel class-name references.
- [x] **AC-4**: All tests pass — `python -m pytest -q tests` => 2242
passed, 22 skipped; `pre-commit run --all-files` passed.
## Follow-up
No follow-up issues or suggestions.
---------
Co-authored-by: Ibuki 🍃 — a wind born from Claude Opus <Ibuki-wind@users.noreply.github.com>
|
|
|
ChongLi
|
93340add26
|
[Test][TESTING] Add missing per-dtype smoke coverage (#952)
## Summary - add representative `smoke` cases for each supported dtype in the issue-listed test fixtures - keep the smoke cases at the front of each parametrized list and preserve `tune=False` - update the current-file equivalents for the issue paths, including `tests/ops/test_engram.py` and `tests/ops/test_pool.py` Closes #945 ## Test plan - `python -m pytest tests --collect-only -q` - `python -m pytest tests/ops/test_activation.py tests/ops/test_fused_gated.py tests/ops/test_binary_arith.py tests/ops/test_engram.py tests/ops/test_pool.py --collect-only -q` - `python scripts/test_node_delta.py tests/ops/test_activation.py tests/ops/test_fused_gated.py tests/ops/test_binary_arith.py tests/ops/test_engram.py tests/ops/test_pool.py` ## Regression - smoke-only relabeling stays scoped to representative dtype-dispatch coverage - all added smoke cases keep `tune=False` - collected `smoke` cases grow from `1072` on `main` to `1597` on this branch (`+525`) - total collected test nodes grow from `2383` on `main` to `2410` on this branch (`+27`, `+1.1%`) ## Additional context - the issue references `tests/ops/test_engram_fwd.py` and `tests/ops/test_avg_pool1d.py`; in the current workspace those correspond to `tests/ops/test_engram.py` and `tests/ops/test_pool.py` ## Follow-up No follow-up issues. Suggestion: `WelfordNonAlignedMultiDimFixture` `flat63_fp16` was promoted to smoke but adds shape coverage not dtype coverage — move back to `full`, keep `flat63_bf16` as the bf16 smoke case. |
|
|
Cao Ying
|
6e0b507b42
|
[Refactor][Workloads] Align workloads/ops/ naming with tileops/ops/ layout (#939)
## Summary
Align `workloads/ops/` file naming and directory structure 1:1 with the
post-#928 `tileops/ops/` layout. Pure file-move/rename refactor — no
workload logic changes.
- Drop `_fwd` suffix from mamba, attention, deltanet, and gated_deltanet
workloads
- Rename `mean_pooling_ops.py` to `mean_pooling.py`
- Move attention workloads into `workloads/ops/attention/` subpackage
- Update all imports across the repo
Closes #931
## Test plan
- [x] **AC-1**: All `from workloads.ops.<name>` imports updated for
moved/renamed files — verified by importing all 20 renamed/moved
workload classes and grep confirming zero old import paths remain
- [x] **AC-2**: `python -c "from workloads.ops.attention import ..."`
resolves for all 15 moved attention workload classes
- [x] **AC-3**: Full test suite passes — 2361 passed, 22 skipped, 0
failed (233.76s at commit
|
|
|
Cao Ying
|
978615d1b3
|
[Refactor][Benchmark] Detach benchmarks from tests via workloads layer (#787)
## Summary Repo-wide refactor to fully detach `benchmarks/` from `tests/` by introducing a shared `workloads/` layer. Resolves #783. Case study documented in #789. ### Problem `benchmarks/` had 81 `tests.*` imports across 67 files, plus 42 `self.test.*` accesses. Benchmark code broke when test-only code changed — the wrong dependency direction. ### Solution Introduce `workloads/` as a neutral shared layer owning only workload parameters, `gen_inputs()`, and fixture metadata. Tests and benchmarks both import from `workloads/` but never from each other. ### Architecture ``` workloads/base.py → WorkloadBase (gen_inputs only), FixtureMeta, FixtureBase workloads/ops/*.py → 55 concrete workload classes (params + gen_inputs) tests/test_base.py → TestBase(WorkloadBase) + @abstractmethod ref_program + check() tests/ops/*.py → Test classes inherit (Workload, TestBase), define ref_program locally benchmarks/benchmark.py → BenchmarkBase(workload: WorkloadBase), self.workload benchmarks/ops/*.py → Import from workloads.ops/*, define baselines locally ``` ### Trust model - `workloads/` contains NO reference implementations, NO correctness logic, NO assertion/tolerance code - `ref_program()` is `@abstractmethod` on `TestBase`, defined independently in each test class - Benchmark baselines are independent local copies — no shared oracle surface between tests and benchmarks - `workloads/` ships in the wheel with lazy `pytest` import (no undeclared runtime dependency) ## Changes **189 files changed** (+5329 / -3732) | Category | Count | Description | |----------|-------|-------------| | `workloads/` (new) | 58 files | `WorkloadBase`, `FixtureMeta`, `FixtureBase`, 55 workload classes | | `benchmarks/` | 73 files | All imports migrated from `tests.*` → `workloads.*`; `self.test` → `self.workload`; local baseline copies | | `tests/` | 58 files | Workload params extracted to `workloads/`; `ref_program` kept local; `TestBase` enforces abstract contract | | `pyproject.toml` | 1 file | `workloads` included in wheel package discovery | **Zero changes** to `tileops/ops/`, `tileops/kernels/`, or `tileops/ops_manifest.yaml`. ## Acceptance criteria - [x] `rg -n "from tests\.|import tests\." benchmarks` → no matches - [x] `rg -n "self\.test\." benchmarks` → no matches - [x] `benchmarks/benchmark.py` does not import or reference `TestBase` - [x] `BenchmarkBase` stores `self.workload`, not `self.test` - [x] `workloads/base.py` exports `WorkloadBase`, `FixtureMeta`, `FixtureBase` - [x] `tests/test_base.py` no longer defines `FixtureMeta` or `FixtureBase` - [x] `workloads/` contains no correctness-only methods or assertion/tolerance logic - [x] No shared oracle surface — reference functions duplicated independently - [x] Op implementation files identical to upstream/main (pure refactor) - [x] Manifest unchanged from upstream/main - [x] Test coverage preserved (fp32/fp16/bf16, 1D-4D, all original params) - [x] Representative benchmark smoke runs produce numeric output - [x] Representative correctness smoke tests pass ## Test plan - [x] Structural audit: zero `tests.*` imports in `benchmarks/` - [x] Structural audit: zero `self.test` in `benchmarks/` - [x] `pytest --collect-only benchmarks` — 1004 tests collected, 0 import errors - [x] `pytest tests/ops/test_softmax.py -m smoke` — 15 passed - [x] `pytest tests/ops/test_activation.py -m smoke` — 20 passed - [x] `pytest benchmarks/ops/bench_softmax.py -m smoke` — passed with numeric output - [x] `pytest benchmarks/ops/bench_activation.py -m smoke` — passed with numeric output - [x] Wheel build: `workloads` importable without pytest installed 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
81c0eafb99
|
[Perf][Elementwise] Optimize binary max/min kernels to close bandwidth gap with PyTorch (#539)
## Summary
- Add `register_copy` strategy to `BinaryKernel` with vectorized 128-bit
fragment loads (`uint4`) for same-shape inputs, achieving 2-3x bandwidth
improvement on complex op_funcs (maximum/minimum)
- Simplify `maximum`/`minimum` `op_func` from 5-node comparison chain
with copysign to `T.max`/`T.min` + `isnan` guards, reducing IR node
count
- Add `autotune_configs` to `BinaryKernel` with 9 configurations
(threads in {128, 256, 512} x num_per_thread in {2, 4, 8}) for
fp16/bf16/fp32 and 6 for fp8
- Cache compiled kernel to eliminate JIT lookup overhead on each
`forward()` call
Closes #512
## Benchmark
### Environment
- **GPU**: NVIDIA H200
- **PyTorch**: 2.9.1+cu128
- **CUDA**: 12.8
- **dtype**: float16
- **Measurement**: 1000 iterations, CUDA event timing, 50-iteration
warmup
### Before (main branch)
| op | shape | TileOPs (TB/s) | PyTorch (TB/s) | ratio |
|---|---|---|---|---|
| maximum | (1024, 4096) | 1.18 | 1.83 | 65% |
| maximum | (1024, 10240) | 1.32 | 2.79 | 47% |
| maximum | (1024, 20480) | 1.42 | 3.37 | 42% |
| minimum | (1024, 4096) | 1.18 | 1.83 | 65% |
| minimum | (1024, 10240) | 1.31 | 2.76 | 48% |
| minimum | (1024, 20480) | 1.42 | 3.32 | 43% |
### After (this PR, independent H200 re-measurement by Reviewer)
| op | shape | TileOPs (TB/s) | PyTorch (TB/s) | ratio |
|---|---|---|---|---|
| maximum | (1024, 4096) | 2.67 | 1.83 | **146%** |
| maximum | (1024, 10240) | 3.36 | 2.77 | **121%** |
| maximum | (1024, 20480) | 3.79 | 3.35 | **113%** |
| minimum | (1024, 4096) | 2.66 | 1.90 | **140%** |
| minimum | (1024, 10240) | 3.36 | 2.72 | **123%** |
| minimum | (1024, 20480) | 3.79 | 3.35 | **113%** |
All target shapes exceed the 80% acceptance threshold, with large shapes
exceeding PyTorch bandwidth.
### Kernel-level (bypassing Op layer Python overhead)
| op | shape | TileOPs (TB/s) | PyTorch (TB/s) | ratio |
|---|---|---|---|---|
| maximum | (1024, 4096) | 2.64 | 4.10 | 64% |
| maximum | (1024, 10240) | 3.55 | 3.52 | **101%** |
| maximum | (1024, 20480) | 3.84 | 3.82 | **101%** |
## Test plan
- [x] 123 unit tests pass (`tests/ops/test_binary_arith.py` +
`tests/test_elementwise_strategy_bench.py`)
- [x] NaN propagation tests pass: `test_maximum_nan_propagation`,
`test_minimum_nan_propagation` (16 tests)
- [x] Signed-zero semantics tests pass: `test_maximum_signed_zero`,
`test_minimum_signed_zero` (16 tests)
- [x] Autotune config validation:
`test_binary_kernel_has_autotune_configs`,
`test_binary_kernel_autotune_configs_distinct`
- [x] 6 benchmark runs pass on H200
## Acceptance Criteria
| AC | Description | Status |
|---|---|---|
| AC-1 | Modified files pass unit tests | Pass (123/123) |
| AC-2 | maximum/minimum >= 80% PyTorch BW on target shapes | Pass
(113-146%) |
| AC-3 | Before/after benchmark comparison in PR | Pass (see above) |
| AC-4 | No regression in signed-zero/NaN tests | Pass (16/16) |
| AC-5 | autotune_configs with >= 3 configs | Pass (9 configs for
fp16/bf16/fp32) |
🤖 Generated with [Claude Code](https://claude.com/claude-code)
---------
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
|
|
|
Cao Ying
|
d9b83d0456
|
[Fix][Elementwise] Correct maximum/minimum signed-zero tie-break (#508)
Closes #469 ## Summary Fix `maximum` / `minimum` elementwise ops to match PyTorch / IEEE 754 signed-zero semantics: - `maximum(+0.0, -0.0)` now correctly returns `+0.0` (was returning `-0.0`) - `minimum(-0.0, +0.0)` now correctly returns `-0.0` (was returning `+0.0`) The fix adds a tie-break path using `copysign` comparison that activates only when `a == b`, preserving the existing NaN-propagation behavior from PR #463. **Dtype support matrix:** | dtype | supported | |-------|-----------| | float16 | yes | | bfloat16 | yes | | float32 | yes | ## Structural Compliance All checks passed. ## Test plan - [x] `maximum(+0.0, -0.0)` returns `+0.0` for fp16, bf16, fp32 - [x] `minimum(-0.0, +0.0)` returns `-0.0` for fp16, bf16, fp32 - [x] NaN propagation behavior from PR #463 is preserved (not regressed) - [x] 12 new signed-zero regression tests added and passing - [x] No performance regression — tie-break path only activates on `a == b` ## Benchmark **Environment**: H200 GPU, CUDA 12.8, PyTorch 2.9.1+cu128, TileLang 0.1.8 Before/after comparison (`main` vs this branch, fp16): | Op | Shape | Before Latency (ms) | Before BW (TB/s) | After Latency (ms) | After BW (TB/s) | Delta | |----|-------|---------------------|-------------------|---------------------|-------------------|-------| | maximum | (1024, 4096) | 0.02 | 1.18 | 0.02 | 1.18 | 0% | | maximum | (1024, 10240) | 0.05 | 1.33 | 0.05 | 1.32 | -0.8% | | maximum | (1024, 20480) | 0.09 | 1.43 | 0.09 | 1.42 | -0.7% | | minimum | (1024, 4096) | 0.02 | 1.19 | 0.02 | 1.18 | -0.8% | | minimum | (1024, 10240) | 0.05 | 1.32 | 0.05 | 1.31 | -0.8% | | minimum | (1024, 20480) | 0.09 | 1.43 | 0.09 | 1.42 | -0.7% | **Takeaways:** - No measurable regression — deltas are within noise margin (<1%) - The `copysign`-based tie-break only activates when `a == b`, which is rare in real workloads - Existing gap vs PyTorch baseline (42–65% efficiency) is a pre-existing issue tracked in #512 **Benchmark command:** ```bash PYTHONPATH="$PWD" python -m pytest benchmarks/ops/bench_binary_elementwise.py -v -k "maximum or minimum" ``` ## Regression - 101 tests pass (12 signed-zero + 10 NaN propagation + existing suite) - Lint clean (`pre-commit run --all-files` passes) - Signed-zero fix does not alter NaN-propagation codepath — verified by `test_maximum_signed_zero_with_nan` and `test_minimum_signed_zero_with_nan` --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
783fbe5f7c
|
[Feat][Elementwise] Register elementwise ops for torch.compile (#482)
## Summary Register all elementwise ops for `torch.compile` via `@torch.library.custom_op` + `register_fake`. Three factory functions handle 55 template ops; independent ops use hand-written registration. After this change, `Kernel.forward` switches from `self.kernel(x)` to `type(self)._wrapped(...)`. Closes #440 ## Changes - **`tileops/ops/elementwise.py`** — Add `_register_unary_custom_op`, `_register_binary_custom_op`, `_register_fused_gated_custom_op` factory functions; `_wrapped` binding on all kernel classes at module load time; forward path switch to use `_wrapped`; `_OP_REGISTRY` as `WeakValueDictionary` to avoid leaks - **`tileops/kernels/elementwise.py`** — Fix `RemainderKernel` and `FloorDivideKernel` fp16 precision: promote division to fp32 before `T.floor` to avoid rounding near integer quotients - **`tests/test_elementwise_compile.py`** — 60 torch.compile correctness tests covering every registered op (unary, binary, comparison, logical, bitwise, fused gated) - **`benchmarks/ops/bench_compile_elementwise.py`** — New benchmark: eager vs torch.compile latency and bandwidth ## Test plan - [x] **AC-1**: All 55 elementwise ops work under `torch.compile` — full-matrix smoke with `fullgraph=True`; 60 pytest cases passed - [x] **AC-2**: `register_fake` returns correct shape and dtype (including `torch.bool` for comparison) — verified via dtype/shape assertions - [x] **AC-3**: Eager mode L1 tests still pass — 196 passed across binary_arith, unary_math, fused_gated, special_elementwise, bitwise suites - [x] **AC-4**: Import time measured and acceptable — average 160.60ms (max 162.43ms), well below 500ms constraint ## Benchmark Environment: Torch 2.9.1+cu128, CUDA 12.8, NVIDIA H200, Driver 575.57.08 ### Unary Ops (fp16, latency us + bandwidth TB/s) | Op | N | Eager (us) | Compile (us) | Overhead | Eager BW (TB/s) | Compile BW (TB/s) | |---|---|---|---|---|---|---| | relu | 262,144 | 2.3 | 2.1 | 0.92x | 0.47 | 0.51 | | relu | 1,048,576 | 3.1 | 3.0 | 0.95x | 1.34 | 1.41 | | relu | 4,000,000 | 6.4 | 6.2 | 0.97x | 2.50 | 2.58 | | exp | 262,144 | 2.3 | 2.1 | 0.92x | 0.46 | 0.50 | | exp | 1,048,576 | 3.2 | 3.0 | 0.95x | 1.33 | 1.40 | | exp | 4,000,000 | 6.4 | 6.2 | 0.97x | 2.48 | 2.57 | | gelu | 262,144 | 2.4 | 2.2 | 0.93x | 0.44 | 0.47 | | gelu | 1,048,576 | 3.5 | 3.4 | 0.96x | 1.20 | 1.25 | | gelu | 4,000,000 | 8.1 | 8.0 | 0.99x | 1.99 | 2.01 | | silu | 262,144 | 2.4 | 2.2 | 0.92x | 0.44 | 0.48 | | silu | 1,048,576 | 3.4 | 3.2 | 0.95x | 1.23 | 1.29 | | silu | 4,000,000 | 7.9 | 7.7 | 0.97x | 2.01 | 2.07 | | abs | 262,144 | 2.2 | 2.1 | 0.92x | 0.47 | 0.51 | | abs | 1,048,576 | 3.1 | 3.0 | 0.94x | 1.33 | 1.41 | | abs | 4,000,000 | 6.4 | 6.2 | 0.97x | 2.50 | 2.57 | | sqrt | 262,144 | 2.4 | 2.3 | 0.94x | 0.43 | 0.46 | | sqrt | 1,048,576 | 3.2 | 3.0 | 0.95x | 1.31 | 1.38 | | sqrt | 4,000,000 | 6.5 | 6.2 | 0.97x | 2.48 | 2.56 | | log | 262,144 | 2.4 | 2.2 | 0.92x | 0.45 | 0.48 | | log | 1,048,576 | 3.4 | 3.2 | 0.94x | 1.24 | 1.31 | | log | 4,000,000 | 7.9 | 7.7 | 0.98x | 2.03 | 2.06 | | tanh | 262,144 | 2.3 | 2.1 | 0.92x | 0.45 | 0.49 | | tanh | 1,048,576 | 3.2 | 3.0 | 0.95x | 1.32 | 1.40 | | tanh | 4,000,000 | 6.5 | 6.3 | 0.97x | 2.44 | 2.53 | ### Binary Ops (fp16, latency us + bandwidth TB/s) | Op | Shape (MxN) | Eager (us) | Compile (us) | Overhead | Eager BW (TB/s) | Compile BW (TB/s) | |---|---|---|---|---|---|---| | add | 1024x4096 | 8.9 | 8.7 | 0.98x | 2.84 | 2.90 | | add | 1024x10240 | 18.4 | 18.2 | 0.99x | 3.41 | 3.46 | | add | 1024x20480 | 33.1 | 32.9 | 0.99x | 3.80 | 3.82 | | mul | 1024x4096 | 8.9 | 8.6 | 0.98x | 2.84 | 2.91 | | mul | 1024x10240 | 18.4 | 18.2 | 0.99x | 3.42 | 3.46 | | mul | 1024x20480 | 33.1 | 32.9 | 0.99x | 3.80 | 3.83 | | sub | 1024x4096 | 8.8 | 8.7 | 0.98x | 2.84 | 2.91 | | sub | 1024x10240 | 18.4 | 18.2 | 0.99x | 3.42 | 3.46 | | sub | 1024x20480 | 33.1 | 32.9 | 0.99x | 3.80 | 3.83 | | div | 1024x4096 | 9.1 | 8.8 | 0.98x | 2.77 | 2.84 | | div | 1024x10240 | 18.2 | 18.0 | 0.99x | 3.45 | 3.49 | | div | 1024x20480 | 33.2 | 32.8 | 0.99x | 3.80 | 3.83 | | remainder | 1024x4096 | 9.5 | 9.3 | 0.98x | 2.65 | 2.71 | | remainder | 1024x10240 | 19.3 | 19.0 | 0.99x | 3.27 | 3.31 | | remainder | 1024x20480 | 34.6 | 34.3 | 0.99x | 3.64 | 3.67 | | floor_divide | 1024x4096 | 9.5 | 9.3 | 0.98x | 2.65 | 2.71 | | floor_divide | 1024x10240 | 19.2 | 19.0 | 0.99x | 3.28 | 3.31 | | floor_divide | 1024x20480 | 34.6 | 34.2 | 0.99x | 3.64 | 3.68 | | pow | 1024x4096 | 16.2 | 16.1 | 0.99x | 1.55 | 1.56 | | pow | 1024x10240 | 35.4 | 35.4 | 1.00x | 1.78 | 1.78 | | pow | 1024x20480 | 67.2 | 67.2 | 1.00x | 1.87 | 1.87 | | maximum | 1024x4096 | 20.0 | 19.7 | 0.98x | 1.26 | 1.28 | | maximum | 1024x10240 | 44.7 | 44.2 | 0.99x | 1.41 | 1.42 | | maximum | 1024x20480 | 81.8 | 81.3 | 0.99x | 1.54 | 1.55 | | eq | 1024x4096 | 19.0 | 18.8 | 0.99x | 1.10 | 1.11 | | eq | 1024x10240 | 40.7 | 40.4 | 0.99x | 1.29 | 1.30 | | eq | 1024x20480 | 85.2 | 84.9 | 1.00x | 1.23 | 1.24 | ### Fused Gated Ops (fp16, latency us + bandwidth TB/s) | Op | Shape (MxN) | Eager (us) | Compile (us) | Overhead | Eager BW (TB/s) | Compile BW (TB/s) | |---|---|---|---|---|---|---| | silu_and_mul | 1024x4096 | 9.8 | 9.5 | 0.97x | 2.57 | 2.65 | | silu_and_mul | 1024x10240 | 20.8 | 20.5 | 0.99x | 3.03 | 3.06 | | silu_and_mul | 1024x20480 | 37.6 | 37.2 | 0.99x | 3.35 | 3.38 | | gelu_and_mul | 1024x4096 | 10.6 | 10.3 | 0.97x | 2.38 | 2.45 | | gelu_and_mul | 1024x10240 | 22.2 | 22.0 | 0.99x | 2.83 | 2.86 | | gelu_and_mul | 1024x20480 | 39.4 | 39.2 | 1.00x | 3.20 | 3.21 | | gelu_tanh_and_mul | 1024x4096 | 9.1 | 8.8 | 0.97x | 2.77 | 2.85 | | gelu_tanh_and_mul | 1024x10240 | 18.1 | 17.9 | 0.99x | 3.48 | 3.52 | | gelu_tanh_and_mul | 1024x20480 | 32.9 | 32.7 | 0.99x | 3.82 | 3.84 | ### Takeaways - **Zero overhead**: `torch.compile` adds no runtime cost — compiled paths are consistently 1-8% **faster** than eager due to eliminated Python dispatch overhead - **Unary ops**: 3-8% faster under compile at small sizes (262K elements), converging to ~1-3% at large sizes (4M) - **Binary/fused ops**: 1-2% faster under compile across all DNN-realistic shapes - **Compute-bound ops** (pow, maximum, eq): negligible difference as expected — kernel time dominates 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
|
|
Cao Ying
|
d820ab34f3
|
[Feat][Elementwise] Add BinaryKernel + FusedGatedKernel template ops (22 ops) (#463)
## Summary Implements 22 elementwise operator subclasses using the template base classes from #436: - **20 BinaryKernel ops**: arithmetic (Sub, Mul, Div, Remainder, FloorDivide, Pow, Maximum, Minimum, Lerp), comparison (Eq, Ne, Gt, Lt, Ge, Le), logical (LogicalAnd, LogicalOr), bitwise (BitwiseAnd, BitwiseOr, BitwiseXor) - **2 FusedGatedKernel ops**: GeluAndMul, GeluTanhAndMul - NaN propagation fix for MaximumKernel/MinimumKernel (round 1) Closes #438 ## Dtype Support Matrix | Op Family | Supported input dtypes | Output dtype | Semantics | |---|---|---|---| | `sub`, `mul` | validated in this PR on `float16`, `bfloat16`, `float32` | same as input | PyTorch-style elementwise arithmetic with broadcast | | `div`, `remainder`, `pow`, `floor_divide`, `lerp`, `maximum`, `minimum` | `float16`, `bfloat16`, `float32` | same as input | PyTorch-aligned float arithmetic; unsupported integer dtypes are rejected at construction time | | `eq`, `ne`, `gt`, `lt`, `ge`, `le` | `float16`, `bfloat16`, `float32` | `torch.bool` | Kernel emits int8 1/0 flags, Op layer casts to bool | | `logical_and`, `logical_or` | `bool`, `uint8`, `int8`, `int16`, `int32`, `int64`, `float16`, `bfloat16`, `float32` | `torch.bool` | Non-zero truthiness, matching `torch.logical_*` semantics | | `bitwise_and`, `bitwise_or`, `bitwise_xor` | `uint8`, `int8`, `int16`, `int32`, `int64` | same as input | PyTorch-style bitwise semantics; float dtypes rejected | | `gelu_and_mul`, `gelu_tanh_and_mul` | `float16`, `bfloat16`, `float32` | same as input | Fused gated activation on `(M, 2N) -> (M, N)` | ## Test plan - [x] **AC1**: 20 BinaryKernel + 2 FusedGatedKernel subclasses implemented - [x] **AC2**: Comparison ops output torch.bool - [x] **AC3**: `__all__` updated in both elementwise.py files - [x] **AC4**: L1: 22 ops pass (fp16, same-shape) - [x] **AC5**: L3: 60 broadcast tests pass (bias-add, row, scalar patterns x 20 ops) - [x] **AC6**: L4: 6 edge cases pass (fp32, 4K) - [x] **AC7**: dtype rejection/runtime mismatch paths covered for the float-only binary and fused-gated contracts - [x] **AC8**: `__init__.py` not modified **Test results**: 196/196 passed, 0 failed ## Benchmark Environment: Torch 2.9.1+cu128, CUDA 12.8, NVIDIA H200, Driver 575.57.08 Shapes: (tokens=1024, hidden_dim) - DNN-realistic hidden dimensions ### Binary Arithmetic (fp16, bandwidth TB/s) | Op | Shape (MxN) | TileOPs | PyTorch | Speedup | |---|---|---|---|---| | sub | 1024x4096 | 1.97 | 1.84 | 1.07x | | sub | 1024x10240 | 2.73 | 2.77 | 0.99x | | sub | 1024x20480 | 3.37 | 3.37 | 1.00x | | mul | 1024x4096 | 1.87 | 1.83 | 1.02x | | mul | 1024x10240 | 2.82 | 2.83 | 0.99x | | mul | 1024x20480 | 3.40 | 3.40 | 1.00x | | div | 1024x4096 | 1.82 | 1.87 | 0.98x | | div | 1024x10240 | 2.72 | 2.59 | 1.05x | | div | 1024x20480 | 3.36 | 3.31 | 1.02x | | remainder | 1024x4096 | 1.81 | 1.61 | 1.13x | | remainder | 1024x10240 | 2.73 | 2.40 | 1.14x | | remainder | 1024x20480 | 3.35 | 2.82 | 1.19x | | pow | 1024x4096 | 1.05 | 1.02 | 1.02x | | pow | 1024x10240 | 1.29 | 1.26 | 1.02x | | pow | 1024x20480 | 1.42 | 1.39 | 1.02x | | floor_divide | 1024x4096 | 1.92 | 0.82 | 2.33x | | floor_divide | 1024x10240 | 2.74 | 0.98 | 2.79x | | floor_divide | 1024x20480 | 3.41 | 1.05 | 3.25x | | lerp | 1024x4096 | 1.86 | 1.85 | 1.00x | | lerp | 1024x10240 | 2.73 | 2.81 | 0.97x | | lerp | 1024x20480 | 3.36 | 3.42 | 0.98x | | maximum | 1024x4096 | 1.01 | 1.82 | 0.56x | | maximum | 1024x10240 | 1.22 | 2.71 | 0.45x | | maximum | 1024x20480 | 1.40 | 3.38 | 0.41x | | minimum | 1024x4096 | 1.02 | 1.89 | 0.54x | | minimum | 1024x10240 | 1.24 | 2.70 | 0.46x | | minimum | 1024x20480 | 1.38 | 3.34 | 0.41x | ### Comparison (fp16, bandwidth TB/s) | Op | Shape (MxN) | TileOPs | PyTorch | Speedup | |---|---|---|---|---| | eq | 1024x4096 | 0.78 | 1.72 | 0.45x | | eq | 1024x10240 | 1.02 | 2.49 | 0.41x | | eq | 1024x20480 | 1.05 | 3.17 | 0.33x | | ne | 1024x4096 | 0.79 | 1.65 | 0.48x | | ne | 1024x10240 | 1.02 | 2.56 | 0.40x | | gt | 1024x4096 | 0.78 | 1.72 | 0.45x | | gt | 1024x10240 | 1.02 | 2.48 | 0.41x | | lt | 1024x4096 | 0.79 | 1.72 | 0.46x | | lt | 1024x10240 | 1.02 | 2.55 | 0.40x | | ge | 1024x4096 | 0.78 | 1.72 | 0.45x | | ge | 1024x10240 | 1.02 | 2.56 | 0.40x | | le | 1024x4096 | 0.79 | 1.64 | 0.48x | | le | 1024x10240 | 1.02 | 2.55 | 0.40x | ### Logical (fp16 input, bandwidth TB/s) | Op | Shape (MxN) | TileOPs | PyTorch | Speedup | |---|---|---|---|---| | logical_and | 1024x4096 | 0.78 | 2.17 | 0.36x | | logical_and | 1024x10240 | 1.02 | 3.37 | 0.30x | | logical_or | 1024x4096 | 0.85 | 2.05 | 0.41x | | logical_or | 1024x10240 | 1.15 | 3.41 | 0.34x | ### Bitwise (int32, bandwidth TB/s) | Op | Shape (MxN) | TileOPs | PyTorch | Speedup | |---|---|---|---|---| | bitwise_and | 1024x4096 | 2.59 | 2.55 | 1.01x | | bitwise_and | 1024x10240 | 3.31 | 3.42 | 0.97x | | bitwise_or | 1024x4096 | 2.51 | 2.54 | 0.99x | | bitwise_or | 1024x10240 | 3.30 | 3.41 | 0.97x | | bitwise_xor | 1024x4096 | 2.51 | 2.61 | 0.96x | | bitwise_xor | 1024x10240 | 3.37 | 3.36 | 1.00x | ### Fused Gated (fp16, bandwidth TB/s) | Op | Shape (MxN) | TileOPs | PyTorch | Speedup | |---|---|---|---|---| | gelu_and_mul | 1024x4096 | 1.61 | 0.63 | 2.53x | | gelu_and_mul | 1024x10240 | 2.27 | 0.77 | 2.96x | | gelu_and_mul | 1024x20480 | 2.75 | 0.84 | 3.29x | | gelu_tanh_and_mul | 1024x4096 | 1.87 | 0.66 | 2.82x | | gelu_tanh_and_mul | 1024x10240 | 2.76 | 0.81 | 3.41x | | gelu_tanh_and_mul | 1024x20480 | 3.36 | 0.89 | 3.79x | ### Broadcast: Bias-Add Pattern (fp16, bandwidth TB/s) | Op | a_shape | b_shape | TileOPs | PyTorch | Speedup | |---|---|---|---|---|---| | sub | 1024x4096 | 1x4096 | 1.53 | 0.87 | 1.75x | | sub | 1024x10240 | 1x10240 | 2.53 | 1.14 | 2.22x | | sub | 1024x20480 | 1x20480 | 3.11 | 1.27 | 2.45x | | mul | 1024x4096 | 1x4096 | 1.54 | 0.88 | 1.75x | | mul | 1024x10240 | 1x10240 | 2.42 | 1.12 | 2.16x | | mul | 1024x20480 | 1x20480 | 3.18 | 1.28 | 2.47x | | div | 1024x4096 | 1x4096 | 1.47 | 0.82 | 1.80x | | div | 1024x10240 | 1x10240 | 2.37 | 1.06 | 2.24x | | div | 1024x20480 | 1x20480 | 2.82 | 1.16 | 2.43x | ### Takeaways - **Broadcast (bias-add)**: 1.75-2.47x faster (stride-based codegen beats PyTorch broadcast) - **Fused gated**: 2.5-3.8x faster (kernel fusion eliminates intermediate writes) - **floor_divide**: 2.3-3.3x faster (TileLang codegen advantage) - **remainder**: 1.1-1.2x faster - **Binary arith** (sub/mul/div/lerp/pow): ~parity with PyTorch at DNN-realistic sizes - **Bitwise**: ~parity (0.96-1.01x) - **Comparison**: ~0.3-0.5x slower (int8 intermediate + bool cast overhead; optimization opportunity) - **Logical**: ~0.3-0.4x slower (int8 intermediate + bool cast; optimization opportunity) - **Maximum/minimum**: ~0.4-0.6x slower (T.isnan NaN propagation overhead; optimization opportunity) ## Changes | File | Description | |---|---| | `tileops/kernels/elementwise.py` | 20 BinaryKernel + 2 FusedGatedKernel subclasses | | `tileops/ops/elementwise.py` | 22 matching Op subclasses | | `tests/ops/test_binary_arith.py` | Arithmetic smoke + broadcast + edge case tests | | `tests/ops/test_comparison.py` | Comparison smoke + broadcast + bool output tests | | `tests/ops/test_logical.py` | Logical op smoke + broadcast tests | | `tests/ops/test_bitwise.py` | Bitwise op smoke + broadcast tests | | `tests/ops/test_fused_gated.py` | Fused gated op tests | | `benchmarks/ops/bench_binary_elementwise.py` | Benchmark: same-shape + broadcast (DNN-realistic 2D shapes) | ## Follow-up Issue The next issue that should be taken immediately after this PR is: - #440: add `torch.compile` / fake-tensor support for the shared elementwise template via `custom_op` registration --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |
|
|
Cao Ying
|
ec595f03ca
|
[Feat][Elementwise] Add template infrastructure: base classes, smoke tests, and strategy benchmark (#453)
## Summary Build the complete template infrastructure for 66 elementwise ops: - **3 kernel base classes** — `UnaryKernel`, `BinaryKernel`, `FusedGatedKernel` with strategy support (`direct`, `explicit_parallel`, `register_copy`) - **3 op base classes** — `UnaryOp`, `BinaryOp`, `FusedGatedOp` with `coalesce_broadcast_dims` utility for stride-based broadcast - **3 concrete example ops** — `relu`, `add`, `silu_and_mul` (one per template family) - **Strategy benchmark** — determines `DEFAULT_STRATEGY` per kernel type with bandwidth rationale Closes #436 ## Test plan - [x] AC1: UnaryKernel, BinaryKernel, FusedGatedKernel implemented with all strategies - [x] AC2: `coalesce_broadcast_dims` handles all 6+ broadcast patterns correctly (6 unit tests) - [x] AC3: relu, add, silu_and_mul pass L1 correctness tests (14 configs total) - [x] AC4: 4 broadcast patterns pass L3 tests (bias-add, row, scalar, interleaved) - [x] AC5: Strategy benchmark run; DEFAULT_STRATEGY updated with rationale - [x] AC6: `__init__.py` re-exports base classes only (not concrete ops) **29 tests pass, 14 benchmarks pass.** ## Benchmark ### Unary (relu) Strategy | strategy | bandwidth (TB/s) | |---|---| | register_copy | 2.40 | | direct | 1.04 | | explicit_parallel | 0.86 | ### Binary (add) Strategy | strategy | bandwidth (TB/s) | |---|---| | explicit_parallel | 2.62 | | direct | 1.42 | **Decision**: Unary `DEFAULT_STRATEGY = register_copy`, Binary `DEFAULT_STRATEGY = explicit_parallel`. ## Files changed (9 files, +1177 lines) | File | Description | |---|---| | `tileops/kernels/elementwise.py` | Kernel base classes + strategy factory | | `tileops/ops/elementwise.py` | Op base classes + coalesce_broadcast_dims | | `tileops/kernels/__init__.py` | Re-export base kernel classes | | `tileops/ops/__init__.py` | Re-export base op classes | | `tests/ops/test_activation.py` | Unary op tests (relu) | | `tests/ops/test_binary_arith.py` | Binary op + broadcast tests (add) | | `tests/ops/test_fused_gated.py` | Fused gated op tests (silu_and_mul) | | `benchmarks/ops/bench_activation.py` | Unary strategy benchmark | | `benchmarks/ops/bench_binary_arith.py` | Binary strategy benchmark | --------- Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> |