Raisetolinalg#412
Draft
arpitj1 wants to merge 171 commits into
Draft
Conversation
…erands and map for linalg.generic
…ded extra tests in lit test
…f debufferizing added which works for tiling and fusion
Adds the darknet (pjreddie/darknet) third-party clone as a fifth
benchmark suite in the IR explorer. The "kernels" are individual .c
files in src/; the bake runs cgeist + raise + match on each.
Approach:
bake_darknet_mlir.sh iterates over third_party/darknet/src/*.c,
baking each through:
cgeist --function='*' --no-inline ...
polygeist-opt --raise-affine-to-linalg-pipeline --linalg-debufferize
kernel_match_rewrite.py
Files use --function='*' because darknet's compute is spread across
many entry points (gemm_nn/nt/tn/tt all need to lift); --no-inline
prevents the raise pass from collapsing init-into-kernel boilerplate
the way it used to on polybenchGpu.
Results (46 .c files, ~25K LOC total):
cgeist OK: 28 (61%)
raise OK: 23 (50%)
produced ≥1 linalg.generic: 18 (39%)
produced ≥1 kernel.launch: 1 ( 2%)
The 1 file that matches: src/gemm.c (6 launches across gemm_nn / nt /
tn / tt / bin). The 17 raise-OK-but-no-match files are an actionable
list of missing matcher templates: pooling (avg/max), batchnorm, LRN,
residual-add, GRU/LSTM gates, transposed conv, locally-connected, dense
+ bias, softmax-with-control-flow, l2norm. The 18 cgeist-fails are
mostly framework code (parser, image, data, network) with no compute.
darknet's actual production hot path is gemm_nn (TA=TB=0). The matcher
hits it as @cublasDaxpy (the inner loop has the scalar-hoisted axpy
shape) but doesn't compose the outer two loops back up into gemm.
gemm_nt and gemm_tt use the conventional sum-accumulator form and do
match as @cublasDgemm_alpha_only. Fixing gemm_nn composition is a
high-value matcher follow-up — it would auto-cover every conv layer
darknet runs at inference time (since every conv goes through gemm_nn
via im2col).
New section in build_ce_viewer.py:
- DARKNET_ROOT / DARKNET_MLIR_DIR path constants
- DARKNET_KERNELS dict (45 .c files)
- DARKNET_NOTES per-file with parallelism tag + characterisation
- DARKNET_BLOCKERS per-file mapped to existing taxonomy
(matcher-gap, cgeist-gap, debuf-bug, none)
- find_kernel_c dispatch for kset="darknet"
- build_index gains darknet_stats parameter
- new section + nav link to "#darknet"
The third_party/darknet/ clone itself is NOT committed (it's a vendored
upstream, would bloat the repo to ~25K LOC for the framework + cfgs).
The bake script's PATH is hardcoded so a fresh clone reproduces the
results.
…end on Jetson Orin
Polybench-style C kernels in third_party/cnn-extracted/, each lifted through the
full Polygeist pipeline (cgeist → raise → debufferize → matcher → ABI lowering →
LLVM IR → aarch64 cross-compile → Jetson silicon).
Five extracted-darknet baseline kernels (matcher templates + lowering branches +
cuDNN/cuBLAS shims + harness + per-kernel HTML page in the IR explorer):
conv2d_batched → cudnnConvolutionFwd_batched 23.8x LARGE
maxpool_batched → cudnnMaxPoolFwd_batched 1.29x LARGE
batchnorm_batched → cudnnBatchNormalizationForwardInference 0.38x LARGE
shortcut_batched → cudnnAddTensor_batched 0.08x LARGE
conv_bn_relu_batched → cudnnConvolutionBiasActivationForward
(with host-side BN folding) 23.5x LARGE
Four fusion-optimization kernels (algebraic rewrites + faster cuBLAS/cublasLt/
cuDNN entry points):
conv_bias_relu_add_batched → cudnnConvolutionBiasActivationForward
(α2*Z addend for ResNet skip) 23x LARGE
gemm_bias_relu → cublasLtMatmul EPILOGUE_RELU_BIAS 901x LARGE
ata_gemm → cublasSsyrk (operand-alias discriminator
detects AᵀA pattern; half the flops) 3393x LARGE
conv1x1_batched → cublasSgemmStridedBatched (4-par+1-red
shape distinguishes K=1 from K×K) 105x LARGE
Cross-cutting infrastructure additions:
* Matcher: ~9 new CompositionEntry templates + AᵀA→syrk post-unify operand-alias
discriminator in kernel_match_rewrite.py. Per-step span replacement preserves
intervening polygeist.submap ops between matched generics.
* Lowering pass: resolveSubmapBase now chains through both polygeist.submap and
polygeist.submapInverse (up to 16 hops). New pre-pass elides redundant
memset_zero_{1D,2D} launches preceding any β=0 op (syrk). Dtype-suffixed
memset dispatch (f32 alongside f64).
* Runtime: cublasLt linkage (libcublasLt.so.12); ensure_cublaslt() helper.
Host-side BN-folding for fused conv+bn+relu (precompute scaled filter + bias).
All cuDNN algo-selection loops use array-sized cudnnConvolutionFwdAlgoPerf_t
buffers (avoiding the stack-smash that bit single-struct attempts).
* Build: scripts/correctness/extracted_darknet_jetson.sh handles all 9 kernels;
bake_extracted_darknet_mlir.sh produces per-stage MLIR snapshots for the
IR explorer; -lcublasLt added to link line.
* IR explorer: two new sections (extracted darknet, Fusion optimization) with
Compiler Explorer deep-links + per-kernel raised/debuf/matched IR preview
pages.
All four fusion optimizations are 100% bit-exact (or FP-noise within 1e-4 print
precision); LARGE speedups range 23x→3393x over the CPU 3-loop reference on the
Jetson Orin (Tegra Ampere, FP32, cuDNN 9.x, CUDA 12.6).
…nv2d + 4 image filters on Jetson Orin
* New LowerKernelLaunchToPVA pass — owns the matcher's i8/i16
@cudnnConvolution2D_9tap_* launches plus new
@pvaBoxFilter_3x3_i{8,16}, @pvaGaussianFilter_3x3_i{8,16},
@pvaBilateralFilter_3x3_i{8,16}, @pvaHistogramEqualization_i8
symbols. Each routes to a polygeist_pva_* runtime shim. Disjoint
symbol set from --lower-kernel-launch-to-cublas; the two passes
run side by side; either order works.
* Shared 9-tap conv lowering helper extracted out of
LowerKernelLaunchToCuBLAS.cpp into KernelLaunchLoweringUtils.{h,cpp}
so both backend passes call the same body. Added a parallel
lowerImageFilter2Operand helper for the 2-memref filter launch
shape (Box/Gaussian/Bilateral/HistogramEq).
* cuBLAS pass: dropped i8/i16 from shimSymbolFor + the dispatch
switch; PVA-claimed launches fall through with a `continue`
instead of erroring out. Net diff is small in the cuBLAS pass
file (the 3 helpers moved out are the bulk of the delta).
* New PVA runtime shim runtime/polygeist_pva_rt.c with:
- cudaSetDevice + nvcvAllocatorConstructPva + non-blocking
stream init (idempotent, lazy, persistent for process lifetime)
- make_pva_image_tensor_dtype: HWC tensor alloc through the PVA
allocator with arbitrary NVCV dtype (needed because half the
PVA ops are U8-only; we reinterpret i8 bytes as U8)
- CupvaMemGetHostPointer-mediated host I/O (raw cudaMemcpy
segfaults on cuPVA-allocated pages; the host-pointer mapping
is mandatory)
- One pva<Op>Create / pva<Op>Submit wrapper per op
- (M-2)×(N-2) interior copy from PVA output back to caller B
to honour the matcher's &B[1][1] pointer-shift convention
(writing the full M×N overflows B by N+1 bytes)
* Matching CPU reference stubs in polygeist_cublas_rt_cpu.c modelled
to mirror PVA hardware semantics: centred kernel anchor, REPLICATE
border, Q-format >>qbits shift, unsigned-kernel reinterpretation
for Conv2d; rounded-mean (sum + 4) / 9 for BoxFilter; canonical
[1,2,1;2,4,2;1,2,1] / 16 for Gaussian; textbook 256-bin CDF-LUT
for HistogramEq. Bilateral has a pass-through stub (the
non-linear hardware semantics aren't worth mirroring bit-exactly).
* third_party/polybenchGpu-extracted/conv2d_i8.c — i8 variant of
the 9-tap stencil (i16 already existed). Matcher fires on it via
the existing dtype-suffix template + emits
@cudnnConvolution2D_9tap_i8, which the new PVA pass claims.
* Cross-compile script conv2d_cudnn_jetson_dtype.sh: i8 dtype
branch added; PVA-library link line (-lpva_operator -lcvcuda
-lnvcv_types -lcupva_host) plus direct DT_NEEDEDs for
-lnvscibuf -lnvscisync via -Wl,--no-as-needed (deferred
resolution segfaults during libcupva_host init constructors);
step (5) now invokes both --lower-kernel-launch-to-cublas
and --lower-kernel-launch-to-pva.
* Four hand-authored kernel.launch test scaffolds in
scripts/correctness/pva_{boxfilter,gaussian,bilateral,histeq}_jetson.sh.
Matcher templates for these C-level patterns aren't written yet,
so each script synthesises the kernel.launch MLIR directly and
runs the rest of the pipeline normally — same harness, wrapper,
ABI lowering, and link line.
* IR explorer (scripts/correctness/build_ce_viewer.py): new "PVA
backend" section at the bottom. Shows the 6 PVA-routed kernels
with their op name, libpva_operator entry points, shim symbol,
and Jetson PVA wall-clock at each size we benchmarked. No CPU
comparison in this view (CPU stubs exist for separate per-op
bit-exact validation).
* CLAUDE.md: "point, don't copy" rule for gated-distribution NVIDIA
SDKs. PVA Solutions / cuPVA SDK headers consumed via -I at build
time; never copied into the Polygeist tree.
End-to-end silicon validation on Jetson Orin: bit-exact PVA-vs-CPU
diff for Conv2d i8/i16, BoxFilter, Gaussian, and HistogramEq at 256².
Bilateral runs cleanly; visual spot-check only (non-linear).
Conv2d at 10240×10240: PVA 216 ms vs CPU 499 ms (2.3× speedup for i8).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Some modifications to fuse linalg.generic op with for op