Function main

python/2_CoreConcepts/parallelReduction/parallelReduction.py:293–379

Main function demonstrating parallel reduction.

()

def main() -> int:
    """Main function demonstrating parallel reduction."""
    print("=" * 70)
    print("Parallel Reduction - Efficient GPU Array Summation")
    print("=" * 70)

    device = Device(0)
    device.set_current()

    # cuda.compute allocates temporary storage from the device's default memory
    # pool, which requires CUDA memory-pool support. This is not available on
    # every platform (for example, Windows in TCC mode).
    if not device.properties.memory_pools_supported:
        print("CUDA memory pools are not supported on this platform.")
        return 2

    stream = device.create_stream()
    cp_stream = cp.cuda.Stream.from_external(stream)

    print()
    print_gpu_info(device)

    array_size = 1 << 20  # 1M elements
    h_input = np.random.rand(array_size).astype(np.float32)
    expected_sum = float(np.sum(h_input))

    print(f"\nArray size: {array_size:,} elements ({array_size * 4 / 1e6:.1f} MB)")
    print(f"Expected sum: {expected_sum:.6f}")

    print("\nCompiling custom CUDA kernel...")
    kernel = compile_kernel(device)

    try:
        with cp_stream:
            d_input = cp.asarray(h_input)

        # ======================================================================
        # Part 1: Custom Kernel
        # ======================================================================
        print("\n" + "=" * 70)
        print("PART 1: Custom Kernel (Educational)")
        print("=" * 70)

        result, time_ms = benchmark_custom(stream, kernel, d_input)

        print(f"\nReduction tree kernel:  {result:>14.2f}")
        print(f"Expected:               {expected_sum:>14.2f}")
        print(f"Time:                   {time_ms:>14.3f} ms")

        # ======================================================================
        # Part 2: cuda.compute (Production)
        # ======================================================================
        print("\n" + "=" * 70)
        print("PART 2: cuda.compute.reduce_into() (Production)")
        print("=" * 70)

        result_cc, time_cc = benchmark_cuda_compute(stream, d_input)