#set document(title: "7.2 GPU programming fundamentals", author: "Modular Inc. / XYZ Homework") #set page(width: 8.5in, height: auto, margin: 1in) #import "@preview/cetz:0.5.2" #set text(font: ("STIX Two Text", "Libertinus Serif", "New Computer Modern"), size: 10.5pt, lang: "en") #show math.equation: set text(font: ("STIX Two Math", "New Computer Modern Math")) #set par(justify: true, leading: 0.62em, spacing: 0.9em) #set enum(spacing: 1.1em) // room between list items so tall inline fractions don't collide #set list(spacing: 1.1em) #set table(stroke: 0.5pt + rgb("#c7ccd3")) #let BLUE = rgb("#183B6F") // brand navy — section bars + example/solution labels (white on navy 11.09:1) #let ORANGE = rgb("#A94509") // brand primary-700 — AA-safe deep orange for TEXT (5.93:1 on white; raw brand #F37021 is 2.94:1 and must never carry text) #let RED = rgb("#DC2626") // brand error-600 #let GREEN = rgb("#059669") // brand success-600 (decoration only; small green text uses green-text #007942) #show heading.where(level: 1): it => block(width: 100%, above: 0pt, below: 16pt, fill: gradient.linear(BLUE, rgb("#2C5AA0")), inset: (x: 14pt, y: 12pt), radius: 3pt, text(fill: white, weight: "bold", size: 19pt, it.body)) #show heading.where(level: 2): it => block(width: 100%, above: 18pt, below: 10pt, fill: BLUE, inset: (x: 10pt, y: 6pt), radius: 2pt, text(fill: white, weight: "bold", size: 12pt, it.body)) #show heading.where(level: 3): it => text(fill: ORANGE, weight: "bold", size: 12.5pt, it.body) #show heading.where(level: 4): it => text(fill: BLUE, weight: "bold", size: 10.5pt, it.body) #let examplebox(label, title, body) = block(width: 100%, breakable: true, fill: rgb("#EFF1F5"), stroke: 0.5pt + rgb("#CFDDF0"), radius: 4pt, inset: 10pt, above: 12pt, below: 12pt)[ #block(below: 6pt)[#box(fill: BLUE, inset: (x: 6pt, y: 2pt), radius: 2pt, text(fill: white, weight: "bold", size: 8.5pt, label)) #h(0.4em) #strong[#title]] #body] // rail = decorative left rule (raw brand token); labelcolor = AA-safe label text shade #let notebox(label, rail, labelcolor, tint, body) = block(width: 100%, breakable: true, fill: tint, stroke: (left: 3pt + rail), inset: (left: 10pt, rest: 8pt), radius: (right: 4pt), above: 11pt, below: 11pt)[ #text(fill: labelcolor, weight: "bold", size: 7.5pt, tracking: 0.5pt)[#upper(label)] #linebreak() #body] #let solutionbox(body) = block(above: 4pt, below: 8pt)[ #text(fill: BLUE, weight: "bold", size: 8.5pt)[Solution] #linebreak() #body] #let figph(msg) = block(width: 100%, height: 60pt, fill: rgb("#f6f7f9"), stroke: (paint: rgb("#c7ccd3"), dash: "dashed"), radius: 4pt, inset: 10pt)[ #align(center + horizon, text(fill: rgb("#889"), style: "italic", size: 9pt, msg))] // Standardize inlined figure sizes: measure the natural CeTZ canvas, then scale to a // consistent envelope (aspect-aware; see build_typst.py FIG_* constants). Unlike the // print preamble, dimensions are FLOORED: in an editor a user can trim a figure to a // degenerate 1-D shape (a bare line), and w/h or tw/w would then divide by zero. #let _STD_W = 3.5 #let _WIDE_W = 5.6 #let _MAX_H = 3.4 #let _ASPECT_WIDE = 2.2 #let _UPSCALE_MAX = 1.15 #let stdfig(body) = context { let m = measure(body) let w = calc.max(m.width / 1in, 0.01) let h = calc.max(m.height / 1in, 0.01) let tw = if w / h > _ASPECT_WIDE { _WIDE_W } else { _STD_W } let s = calc.min(tw / w, _MAX_H / h, _UPSCALE_MAX) align(center, box(scale(x: s * 100%, y: s * 100%, reflow: true, body))) } #show figure: set block(breakable: false) #set figure(gap: 8pt) #show figure.caption: set text(size: 8.5pt, fill: rgb("#555")) == 7.2#h(0.6em)GPU programming fundamentals This guide explores the fundamentals of GPU programming using the Mojo programming language, covering essential concepts and techniques for developing GPU-accelerated applications that can work on a variety of supported GPUs from different vendors. Key topics covered in this guide: - Understanding the CPU-GPU programming model. - Working with Mojo's GPU support through the standard library. - Managing GPU devices and contexts using DeviceContext. - Writing and executing kernel functions for parallel computation. - Memory management and data transfer between CPU and GPU. - Organizing threads and thread blocks for optimal performance. Before diving into GPU programming, ensure you have a #link("https://mojolang.org/docs/requirements#gpu-compatibility")[compatible GPU] and the necessary development environment installed. And if you're new to GPU programming, we suggest you read the Intro to GPU architectures. === Overview of GPU programming in Mojo The Mojo language, including its standard library and open source MAX kernels library, allow you to develop GPU-enabled applications. ==== GPU support in the Mojo standard library The #link("https://mojolang.org/docs/std/gpu/")[gpu] package of the Mojo standard library includes several subpackages for interacting with GPUs, with the #link("https://mojolang.org/docs/std/gpu/host/")[gpu.host] package providing most of the commonly used APIs. However, the #link("https://mojolang.org/docs/std/sys/info/")[sys] package contains a few basic introspection functions for determining whether a system has a supported GPU: - #link("https://mojolang.org/docs/std/sys/info/has_accelerator/")[has\_accelerator()]: Returns True if the host system has an accelerator and False otherwise. - #link("https://mojolang.org/docs/std/sys/info/has_amd_gpu_accelerator/")[has\_amd\_gpu\_accelerator()]: Returns True if the host system has an AMD GPU and False otherwise. - #link("https://mojolang.org/docs/std/sys/info/has_apple_gpu_accelerator/")[has\_apple\_gpu\_accelerator()]: Returns True if the host system has an Apple silicon GPU and False otherwise. - #link("https://mojolang.org/docs/std/sys/info/has_nvidia_gpu_accelerator/")[has\_nvidia\_gpu\_accelerator()]: Returns True if the host system has an NVIDIA GPU and False otherwise. These functions are useful for conditional compilation or execution depending on whether a supported GPU is available. from std.sys import has\_accelerator def main(): comptime if has\_accelerator(): print("GPU detected") \# Enable GPU processing else: print("No GPU detected") \# Print error or fall back to CPU-only execution#notebox("Note", rgb("#8a94a6"), rgb("#556666"), rgb("#f7f8fa"))[ Mojo requires a #link("https://mojolang.org/docs/requirements/#gpu-compatibility")[compatible GPU development environment] to compile kernel functions, otherwise it raises a compile-time error. In this example, we're using the comptime if statement to evaluate the has\_accelerator() function at compile time and compile only the corresponding branch of the if statement. As a result, if you don't have a compatible GPU development environment, you'll see the following message when you run the program: No GPU detected ] ==== GPU programming model GPU programming follows a distinct pattern where work is divided between the CPU and GPU: - The CPU (host) manages program flow and coordinates GPU operations. - The GPU (device) executes parallel computations across many threads. - You must explicitly manage data exchange between host and device memory. A GPU program generally follows these steps: + Initialize data in host (CPU) memory. + Allocate device (GPU) memory and transfer data from host to device memory. + Execute a kernel function on the GPU to process the data. + Transfer results back from device to host memory. This process typically runs asynchronously, allowing the CPU to perform other tasks while the GPU processes data. Any time that the CPU needs to ensure that the GPU has completed an operation, such as before it copies kernel results from device memory, it must first explicitly synchronize with the GPU as described in Asynchronous operation and synchronizing the CPU and GPU. A simple example helps to understand this programming model. We'll not go into detail about the specific APIs at this point other than the included comments, but all of the types, functions, and methods are discussed in more detail in later sections of this document. from std.math import iota from std.sys import exit, has\_accelerator from std.gpu.host import DeviceContext from std.gpu import block\_dim, block\_idx, thread\_idx comptime num\_elements = 20 def scalar\_add( vector: UnsafePointer\[Float32, MutAnyOrigin\], size: Int, scalar: Float32 ): """ Kernel function to add a scalar to all elements of a vector. This kernel function adds a scalar value to each element of a vector stored in GPU memory. The input vector is modified in place. Args: vector: Pointer to the input vector. size: Number of elements in the vector. scalar: Scalar to add to the vector. """ \# Calculate the global thread index within the entire grid. Each thread \# processes one element of the vector. \# \# block\_idx.x: index of the current thread block. \# block\_dim.x: number of threads per block. \# thread\_idx.x: index of the current thread within its block. idx = block\_idx.x \* block\_dim.x + thread\_idx.x \# Bounds checking: ensure we don't access memory beyond the vector size. \# This is crucial when the number of threads doesn't exactly match vector \# size. if idx \< size: \# Each thread adds the scalar to its corresponding vector element \# This operation happens in parallel across all GPU threads vector\[idx\] += scalar def main() raises: comptime if not has\_accelerator(): print("No GPUs detected") exit(0) else: \# Initialize GPU context for device 0 (default GPU device). ctx = DeviceContext() \# Create a buffer in host (CPU) memory to store our input data host\_buffer = ctx.enqueue\_create\_host\_buffer\[DType.float32\]( num\_elements ) \# Wait for buffer creation to complete. ctx.synchronize() \# Fill the host buffer with sequential numbers (0, 1, 2, ..., size-1). iota(host\_buffer.as\_span()) print("Original host buffer:", host\_buffer) \# Create a buffer in device (GPU) memory to store data for computation. device\_buffer = ctx.enqueue\_create\_buffer\[DType.float32\](num\_elements) \# Copy data from host memory to device memory for GPU processing. ctx.enqueue\_copy(src\_buf=host\_buffer, dst\_buf=device\_buffer) \# Compile the scalar\_add kernel function for execution on the GPU. scalar\_add\_kernel = ctx.compile\_function\[scalar\_add\]() \# Launch the GPU kernel with the following arguments: \# \# - device\_buffer: GPU memory containing our vector data \# - num\_elements: number of elements in the vector \# - Float32(20.0): the scalar value to add to each element \# - grid\_dim=1: use 1 thread block \# - block\_dim=num\_elements: use 'num\_elements' threads per block (one \# thread per vector element) ctx.enqueue\_function( scalar\_add\_kernel, device\_buffer, num\_elements, Float32(20.0), grid\_dim=1, block\_dim=num\_elements, ) \# Copy the computed results back from device memory to host memory. ctx.enqueue\_copy(src\_buf=device\_buffer, dst\_buf=host\_buffer) \# Wait for all GPU operations to complete. ctx.synchronize() \# Display the final results after GPU computation. print("Modified host buffer:", host\_buffer)Original host buffer: HostBuffer(\[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, 17.0, 18.0, 19.0\]) Modified host buffer: HostBuffer(\[20.0, 21.0, 22.0, 23.0, 24.0, 25.0, 26.0, 27.0, 28.0, 29.0, 30.0, 31.0, 32.0, 33.0, 34.0, 35.0, 36.0, 37.0, 38.0, 39.0\])This application produces the following output: === Accessing and managing GPUs with DeviceContext The #link("https://mojolang.org/docs/std/gpu/host/")[gpu.host] package includes the #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceContext/")[DeviceContext] struct, which represents a logical instance of a GPU device. It provides methods for allocating memory on the device, copying data between the host CPU and the GPU, and compiling and running functions (also known as #emph[kernels]) on the device. ==== Creating an instance of DeviceContext to access a GPU Mojo supports systems with multiple GPUs. GPUs are uniquely identified by integer indices starting with 0, which is considered the "default" device. You can determine the number of GPUs available by invoking the #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceContext/#number_of_devices")[DeviceContext.number\_of\_devices()] static method. The DeviceContext() constructor returns an instance for interacting with a specified GPU. It accepts two optional arguments: - device\_id: An integer index of a specific GPU on the system. The default value of 0 refers to the "default" GPU for the system. - api: A String specifying a particular vendor's API. "cuda" (NVIDIA), "hip" (AMD), and "metal" (Apple) are currently supported. If your system doesn't have a supported GPU—or doesn't have a GPU matching the device\_id or api, if provided—then the constructor raises an error. ==== Asynchronous operation and synchronizing the CPU and GPU Typical CPU-GPU interaction is asynchronous, allowing the GPU to process tasks while the CPU is busy with other work. Each DeviceContext has an associated stream of queued operations to execute on the GPU. Operations within a stream execute in the order they are enqueued. The #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceContext/#synchronize")[synchronize()] method blocks execution of the current CPU thread until all queued operations on the associated DeviceContext stream have completed. Most commonly, you use this to wait until the result of a kernel function is copied from device memory to host memory before accessing it on the host. === Kernel functions A GPU #emph[kernel] is simply a function that runs on a GPU, executing a specific computation on a large dataset in parallel across thousands or millions of #emph[threads]. You specify the number of threads when you execute a kernel function, and all threads run the same kernel function. However, the GPU assigns a unique thread index for each thread, and you use the thread index to determine which data elements an individual thread should process. ==== Multidimensional grids and thread organization As discussed in GPU execution model, a #emph[grid] is the top-level organizational structure of the threads executing a kernel function on a GPU. A grid consists of multiple #emph[thread blocks], which are organized across one, two, or three dimensions. Each thread block is further divided into individual threads, which are in turn organized across one, two, or three dimensions. You specify the grid and thread block dimensions with the grid\_dim and block\_dim keyword arguments when you enqueue a kernel function to execute using the #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceContext/#enqueue_function")[enqueue\_function()] method. For example: \# Enqueue the print\_threads() kernel function ctx.enqueue\_function\[print\_threads\]( grid\_dim=(2, 2, 1), \# 2x2x1 blocks per grid block\_dim=(4, 4, 2), \# 4x4x2 threads per block )For both grid\_dim and block\_dim, you express the size in the x, y, and z dimensions as a #link("https://mojolang.org/docs/std/gpu/host/dim/Dim/")[Dim] or a Tuple. The y and z dimensions default to 1 if you don't explicitly provide them (that is, (2, 2) is treated as (2, 2, 1) and (8,) is treated as (8, 1, 1)). You can also provide just an Int value to specify only the x dimension (that is, 64 is treated as (64, 1, 1)). #figure(figph[Diagram: images gpu multidimensional grid], alt: "Diagram: images gpu multidimensional grid", caption: [grid.]) From within a kernel function, you can access the grid and thread block dimensions and the assigned thread block and thread indices of the individual threads executing the kernel using the following structures: #figure(table( columns: 2, align: left, inset: 6pt, table.header([Comptime value], [Description]), [#link("https://mojolang.org/docs/std/gpu/primitives/id/#grid_dim")[grid\_dim]], [Dimensions of the grid as x, y, and z values (for example, grid\_dim.y).], [#link("https://mojolang.org/docs/std/gpu/primitives/id/#block_dim")[block\_dim]], [Dimensions of the thread block as x, y, and z values.], [#link("https://mojolang.org/docs/std/gpu/primitives/id/#block_idx")[block\_idx]], [Index of the block within the grid as x, y, and z values.], [#link("https://mojolang.org/docs/std/gpu/primitives/id/#thread_idx")[thread\_idx]], [Index of the thread within the block as x, y, and z values.], [#link("https://mojolang.org/docs/std/gpu/primitives/id/#global_idx")[global\_idx]], [The global offset of the thread as x, y, and z values. That is, global\_idx.x = block\_dim.x \* block\_idx.x + thread\_idx.x, global\_idx.y = block\_dim.y \* block\_idx.y + thread\_idx.y, and global\_idx.z = block\_dim.z \* block\_idx.z + thread\_idx.z.], )) #notebox("Note", rgb("#8a94a6"), rgb("#556666"), rgb("#f7f8fa"))[ All of these dimensions and indices are #link("https://mojolang.org/docs/std/builtin/int/Int/")[Int] values. ] The following example shows a kernel function that prints the thread block index, thread index, and global index for each thread. from std.sys import has\_accelerator from std.gpu.host import DeviceContext from std.gpu import block\_dim, block\_idx, global\_idx, thread\_idx def print\_threads(): """Print thread block and thread indices.""" print( block\_idx.x, block\_idx.y, block\_idx.z, thread\_idx.x, thread\_idx.y, thread\_idx.z, global\_idx.x, global\_idx.y, global\_idx.z, block\_dim.x \* block\_idx.x + thread\_idx.x, block\_dim.y \* block\_idx.y + thread\_idx.y, block\_dim.z \* block\_idx.z + thread\_idx.z, sep="\\t", ) def main() raises: comptime if not has\_accelerator(): print("No compatible GPU found") else: \# Initialize GPU context for device 0 (default GPU device). ctx = DeviceContext() print("block\_idx\\t\\tthread\_idx\\t\\tglobal\_idx\\t\\tcalculated global\_idx") print("x\\ty\\tz", "x\\ty\\tz", "x\\ty\\tz", "x\\ty\\tz", sep="\\t") print("-" \* 20, "-" \* 20, "-" \* 20, "-" \* 20, sep="\\t") ctx.enqueue\_function\[print\_threads\]( grid\_dim=(2, 2, 1), \# 2x2x1 blocks per grid block\_dim=(4, 4, 2), \# 4x4x2 threads per block ) ctx.synchronize() print("Done")block\_idx thread\_idx global\_idx calculated global\_idx x y z x y z x y z x y z -------------------- -------------------- -------------------- -------------------- 0 1 0 0 0 0 0 4 0 0 4 0 0 1 0 1 0 0 1 4 0 1 4 0 ... 1 1 0 2 3 1 6 7 1 6 7 1 1 1 0 3 3 1 7 7 1 7 7 1 DoneThis example produces output similar to the following (with the output order indeterminate because of the concurrent execution of multiple threads): #notebox("Note", rgb("#8a94a6"), rgb("#556666"), rgb("#f7f8fa"))[ #emph[Known limitation] On Apple silicon GPUs, each print() call inside a GPU kernel currently supports at most one string literal argument. To ensure readable output across all GPU platforms, the examples here print numeric values only, with the host printing column headers before launching the kernel. This restriction doesn't apply to NVIDIA or AMD GPUs. ] ==== Writing a kernel function Kernel functions must be non-raising. #notebox("Note", rgb("#8a94a6"), rgb("#556666"), rgb("#f7f8fa"))[ For functions called from GPU kernels that need error handling, Mojo supports typed errors with types that don't allocate heap memory. See Define a custom error type for details. ] Argument values must be of types that conform to the #link("https://mojolang.org/docs/std/builtin/device_passable/DevicePassable/")[DevicePassable] trait. Additionally, a kernel function can't have a return value. Instead, you must write any result of a kernel function to a memory buffer passed in as an argument. The next two sections, Passing data between CPU and GPU and DeviceBuffer and HostBuffer go into more detail on how to pass values to a kernel function and get back results. As discussed in GPU execution model, when the GPU executes a kernel, it assigns the grid's thread blocks to various streaming multiprocessors (SMs) for execution. The SM then divides the thread block into subsets of threads called a #emph[warp]. The size of a warp depends on the GPU architecture, but most modern GPUs currently use a warp size of 32 or 64 threads. #figure(figph[Diagram: images gpu grid hierarchy], alt: "Diagram: images gpu grid hierarchy", caption: [relationship of the grid, thread blocks, warps, and individual threads, based on HIP Programming Guide ©2023-2025 Advanced Micro Devices, Inc.]) If a thread block contains a number of threads not evenly divisible by the warp size, the SM creates a partially filled final warp that still consumes the full warp's resources. For example, if a thread block has 100 threads and the warp size is 32, the SM creates: - 3 full warps of 32 threads each (96 threads total). - 1 partial warp with only 4 active threads but still occupying a full warp's worth of resources (32 thread slots). Because of this execution model, you must ensure that the threads in your kernel don't attempt to access out-of-bounds data. Otherwise, your kernel might crash or produce incorrect results. For example, if you pass a 2,000-element vector to a kernel that you execute with single-dimension thread blocks of 512 threads each, and each thread is responsible for processing one element, your kernel could perform a boundary check like this to ensure that it doesn't attempt to process out-of-bounds elements: from std.gpu import global\_idx def process\_vector(vector: UnsafePointer\[Float32, MutAnyOrigin\], size: Int): if global\_idx.x \< size: \# Process vector\[global\_idx.x\] in some way ==== Passing data between CPU and GPU All values passed to a kernel function must be of types that conform to the #link("https://mojolang.org/docs/std/builtin/device_passable/DevicePassable/")[DevicePassable] trait. The trait declares a comptime member named device\_type that maps the type as used on the CPU host to a corresponding type used on the GPU device. As an example, #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceBuffer/")[DeviceBuffer] is a host-side representation of a buffer located in the GPU's global memory space. But it defines its device\_type member as UnsafePointer, so the data represented by a DeviceBuffer is actually passed to the kernel function as a value of type #link("https://mojolang.org/docs/std/memory/unsafe_pointer/UnsafePointer/")[UnsafePointer]. The next section, DeviceBuffer and HostBuffer, describes in more detail how to allocate memory buffers on the host and device and to exchange blocks of data between host and device. The following table lists the most commonly used types in the Mojo libraries that conform to the DevicePassable trait. #figure(table( columns: 3, align: left, inset: 6pt, table.header([Host type], [Device Type], [Description]), [#link("https://mojolang.org/docs/std/builtin/int/Int/")[Int]], [Int], [Signed integer], [#link("https://mojolang.org/docs/std/builtin/simd/SIMD/")[SIMD\[dtype, width\]]], [SIMD\[dtype, width\]], [Small vector backed by a hardware vector element], [#link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceBuffer/")[DeviceBuffer\[dtype\]]], [\[UnsafePointer\[SIMD\[dtype, 1\]\]\](/docs/std/memory/unsafe\_pointer/UnsafePointer/)], [Memory buffer of dtype values], [#link("https://mojolang.org/docs/layout/tile_tensor/TileTensor/")[TileTensor]], [TileTensor], [A powerful abstraction for multi-dimensional data. See Using TileTensor.], )) ==== DeviceBuffer and HostBuffer This section describes how to use DeviceBuffer and HostBuffer to allocate memory on the device and host respectively, and to copy data between device and host memory. Creating a DeviceBufferThe #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceBuffer/")[DeviceBuffer] type represents a block of device memory associated with a particular DeviceContext. Specifically, the buffer is located in the device's #emph[global memory] space. As such, the buffer is accessible by all threads of all kernel functions executed by the DeviceContext. As discussed in Passing data between CPU and GPU, DeviceBuffer is the type used by the #strong[host] to allocate the buffer and to copy data between the host and device. But when you pass a DeviceBuffer to a kernel function, the argument received by the function is of type UnsafePointer. Attempting to use the DeviceBuffer type directly from within a kernel function results in an error. The #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceContext/#enqueue_create_buffer")[DeviceContext.enqueue\_create\_buffer()] method creates a DeviceBuffer associated with that DeviceContext. It accepts the data type as a compile-time #link("https://mojolang.org/docs/std/builtin/dtype/DType/")[DType] parameter and the size of the buffer as a run-time argument. So to create a buffer for 1,024 #link("https://mojolang.org/docs/std/builtin/simd/#float32")[Float32] values, you would execute: device\_buffer = ctx.enqueue\_create\_buffer\[DType.float32\](1024)As the method name implies, this method is asynchronous and enqueues the operation on the DeviceContext's associated stream of queued operations. Creating a HostBufferThe #link("https://mojolang.org/docs/std/gpu/host/device_context/HostBuffer/")[HostBuffer] type is analogous to DeviceBuffer, but represents a block of host memory associated with a particular DeviceContext. It supports methods for transferring data between host and device memory, as well as a basic set of methods for accessing data elements by index and for printing the buffer. The #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceContext/#enqueue_create_host_buffer")[DeviceContext.enqueue\_create\_host\_buffer()] method accepts the data type as a compile-time #link("https://mojolang.org/docs/std/builtin/dtype/DType/")[DType] parameter and the size of the buffer as a run-time argument and returns a HostBuffer. As with all DeviceContext methods whose name starts with enqueue\_, the method is asynchronous and returns immediately, adding the operation to the queue to be executed by the DeviceContext. Therefore, you need to call the synchronize() method to ensure that the operation has completed before you write to or read from the HostBuffer object. device\_buffer = ctx.enqueue\_create\_host\_buffer\[DType.float32\](1024) \# Synchronize to wait until buffer is created before attempting to write to it ctx.synchronize() \# Now it's safe to write to the buffer for i in range(1024): device\_buffer\[i\] = Float32(i \* i)Copying data between host and device memoryThe #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceContext/#enqueue_copy")[enqueue\_copy()] method is overloaded to support copying from host to device, device to host, or even device to device for systems that have multiple GPUs. Typically, you'll use it to copy data that you've staged in a HostBuffer to a DeviceBuffer before executing a kernel, and then from a DeviceBuffer to a HostBuffer to retrieve the results of kernel execution. The scalar\_add.mojo example in GPU programming model shows this pattern in action. In it, the kernel function does an in-place modification of the buffer it receives as an argument and then reuses the original HostBuffer to copy the results back from the device. However, you can allocate a separate DeviceBuffer and HostBuffer for the result of a kernel function if you want to retain the original data. In addition to copying data between a HostBuffer to a DeviceBuffer, you can use an #link("https://mojolang.org/docs/std/memory/unsafe_pointer/UnsafePointer/")[UnsafePointer] as the source or destination of a copy. However, the UnsafePointer must reference host memory for this operation. Attempting to use an UnsafePointer referencing device memory results in an error. For example, this is useful if you have data already staged in a data structure on the host that can expose the data through an UnsafePointer. In that case you would not need to copy the data from the data structure to a HostBuffer before copying it to the DeviceBuffer. Both DeviceBuffer and HostBuffer also include #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceBuffer/#enqueue_copy_to")[enqueue\_copy\_to()] and #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceBuffer/#enqueue_copy_from")[enqueue\_copy\_from()] methods. These are simply convenience methods that call the enqueue\_copy() method on their corresponding DeviceContext. For example, the following two method calls are interchangeable: ctx.enqueue\_copy(src\_buf=host\_buffer, dst\_buf=device\_buffer) \# Equivalent to: host\_buffer.enqueue\_copy\_to(dst=device\_buffer)Finally, as a convenience for testing or prototyping, you can use the #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceBuffer/#map_to_host")[DeviceBuffer.map\_to\_host()] method to create a host-accessible view of the device buffer's contents. This returns HostBuffer as a context manager that contains a copy of the data from the corresponding DeviceBuffer. Additionally, any modifications that you make to the HostBuffer are automatically copied back to the DeviceBuffer when the with statement exits. For example: ctx = DeviceContext() length = 1024 input\_device = ctx.enqueue\_create\_buffer\[DType.float32\](length) \# Initialize the input with input\_device.map\_to\_host() as input\_host: for i in range(length): input\_host\[i\] = Float32(i)However, you should not use this in most production code because of the bidirectional copies and synchronization. The example above is equivalent to: ctx = DeviceContext() length = 1024 input\_device = ctx.enqueue\_create\_buffer\[DType.float32\](length) input\_host = ctx.enqueue\_create\_host\_buffer\[DType.float32\](length) input\_device.enqueue\_copy\_to(input\_host) ctx.synchronize() for i in range(length): input\_host\[i\] = Float32(i) input\_host.enqueue\_copy\_to(input\_device) ctx.synchronize()Deallocating memory buffersBoth DeviceBuffer and HostBuffer are subject to Mojo's standard ownership and lifecycle mechanisms. The Mojo compiler analyzes our program to determine the last point that the owner of or a reference to an object is used and automatically adds a call to the object's destructor. This means that you don't explicitly call any method to free the memory represented by a DeviceBuffer or HostBuffer instance. See the Ownership and Intro to value lifecycle sections of the Mojo Manual for more information on Mojo value ownership and value lifecycle management, and the Death of a value section for a detailed explanation of object destruction. ==== Compiling and enqueuing a kernel function for execution The #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceContext/#compile_function")[compile\_function()] method accepts a kernel function as a compile-time parameter and then compiles it for the associated DeviceContext. Then you can enqueue the compiled kernel for execution by passing it to the #link("https://mojolang.org/docs/std/gpu/host/device_context/DeviceContext/#enqueue_function")[enqueue\_function()] method. The example in the GPU programming model demonstrated this pattern: ... scalar\_add\_kernel = ctx.compile\_function\[scalar\_add\]() ctx.enqueue\_function( scalar\_add\_kernel, device\_buffer, num\_elements, Float32(20.0), grid\_dim=1, block\_dim=num\_elements, )When using a compiled kernel function like this, you execute it by calling enqueue\_function() with the following arguments in this order: - The kernel function to execute. - Any additional arguments specified by the kernel function definition in the order specified by the function. - The grid dimensions using the grid\_dim keyword argument. - The thread block dimensions using the block\_dim keyword argument. Refer to the Multidimensional grids and thread organization section for more information on grid and thread block dimensions. #notebox("Note", rgb("#8a94a6"), rgb("#556666"), rgb("#f7f8fa"))[ #emph[Compile-time type checking] The compile\_function() and enqueue\_function() methods type-check the kernel function arguments at compile time. If an argument's type doesn't match the kernel function's parameter list, you'll get a compile-time error instead of a run-time failure. ] The advantage of compiling the kernel as a separate step is that you can execute the same compiled kernel on the same device multiple times. This avoids the overhead of compiling the kernel each time it's executed. If your application needs to execute a kernel function only once, you can use an overloaded version of enqueue\_function() that compiles the kernel and enqueues it in a single step. Therefore, the following is equivalent to the separate calls to compile\_function() and enqueue\_function() shown above (note that the kernel function is provided as a compile-time parameter in this case): ctx.enqueue\_function\[scalar\_add\]( device\_buffer, num\_elements, Float32(20.0), grid\_dim=1, block\_dim=num\_elements, )