CUDA Fused AXPBY Kernel

mediumPROCore Compute KernelsCUDA Basics with Numba

Problem Statement

Implement a fused linear combination of two vectors: out = a * x + b * y with two runtime scalars a and b.

Background

Fusing both scaled adds into one kernel reads each input once and writes the output once, saving memory bandwidth versus separate scale-and-add passes.

Your Task

Implement axpby_kernel and run(n=1024, a=2.0, b=-1.0) comparing to a * x + b * y.

How it is tested

Your solution must define a top-level function run(...) that allocates the inputs, copies them to the GPU, launches your @cuda.jit kernel, and returns a Python bool from np.allclose(gpu_result, reference). The grader prints run(...); the expected output is True.

Example

Input

n = 1024, a = 2.0, b = -1.0

Output

True

Reasoning

The input values are n = 1024, a = 2.0, and b = -1.0, which represent the size of the vectors and the scalar coefficients, respectively.
Two random vectors x and y of size n are generated, and the reference result is calculated as $a \cdot x + b \cdot y = 2.0 \cdot x + (-1.0) \cdot y$ .
The axpby_kernel function is launched on the GPU, performing the fused linear combination of x and y using the scalar coefficients a and b, and storing the result in the out vector.
The result from the GPU is compared to the reference result using np.allclose, which checks if the two arrays are element-wise equal within a tolerance, and returns True if they are equal.

Constraints

Single kernel computing ax + by
Two input arrays, two scalar arguments
Bounds-check the global index

More from CUDA Basics with Numba

CUDA Dot Producthard CUDA Atomic Max Reductionmedium CUDA Two-Stage Reduction with Atomicshard CUDA Atomic Histogrammedium CUDA Atomic Sum Reductionmedium

PRO

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

GPU · T4

Test results