You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
Improve CPU kernel vectorization by hoisting launch state #2177
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:
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.
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.
Kernels without the linear-contiguous specialization keep one generic size_t loop, with the same per-launch copies.
Activity