CUDA SAXPY Kernel
Problem Statement
Implement SAXPY — "Single-precision A·X Plus Y" — the classic BLAS Level-1 kernel: out = a * x + y for two 1D arrays and a runtime scalar a.
Background
SAXPY is the textbook fused multiply-add over vectors. Each thread does one multiply and one add, reading two inputs and writing one output.
Your Task
Implement saxpy_kernel and run(n=1024, a=2.0) returning whether the result matches a * x + 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:
n = 1024, a = 2.0
True
- The input values are
n = 1024anda = 2.0, which are used to generate two 1D arraysxandyof lengthn. - The
saxpy_kernelfunction is launched on the GPU, where each thread performs the calculationout = a * x + yfor corresponding elements inxandy. - The result is stored in the
outarray and copied back to the host, where it is compared to the reference result calculated usingnp.allclose(gpu_result, reference). - The comparison checks if the result of the GPU calculation
out = 2.0 * x + ymatches the reference result within a small tolerance, resulting in the outputTrueif the results match.
Constraints:
- out[i] = a * x[i] + y[i]
- a is a runtime scalar argument
- Bounds-check the global index
Background Knowledge
The CUDA SAXPY kernel problem involves implementing a fundamental operation in linear algebra, specifically the Single-precision A·X Plus Y (SAXPY) operation. This operation is a basic building block in many scientific computing and machine learning applications. The SAXPY operation takes two 1D arrays x and y, and a scalar a, and computes the result out = a * x + y. This operation is a key component of the Basic Linear Algebra Subprograms (BLAS) library, which provides a set of fundamental linear algebra operations.
In the context of CUDA programming, the SAXPY kernel is a parallel operation that can be executed on a Graphics Processing Unit (GPU). The kernel is launched on a 1D grid of threads, where each thread performs a single SAXPY operation on a corresponding element of the input arrays. This allows for a significant speedup over traditional CPU-based implementations, especially for large input arrays. The CUDA programming model provides a heterogeneous computing environment, where the CPU and GPU can work together to execute tasks.
To implement the SAXPY kernel, you will need to understand the basics of CUDA programming, including memory management, thread execution, and synchronization. You will also need to be familiar with the NumPy library, which provides a convenient interface for working with arrays and performing numerical computations. Additionally, you will need to understand the concept of fused multiply-add operations, which can help improve the accuracy and performance of the SAXPY kernel.
Algorithm/Approach
The general approach to solving this problem involves the following steps:
- Allocate memory for the input arrays x and y, and the output array out, on both the host (CPU) and device (GPU) sides.
- Copy the input data from the host to the device.
- Launch the SAXPY kernel on the device, using a 1D grid of threads to perform the parallel computation.
- Copy the result from the device back to the host.
- Compare the result with a reference implementation to verify its accuracy.
The SAXPY kernel itself can be implemented using a simple element-wise operation, where each thread performs a single SAXPY operation on a corresponding element of the input arrays.
Step-by-Step Strategy
To implement the solution, follow these steps:
- Import the necessary libraries, including numpy and numba.
- Define the saxpy_kernel function, which will be launched on the device.
- Allocate memory for the input arrays x and y, and the output array out, on both the host and device sides.
- Copy the input data from the host to the device.
- Launch the SAXPY kernel on the device, using a 1D grid of threads.
- Copy the result from the device back to the host.
- Compare the result with a reference implementation to verify its accuracy.
- Define the run function, which will allocate the inputs, copy them to the GPU, launch the SAXPY kernel, and return a Python bool indicating whether the result matches the reference.
Common Pitfalls
When implementing the solution, watch out for the following common pitfalls:
- Memory management: Make sure to allocate and deallocate memory correctly on both the host and device sides.
- Thread execution: Ensure that the threads are launched correctly and that the kernel is executed on the correct device.
- Synchronization: Make sure to synchronize the threads correctly to avoid data corruption or incorrect results.
- Numerical accuracy: Be aware of the potential for numerical errors due to the use of single-precision floating-point numbers.
Time & Space Complexity
The time complexity of the SAXPY kernel is O(n), where n is the length of the input arrays. This is because each thread performs a single SAXPY operation on a corresponding element of the input arrays, resulting in a linear time complexity. The space complexity is also O(n), as we need to allocate memory for the input and output arrays. Note that the parallel nature of the SAXPY kernel allows for a significant speedup over traditional CPU-based implementations, especially for large input arrays.