FrRay
/
TileOPs-Metax
forked from ccf-ai-infra/TileOPs-Metax
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
429 Commits 6 Branches 0 Tags 15 MiB
Commit Graph

7 Commits

Author SHA1 Message Date
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 4dc54b7)
- [x] **AC-4**: No orphaned files remain after migration — all 11 old
filenames confirmed absent; `mhc_post.py`, `mhc_pre.py` at
`workloads/ops/` root; `nsa_utils.py` at `workloads/` root

## Follow-up

No follow-up issues or suggestions.

---------

Co-authored-by: Ibuki 🍃 — a wind born from Claude Opus <Ibuki-wind@users.noreply.github.com>
 2026-04-13 10:44:06 +08:00
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>
 2026-04-05 11:38:51 +08:00
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>
 2026-03-17 10:43:28 +08:00
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>
 2026-03-16 10:38:13 +08:00
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>
 2026-03-13 22:47:01 +08:00
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>
 2026-03-13 13:21:27 +08:00
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>
 2026-03-11 17:23:03 +08:00