nki(2d_conv): NKI (Trainium) implementation

[bowencui123]

NKI (AWS Trainium) implementation of 2d_conv, split out of the consolidated NKI branch cecilia/feature/nki-vector-add (nki-all-operators, #259) so each operator can be reviewed independently. Supersedes #215 (older per-operator branch: legacy neuronxcc.nki imports; this is the migrated import nki version).

Files: A benchmarks/operators/2d_conv/impl_nki.py

Status: imports and exposes run()/get_last_config() on trn2 (nki 0.6.0); not individually re-benchmarked in this split; imports via importlib (the operator name starts with a digit, so a plain import statement is invalid — the engine uses importlib)

Implementation by @Cecilia123li. Timing/identity infrastructure: #261; Trainium peak/roofline infra: #262.

🤖 Generated with Claude Code

https://claude.ai/code/session_012Q38kGmXvyoeM1qtCbheSL

Autotune (9ec3a35)

NKI tunables	Triton counterpart	note
block_size_oh (output rows per block)	BLOCK_SIZE_BATCH_HEIGHT_WIDTH/..._FEAT	—

autotune=False keeps the previous constants (default numbers unchanged). Validation on trn2, case 0 (default run + autotune code path with the candidate timer stubbed — no sweep; --autotune runs a real sweep):

# initial run
[2d_conv] default : verify=FAIL (1s) Comparing
[2d_conv] autotune: verify=FAIL (1s) last_config={'block_size_oh': 2} trace_records=1 Comparing
STUB_EXIT=1
Params       |    Dtype |  Torch(ms) |      NKI(ms) |  Speedup(N)
n=75497472   | fp16     |     0.0324 |       0.1112 |        0.29
# re-validation (tolerance parsing fixed in the checker; engine line above stands)
[2d_conv] default : verify=OK (0s) 
[2d_conv] autotune: verify=OK (0s) last_config={'block_size_oh': 2} trace_records=1 
STUB_EXIT=0

[bowencui123]

Review pass (Bowen's assignment): Triton/cuTile alignment + scaffold removal + real autotune & default runs on trn2 (post-#309 runtime-trace timing framework, case 0, --warmup 2 --repeat 10; autotune numbers come from a real selector sweep replayed by the profile worker).

#264 2d_conv — NO CODE CHANGE (note)

• Same family as 1d/3d. Note: run() restricts to batch=1, k=3, C_in=C_out=128 (the benchmarked
  config) although the kernel is written general — widening needs its own validation, out of scope.
• Tested ({'batch': 1, 'in_channels': 128, 'out_channels': 128, 'kernel_size': 3, 'stride': 1, 'padding': 1, 'groups': 1, 'H': 16} fp16): default 0.1060 ms (torch 0.0253),
  autotune 0.0992 ms, winner namespace(block_size_oh=4).

🤖 Generated with Claude Code

https://claude.ai/code/session_011Wvs1tztZTGQFD78YdZaha