GPU Fundamentals

const std = @import("std");
const builtin = @import("builtin");

pub fn main() !void {
    // Query the compile-time target information to inspect the environment
    // this binary is being compiled for (host or cross-compilation target)
    const target = builtin.target;

    // Display basic target information: CPU architecture, OS, and object format
    std.debug.print("host architecture: {s}\n", .{@tagName(target.cpu.arch)});
    std.debug.print("host operating system: {s}\n", .{@tagName(target.os.tag)});
    std.debug.print("default object format: {s}\n", .{@tagName(target.ofmt)});

    // Check if we're compiling for a GPU backend by examining the target CPU architecture.
    // GPU architectures include AMD GCN, NVIDIA PTX variants, and SPIR-V targets.
    const is_gpu_backend = switch (target.cpu.arch) {
        .amdgcn, .nvptx, .nvptx64, .spirv32, .spirv64 => true,
        else => false,
    };
    std.debug.print("compiling as GPU backend: {}\n", .{is_gpu_backend});

    // Import address space types for querying GPU-specific memory capabilities
    const AddressSpace = std.builtin.AddressSpace;
    const Context = AddressSpace.Context;

    // Query whether the target supports GPU-specific address spaces:
    // - shared: memory shared within a workgroup/threadblock
    // - constant: read-only memory optimized for uniform access across threads
    const shared_ok = target.cpu.supportsAddressSpace(AddressSpace.shared, null);
    const constant_ok = target.cpu.supportsAddressSpace(AddressSpace.constant, Context.constant);

    std.debug.print("supports shared address space: {}\n", .{shared_ok});
    std.debug.print("supports constant address space: {}\n", .{constant_ok});

    // Construct a custom target query for SPIR-V 64-bit targeting Vulkan
    const gpa = std.heap.page_allocator;
    const query = std.Target.Query{
        .cpu_arch = .spirv64,
        .os_tag = .vulkan,
        .abi = .none,
    };
    
    // Convert the target query to a triple string (e.g., "spirv64-vulkan")
    const triple = try query.zigTriple(gpa);
    defer gpa.free(triple);
    std.debug.print("example SPIR-V triple: {s}\n", .{triple});
}

//! GPU Kernel: Coordinate Capture
//!
//! This module demonstrates a minimal SPIR-V compute kernel that captures GPU dispatch
//! coordinates into a storage buffer. It shows how to use Zig's GPU-specific builtins
//! and address space annotations to write kernels that compile to SPIR-V.

const builtin = @import("builtin");

/// Represents GPU dispatch coordinates for a single invocation
/// 
/// Uses `extern` layout to guarantee memory layout matches host-side expectations,
/// ensuring the kernel's output can be safely interpreted by CPU code reading the buffer.
const Coordinates = extern struct {
    /// Work group ID (which group this invocation belongs to)
    group: u32,
    /// Work group size (number of invocations per group in this dimension)
    group_size: u32,
    /// Local invocation ID within the work group (0 to group_size-1)
    local: u32,
    /// Global linear ID across all invocations (group * group_size + local)
    linear: u32,
};

/// GPU kernel entry point that captures dispatch coordinates
///
/// This function must be exported so the SPIR-V compiler generates an entry point.
/// The `callconv(.kernel)` calling convention tells Zig to emit GPU-specific function
/// attributes and handle parameter passing according to compute shader ABI.
///
/// Parameters:
///   - out: Pointer to storage buffer where coordinates will be written.
///          The `.storage_buffer` address space annotation ensures proper
///          memory access patterns for device-visible GPU memory.
pub export fn captureCoordinates(out: *addrspace(.storage_buffer) Coordinates) callconv(.kernel) void {
    // Query the work group ID in the X dimension (first dimension)
    // @workGroupId is a GPU-specific builtin that returns the current work group's coordinate
    const group = @workGroupId(0);
    
    // Query the work group size (how many invocations per group in this dimension)
    // This is set at dispatch time by the host and queried here for completeness
    const group_size = @workGroupSize(0);
    
    // Query the local invocation ID within this work group (0 to group_size-1)
    // @workItemId is the per-work-group thread index
    const local = @workItemId(0);
    
    // Calculate global linear index across all invocations
    // This formula converts 2D coordinates (group, local) to a flat 1D index
    const linear = group * group_size + local;

    // Write all captured coordinates to the output buffer
    // The GPU will ensure this write is visible to the host after synchronization
    out.* = .{
        .group = group,
        .group_size = group_size,
        .local = local,
        .linear = linear,
    };
}

// Compile-time validation to ensure this module is only compiled for SPIR-V targets
// This prevents accidental compilation for CPU architectures where GPU builtins are unavailable
comptime {
    switch (builtin.target.cpu.arch) {
        // Accept both 32-bit and 64-bit SPIR-V architectures
        .spirv32, .spirv64 => {},
        // Reject all other architectures with a helpful error message
        else => @compileError("captureCoordinates must be compiled with a SPIR-V target, e.g. -target spirv32-vulkan-none"),
    }
}

//! GPU Dispatch Planning Utility
//! 
//! This module demonstrates how to calculate workgroup dispatch parameters for GPU compute shaders.
//! It shows the relationship between total work items, workgroup size, and the resulting dispatch
//! configuration, including handling of "tail" elements that don't fill a complete workgroup.

const std = @import("std");

/// Represents a complete dispatch configuration for parallel execution
/// Contains all necessary parameters to launch a compute kernel or parallel task
const DispatchPlan = struct {
    /// Size of each workgroup (number of threads/invocations per group)
    workgroup_size: u32,
    /// Number of workgroups needed to cover all items
    group_count: u32,
    /// Total invocations including padding (always a multiple of workgroup_size)
    padded_invocations: u32,
    /// Number of padded/unused invocations in the last workgroup
    tail: u32,
};

/// Computes optimal dispatch parameters for a given problem size and workgroup configuration
/// 
/// Calculates how many workgroups are needed to process all items, accounting for the fact
/// that the last workgroup may be partially filled. This is essential for GPU compute shaders
/// where work must be dispatched in multiples of the workgroup size.
fn computeDispatch(total_items: u32, workgroup_size: u32) DispatchPlan {
    // Ensure workgroup size is valid (GPU workgroups cannot be empty)
    std.debug.assert(workgroup_size > 0);
    
    // Calculate number of workgroups needed, rounding up to ensure all items are covered
    const groups = std.math.divCeil(u32, total_items, workgroup_size) catch unreachable;
    
    // Calculate total invocations including padding (GPU always launches complete workgroups)
    const padded = groups * workgroup_size;
    
    return .{
        .workgroup_size = workgroup_size,
        .group_count = groups,
        .padded_invocations = padded,
        // Tail represents wasted invocations that must be handled with bounds checks
        .tail = padded - total_items,
    };
}

/// Simulates CPU-side parallel execution planning using the same dispatch logic
/// 
/// Demonstrates that the workgroup dispatch formula applies equally to CPU thread batching,
/// ensuring consistent behavior between GPU and CPU fallback implementations.
fn simulateCpuFallback(total_items: u32, lanes: u32) DispatchPlan {
    // Reuse the GPU formula so host-side chunking matches device scheduling.
    return computeDispatch(total_items, lanes);
}

pub fn main() !void {
    // Define a sample problem: processing 1000 items
    const problem_size: u32 = 1000;
    
    // Typical GPU workgroup size (often 32, 64, or 256 depending on hardware)
    const workgroup_size: u32 = 64;
    
    // Calculate GPU dispatch configuration
    const plan = computeDispatch(problem_size, workgroup_size);
    std.debug.print(
        "gpu dispatch: {d} groups × {d} lanes => {d} invocations (tail {d})\n",
        .{ plan.group_count, plan.workgroup_size, plan.padded_invocations, plan.tail },
    );

    // Simulate CPU fallback with fewer parallel lanes
    const fallback_threads: u32 = 16;
    const cpu = simulateCpuFallback(problem_size, fallback_threads);
    std.debug.print(
        "cpu chunks: {d} batches × {d} lanes => {d} logical tasks (tail {d})\n",
        .{ cpu.group_count, cpu.workgroup_size, cpu.padded_invocations, cpu.tail },
    );
}