forked from ccf-ai-infra/TileOPs-Metax
12 Commits
| Author | SHA1 | Message | Date |
|---|---|---|---|
Cao Ying
|
c986df5407
|
[Refactor][POOL] strip over-design and scaffolding from the merged cleanup PRs (#1779)
## Summary
- Remove over-engineered pool hooks and private snapshot tests; make
indexed max-pool forwards explicit.
- Deduplicate GLA, Mamba, and formula tests; remove dead helpers,
commented benchmark rows, and process metadata.
- Consolidate benchmarks onto manifest workloads and
`ManifestBenchmark`; delete obsolete benchmark modules and factor
redundant sweep axes.
- Require implemented ops to declare `kernel_map` and manifest-driven
benchmark coverage, filling the corresponding manifest gaps.
- Reduce the repository by 2,415 net lines without changing runtime
operator behavior.
## Test plan
- [x] pre-commit passed
- [x] Pool tests passed: 246
- [x] Perf/formula and validator tests passed: 128
- [x] Changed benchmark modules collected: 538 nodes
- [x] Test node delta: 476 → 391 (-85)
- [x] `python scripts/validate_manifest.py --strict` passed
---------
Co-authored-by: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com>
|
|
|
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
|
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
|
24886ce6d0
|
[Fix][Benchmark] stop losing results on native failures; derive workload keys from signature (#1745)
Closes #1740. Closes #1746. ## Summary - **One failure no longer destroys the session's results.** tilelang's fd-level output suppression corrupted pytest's capture stream; when the 07-24 nightly stalled >900s, the timeout handler crashed and all 1328 completed results were lost. Suppression now applies only when stdout/stderr are the real fds 1/2. The stall itself did not reproduce (7 attempts, latency parity with the last green nightly); mechanism unidentified. Process isolation: #1744. - **Measurement never degrades silently.** CUPTI projection failure raises a dedicated sentinel — genuine CUDA errors/OOM propagate; the CUDA-events fallback and the >1 GiB clone-skip are logged and marked in results. - **Workload keys derive from the manifest signature.** The hardcoded shape-key allowlist (which let #1734 break dropout/fft collection) is gone; the contract lives once in `tileops.manifest`, consumed by both `workloads_to_params` and validator rule R21. 26 manifest entries aligned to their signatures (`x_shape`→`input_shape`; `DropoutOp`/`FFTC2COp` input `x`→`input` incl. `forward()`). Unknown signature/entry/workload keys are rejected; the validator stays total over malformed YAML. - **Cleanup.** "harness" jargon removed repo-wide; benchmark_base tests folded into one file; 24 redundant/tombstone validator tests deleted; `init_dims`/`parity_opt_out` special cases replaced by generic unknown-key rejection. Net −350 lines. ## Test plan - `validate_manifest.py` full pass; 207 validator + 17 benchmark unit tests; `--collect-only benchmarks/ops` 1638, zero errors - GPU (H200): dropout/fft tests 37 passed, benches 6 passed; `bench_moe_grouped_gemm_nopad` 4/4; fallback + error-propagation contract tests ## Benchmark No kernel change and the timing protocol (CUPTI kernel-only, warmup/repeats/trials, L2 flush) is untouched; this PR changes suppression selection and fallback/error semantics around it. Latency parity: `bench_moe_grouped_gemm_nopad` decode-down 2.3061 ms vs 07-17's 2.3088 ms. ## Regression - Bench files keep the `(shape, dtype[, extra])` parametrize contract - Post-merge check: next nightly progresses past `bench_moe_grouped_gemm_nopad.py` with a non-empty `bench_results.xml` --------- Co-authored-by: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com> |
|
|
Cao Ying
|
4e717d3abf
|
[Bench][Elementwise] preserve shape tuples in elementwise bench files (#1202)
Closes #1176 ## Summary - Convert 10 elementwise-family bench files (`benchmarks/ops/bench_{activation,binary_arith,binary_elementwise,binary_strategy,dropout,elementwise_fp8,independent_elementwise,rope,unary_elementwise,unary_strategy}.py`) from scalar `n_total: int` parametrization to shape-tuple parametrization with LLaMA-default geometry. - Route the original shape tuple through `BenchmarkReport.record()` so `profile_run.log` and PR-body benchmark tables carry tuples like `(1024, 4096)` verbatim instead of flattening to a scalar element count. - Extend the shared `benchmarks/benchmark_base.py::_is_serializable()` helper to accept tuples of primitives recursively; existing scalar/dtype handling is unchanged. - Exercise at least one non-pow2 shape per op (`11008`, the LLaMA-7B intermediate dim) so bench harnesses exercise tail handling. Op kernel signatures are untouched — call sites that still need a flat element count compute `math.prod(shape)` locally. - Pre-W2 cleanup tracked under #1142: lets the upcoming W2 bench files inherit a clean shape-tuple convention instead of perpetuating the flattened-`n_total` pattern. ## Test plan - [x] AC-1: Modified files pass unit tests. (`pre-commit run --all-files` all hooks Passed; `pytest -m smoke` across the 10 bench files: 48 passed, 2 pre-existing clamp failures unrelated to this PR, 1 skipped, in 6.69s.) - [x] AC-2: All 10 bench files parametrize over shape tuples with model-relevant LLaMA-default geometry, and at least one non-pow2 shape is exercised per op. (Non-pow2 coverage spans `(1024, 11008)`, `(2048, 300)`, `(3000, 128)`, etc. across the 9 directly modified bench files; `bench_unary_elementwise.py` inherits shapes from `tileops/manifest/` per the trust model — non-pow2 manifest workloads are a separate manifest PR concern.) - [x] AC-3: `profile_run.log` produced by these bench files shows shape tuples in the parametrization columns, not bare element counts. (Verified on `bench_binary_elementwise`, `bench_unary_strategy`, `bench_dropout` smoke runs; example row: `| sub | (1024, 4096) | torch.float16 | torch.float16 | 0.0087 | 0.4795 | 2.8768 |`.) - [x] AC-4: A representative PR-body benchmark section produced from one of the rewritten benches conforms to the per-op-section-with-tuple-shape format from `.foundry/mold/benchmark-template.md`. (See Benchmark section below.) - [x] AC-5: No op kernel signatures change; diff is confined to `benchmarks/`. (`git diff --name-only main..HEAD`: 10 files, all under `benchmarks/`; `benchmark_base.py` change is a small helper extension.) ## Benchmark **Environment**: H200, smoke profile (representative — full sweep is run by `make bench`). ### sub (binary elementwise) | Shape | dtype | TileOPs (ms) | torch (ms) | Speedup | BW (TB/s) | | ----- | ----- | ------------ | ---------- | ------- | --------- | | (1024, 4096) | float16 | 0.0087 | 0.0120 | 1.38× | 2.88 | **Takeaways:** Sample row demonstrates the shape-tuple format conformance; full per-op tables are produced by running `make bench` after merge. The point of this PR is the format/parametrization fix, not new performance numbers. **Command:** `PYTHONPATH="$PWD" python -m pytest benchmarks/ops/bench_binary_elementwise.py -v` 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>
|
|
|
Cao Ying
|
feb92848cd
|
[Refactor][Benchmark] Add explicit type parameters to BenchmarkBase subclasses (#955)
## Summary Add explicit generic type parameters to all 97 `BenchmarkBase` subclasses in `benchmarks/ops/`, replacing raw `BenchmarkBase` with `BenchmarkBase[ConcreteWorkloadType]`. - Concrete workload types used (imported `*Test` classes, local `*BenchCase` classes) - Two local protocols introduced for benchmarks with multiple workload types (`_UnaryWorkload`, `BinaryWorkload`) - Zero pyright regression on changed files vs main Closes #937 ## Test plan - [x] **AC-1**: All `BenchmarkBase` subclasses have explicit type parameter (97/97) - [x] **AC-2**: All existing benchmark tests pass (7/7) - [x] **AC-3**: No raw `BenchmarkBase` (without type parameter) remains in `benchmarks/ops/` ## Follow-up - #956 — Narrow `gen_inputs` return types across workload protocols --------- Co-authored-by: Ibuki 🍃 — a wind born from Claude Opus <Ibuki-wind@users.noreply.github.com> |
|
|
Cao Ying
|
bdbd7a8324
|
[Refactor] Rename base modules to <module>_base.py convention (#950)
## Summary Unify base module file naming to `<module>_base.py` convention across four modules: - `workloads/base.py` -> `workloads/workload_base.py` - `tileops/ops/op.py` -> `tileops/ops/op_base.py` - `tileops/kernels/kernel.py` -> `tileops/kernels/kernel_base.py` - `benchmarks/benchmark.py` -> `benchmarks/benchmark_base.py` All imports, string references, and doc paths updated. No logic changes — pure rename + import propagation. Closes #947 ## Test plan - [x] **AC-1**: All four files renamed to <module>_base.py -- pass - [x] **AC-2**: All imports updated, no broken references -- pass - [x] **AC-3**: docs/testing.md path references updated -- pass - [x] **AC-4**: Pre-commit and existing tests pass -- pass **Test results**: 2368/2368 tests passed, 0 failed --------- Co-authored-by: Ibuki 🍃 — a wind born from Claude Opus <Ibuki-wind@users.noreply.github.com> |
|
|
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> |
|
|
Ang Gao
|
fce542af21
|
[Fix][Bench] Standardize benchmark baseline tags to match actual implementations (#624) (#657)
## Summary - Replace generic `tag="baseline"` across 52 benchmark files with precise tags that reflect the actual backend implementation - Nightly report Baseline column now shows exactly what was used for comparison instead of uninformative "baseline" ## Changes 1. **Generic → specific**: `"baseline"` → descriptive names across all benchmark files 2. **Backend precision**: `torch.matmul` → `"torch-cublas"`, `torch.fft.fft` → `"torch-cufft"`, `F.scaled_dot_product_attention` → `"torch-sdpa"` 3. **Ref classification**: hand-written multi-step references unified as `"torch-ref"` 4. **Naming consistency**: `"FA3"` → `"fa3"`, `"pytorch-*"` → `"torch-*"`, underscores → hyphens, `"fla_bwd_with_recompute"` → `"fla"` 5. **Rules update**: `.claude/rules/benchmark.md` tag registry expanded ## Tag taxonomy | Tag | Meaning | |-----|---------| | `"torch"` | PyTorch built-in API with its own CUDA kernel (softmax, layer_norm, sum, etc.) | | `"torch-cublas"` | PyTorch API → cuBLAS (matmul) | | `"torch-cufft"` | PyTorch API → cuFFT (fft) | | `"torch-sdpa"` | `F.scaled_dot_product_attention` | | `"torch-cudnn"` | PyTorch → cuDNN (batch_norm fwd) | | `"torch-autograd"` | PyTorch built-in fwd + autograd bwd (batch_norm bwd) | | `"torch-ref"` | Hand-written multi-step PyTorch reference | | `"fa3"` / `"fla"` / `"triton"` / `"sgl-kernel"` / `"vllm"` | External libraries | ## Test plan - [ ] `pytest benchmarks/ --collect-only` passes (no import errors) - [ ] Nightly report baseline column shows specific tags instead of "baseline" Closes #624 |
|
|
Cao Ying
|
9020a726c1
|
[Feat][Elementwise] Add strategy selection to FusedGatedKernel (#490)
Closes #441 ## Summary Add strategy selection (direct + explicit_parallel) to FusedGatedKernel, aligning with the Unary/Binary kernel patterns. - Added `_make_fused_gated_direct` kernel builder and strategy dispatch to `FusedGatedKernel` - Threaded `strategy` argument through `FusedGatedOp` in the ops layer - Added 13 new tests covering both strategies for silu_and_mul, gelu_and_mul, gelu_tanh_and_mul - Added strategy benchmark harness to `bench_binary_elementwise.py` **Supported dtypes**: `float16`, `bfloat16`, `float32` (inherited from `FusedGatedKernel.SUPPORTED_DTYPES`). ## Test plan - [x] AC-1: Decision documented with benchmark data — explicit_parallel chosen as default - [x] AC-2: FusedGatedKernel supports both strategies, all tests pass - [x] AC-3: Strategy validation uses ValueError (not assert) **Test commands:** ```bash PYTHONPATH="$PWD" python -m pytest tests/ops/test_fused_gated.py -v ``` ## Structural Compliance All checks passed. ## Benchmark **Environment**: NVIDIA H200, CUDA 12.8, PyTorch 2.9.1+cu128, TileLang 0.1.8 Strategy comparison: `direct` vs `explicit_parallel` (before/after). ### silu_and_mul | Shape | dtype | Direct (ms) | Direct BW (TB/s) | Explicit (ms) | Explicit BW (TB/s) | Speedup | |-------|-------|------------|------------------|--------------|-------------------|---------| | 1024×4096 | fp16 | 0.020 | 1.35 | 0.010 | 2.41 | 1.8x | | 1024×4096 | bf16 | 0.020 | 1.35 | 0.010 | 2.58 | 1.9x | | 1024×4096 | fp32 | 0.020 | 2.38 | 0.020 | 3.25 | 1.4x | | 1024×10240 | fp16 | 0.040 | 1.45 | 0.020 | 2.83 | 2.0x | | 1024×10240 | bf16 | 0.040 | 1.46 | 0.020 | 3.24 | 2.2x | | 4096×4096 | fp16 | 0.070 | 1.50 | 0.030 | 3.04 | 2.0x | | 4096×4096 | bf16 | 0.070 | 1.51 | 0.030 | 3.46 | 2.3x | | 4096×4096 | fp32 | 0.070 | 2.74 | 0.050 | 4.00 | 1.5x | ### gelu_tanh_and_mul | Shape | dtype | Direct (ms) | Direct BW (TB/s) | Explicit (ms) | Explicit BW (TB/s) | Speedup | |-------|-------|------------|------------------|--------------|-------------------|---------| | 1024×4096 | fp16 | 0.020 | 1.38 | 0.010 | 2.73 | 2.0x | | 1024×10240 | fp16 | 0.040 | 1.49 | 0.020 | 3.20 | 2.1x | | 4096×4096 | fp16 | 0.070 | 1.51 | 0.030 | 3.39 | 2.2x | | 4096×4096 | bf16 | 0.070 | 1.51 | 0.030 | 3.84 | 2.5x | | 4096×4096 | fp32 | 0.070 | 2.80 | 0.050 | 4.12 | 1.5x | **Takeaways:** - `explicit_parallel` is **1.4–2.5x faster** than `direct` across all shapes/dtypes - Speedup increases with tensor size — 8 elements/thread amortizes loop overhead better at scale - bf16 ≈ fp16 latency; fp32 shows smaller strategy gap due to higher per-element bandwidth - Default choice `explicit_parallel` validated across 54 benchmark cases **Benchmark command:** ```bash PYTHONPATH="$PWD" python -m pytest benchmarks/ops/bench_binary_elementwise.py::test_fused_gated_strategy_bench -v # 54 passed in 157s ``` --------- 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> |