Skip to content

Improve CPU kernel vectorization by hoisting launch state #2177

Description

@c0d1f1ed

Description

The task loop that Warp generates around every CPU kernel reads the launch bounds and the kernel arguments through pointers on every task:

for (size_t task_index = 0; task_index < dim->size; ++task_index)
{
    kernel_cpu_kernel_forward<wp_linear_contiguous>(*dim, task_index, _wp_args);
}

The kernel body stores through array data pointers that LLVM cannot prove distinct from *dim and *_wp_args, so after every store it reloads dim->size and the array descriptors (data pointer, shape, strides) of every argument before the next task. For small kernel bodies this reload traffic is a large share of the work. The loop also cannot be vectorized: wp.tid() truncates the size_t counter to int, and byte_offset() multiplies that int by the stride, so array addresses are not affine in the induction variable and the loop vectorizer stops before it reaches a cost decision. This holds for the linear-contiguous specialization as well, whose whole purpose is to make the elementwise case cheap.

Proposed change

Read the launch state once per launch instead of once per task, in both the forward and the backward task-loop templates in warp/_src/codegen.py:

  1. Copy *dim and *_wp_args (and *_wp_adj_args for the backward pass) into local objects at the top of the task loop and pass references to the copies to the per-task body. The body never writes to the launch bounds or to the argument structs, so this is exact. Private copies are not reachable from any pointer the body stores through, which lets LLVM keep the descriptors in registers across tasks.
  2. On the linear-contiguous path, iterate with a signed 32-bit counter: for (int task_index = 0; task_index < (int)size; ++task_index), passing (size_t)task_index to the body. A signed induction variable keeps the array addresses affine in the task index, which is what the loop vectorizer needs. The dispatch guard that selects the specialization gains dim->size <= 0x7fffffff, so launches whose linear size does not fit a signed 32-bit integer take the generic size_t loop as today. wp.tid() already returns int, so no coordinate that the body can observe changes.
  3. Kernels without the linear-contiguous specialization keep one generic size_t loop, with the same per-launch copies.

Activity

  1. added this to the 1.20.0 milestone on Oct 28, 2026
  2. self-assigned this
    on Oct 28, 2026
  3. modified the milestones: 1.20.0, 1.21.0 on Nov 1, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

feature requestRequest for something to be added

Projects

No projects

    Milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions