[DLight] Add CPU Reduction schedule rule for softmax-like operators

[swjng]

This PR resolves part of #18569

Description

This PR adds a DLight CPU schedule rule targeting reduction patterns
(softmax, layer norm, RMS norm) that previously had no CPU-specific
schedule.

Without this rule, LLVM auto-vectorization produces suboptimal code for
RVV targets — #18569 reports RVV softmax is 1.34x slower than scalar.
The root cause is:

• LLVM fully unrolls the 185-element reduction loop
• Generates harmful fixed-width vectors (<8 x float> = 256-bit) on a
  128-bit VLEN target
• No parallelization of the batch axis

Schedule strategy

1. Parallelize leading spatial axes (batch dimension)
2. Compute-at all blocks under the spatial loop for data locality
3. Vectorize injective blocks (exp, delta, norm) on inner axis
4. Split + unroll reduction inner axis to VLEN-sized chunks

Results (shape=(14, 185), float32, fast_softmax)

Config	ASM Instructions	Vector Insns	LLVM IR Lines
RV scalar (baseline)	1,463	0	1,792
RVV unscheduled (the bug)	3,282	960	2,345
RVV + this schedule	1,111	105	1,338

• 66% fewer instructions vs unscheduled RVV
• 24% fewer instructions vs scalar baseline
• fast_softmax polynomial exp fully vectorizes into RVV vfsub/vfmul/vfmax instructions with zero scalar exp() calls

Limitations (follow-up work)

• Reduction blocks (max, sum) use split+unroll rather than vectorized
  partial reduction via rfactor, because TVM's rfactor primitive
  requires the reduction block to be the first child of its enclosing
  loop — incompatible with compute_at when multiple blocks share one
  spatial loop. A follow-up PR will register RVV reduction intrinsics
  (vfredmax/vfredusum) and use tensorize to vectorize reductions.

Testing

pytest tests/python/s_tir/dlight/test_cpu_reduction.py -v

• 14 shape x operator applicability tests
• 2 TIR structure verification tests
• 4 RVV codegen quality tests (code size, exp vectorization, instruction count)
• Existing test_cpu_gemv.py unaffected (10/10 pass)

[gemini-code-assist]

Code Review

This pull request introduces a new CPU reduction schedule rule for DLight, specifically targeting operators like softmax, layer norm, and RMS norm. The implementation focuses on parallelizing leading spatial axes, vectorizing injective blocks, and applying split-and-unroll strategies to reduction blocks to optimize LLVM backend performance, particularly for RISC-V Vector (RVV) targets. Feedback suggests improving the rule by dynamically inferring data type bit widths to support half-precision formats and enabling support for dynamic shapes by removing restrictive integer checks on loop extents.

[swjng]

Follow-up plan: Vectorized reduction via RVV intrinsics

This PR improves loop structure and vectorizes injective blocks, but reduction blocks (max, sum) still use scalar accumulation with split+unroll.

If this PR is okay to be merged, I would submit a follow-up PR that vectorizes the reductions themselves:

1. Register RVV reduction intrinsics in s_tir/tensor_intrin/riscv_cpu.py:

   • vfredmax.vs (vector float max reduction)
   • vfredusum.vs (vector float unordered sum reduction)

2. Use tensorize in the Reduction schedule rule to map reduction blocks to these intrinsics,
   bypassing the rfactor first-child constraint discovered during this PR's development.

[tlopex]

A few concerns comments:

1. The exception handling around llvm_get_vector_width is too broad. The docstring says it returns -1 on failure, but that case is not handled here, and except Exception may silently hide unrelated issues.

2. num_leading_s is inferred from only the first reduction block. The current check does not verify that all relevant blocks have the same number of leading spatial axes, which could make the later scheduling invalid.

3. The last block is vectorized in Phase 2 before Phase 3 applies compute_at. Please confirm that this transformation order is safe and does not invalidate the earlier vectorization.

[swjng]

@tlopex Thanks for the comment! All three points addressed in the latest push.

1. Exception handling: Replaced except Exception with a return-value check (<= 0). Verified that llvm_get_vector_width returns -1 for non-LLVM targets and never throws for valid LLVM targets.

2. Leading spatial axes: Now computed as min(leading_S) across ALL blocks rather than reading from just the first reduction block. In softmax-like patterns, reduction blocks have dom_kind=SR (1 leading S) while injective blocks have dom_kind=SS (2 leading S). Taking the minimum ensures the fuse/parallel scope is valid for every block that will be compute_at'd into it.

3. Phase 2 -> 3 order: Confirmed safe. Phase 2 vectorizes only the inner loop of the last block (splitting loop[-1] into [outer, vec]), while Phase 3 compute_at moves other blocks under spatial (the outermost loop of the last block). These operate on disjoint parts of the loop nest. Verified with bit-exact correctness checks across 6 shape × 2 operator combinations.

Also extracted _vectorize_inner and _unroll_reduction_inner as static methods to reduce apply() complexity.