warp-lang 1.5.0__py3-none-manylinux2014_x86_64.whl → 1.6.0__py3-none-manylinux2014_x86_64.whl

warp/examples/benchmarks/benchmark_tile.py

@@ -1,179 +0,0 @@
-# Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
-# NVIDIA CORPORATION and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION is strictly prohibited.
-
-import numpy as np
-import torch
-
-import warp as wp
-
-wp.init()
-wp.set_module_options({"enable_backward": False, "fast_math": True})
-wp.set_device("cuda:0")
-
-wp.build.clear_kernel_cache()
-
-
-@wp.kernel
-def gemm(A: wp.array2d(dtype=float), B: wp.array2d(dtype=float), C: wp.array2d(dtype=float)):
-    # output index
-    i, j = wp.tid()
-
-    sum = float(0.0)
-
-    for k in range(0, A.shape[1]):
-        sum += A[i, k] * B[k, j]
-
-    C[i, j] = sum
-
-
-TILE_M = wp.constant(64)
-TILE_N = wp.constant(64)
-TILE_K = wp.constant(8)
-
-
-@wp.kernel
-def gemm_tiled(A: wp.array2d(dtype=float), B: wp.array2d(dtype=float), C: wp.array2d(dtype=float)):
-    # output tile index
-    i, j = wp.tid()
-
-    sum = wp.tile_zeros(m=TILE_M, n=TILE_N, dtype=wp.float32)
-
-    _M = A.shape[0]
-    _N = B.shape[1]
-    K = A.shape[1]
-
-    count = int(K / 8)  # TODO: code-gen bug if you use a constant before passing it to a kwd arg (in this case TILE_K)
-
-    for k in range(count):
-        a = wp.tile_load(A, i, k, m=TILE_M, n=TILE_K)
-        b = wp.tile_load(B, k, j, m=TILE_K, n=TILE_N)
-
-        # sum += a*b
-        wp.tile_matmul(a, b, sum)
-
-    wp.tile_store(C, i, j, sum)
-
-
-def benchmark_numpy(A, B, C):
-    timers = {}
-    iters = 10
-
-    # warm up
-    for _i in range(10):
-        _C = A @ B
-
-    with wp.ScopedTimer("NumPy", dict=timers):
-        for _i in range(iters):
-            _C = A @ B
-
-    return min(timers["NumPy"])
-
-
-def benchmark_warp_simt(A, B, C):
-    timers = {}
-    iters = 10
-
-    A_wp = wp.array(A)
-    B_wp = wp.array(B)
-    C_wp = wp.array(C)
-
-    # warm up
-    for _i in range(10):
-        wp.launch(gemm, dim=(M, N), inputs=[A_wp, B_wp, C_wp])
-
-    with wp.ScopedTimer("Warp (SIMT)", dict=timers, print=False, synchronize=True):
-        for _i in range(iters):
-            wp.launch(gemm, dim=(M, N), inputs=[A_wp, B_wp, C_wp])
-
-    return min(timers["Warp (SIMT)"])
-
-
-def benchmark_warp_tiled(A, B, C):
-    timers = {}
-    iters = 10
-
-    # must match with the tile_matmul() partition size
-    SUB_TILE_M = 4
-    SUB_TILE_N = 4
-
-    num_threads = int(TILE_M / SUB_TILE_M) * int(TILE_N / SUB_TILE_N)
-    A_wp = wp.array(A)
-    B_wp = wp.array(B)
-    C_wp = wp.array(C)
-
-    # warm up
-    wp.capture_begin()
-
-    for _i in range(iters):
-        wp.launch(gemm_tiled, dim=(int(M / TILE_M), int(N / TILE_N)), inputs=[A_wp, B_wp, C_wp], tile_size=num_threads)
-
-    graph = wp.capture_end()
-
-    with wp.ScopedTimer("Warp (Tiled)", dict=timers, print=False, synchronize=True):
-        # for i in range(iters):
-        #    wp.launch(gemm_tiled, dim=(int(M/TILE_M), int(N/TILE_N)), inputs=[A_wp, B_wp, C_wp], tile_size=num_threads)
-        wp.capture_launch(graph)
-
-    return min(timers["Warp (Tiled)"])
-
-
-def benchmark_torch(A, B, C):
-    A_tc = torch.from_numpy(A).to("cuda:0")
-    B_tc = torch.from_numpy(B).to("cuda:0")
-    C_tc = torch.from_numpy(C).to("cuda:0")
-
-    # warm-up
-    for _i in range(10):
-        torch.matmul(A_tc, B_tc, out=C_tc)
-
-    timers = {}
-    iters = 10
-
-    torch.cuda.synchronize()
-
-    with wp.ScopedTimer("Torch", dict=timers, print=False):
-        for _i in range(iters):
-            torch.matmul(A_tc, B_tc)  # , out=C_tc)
-
-        torch.cuda.synchronize()
-
-    return min(timers["Torch"])
-
-
-results_torch = []
-results_warp_simt = []
-results_warp_tiled = []
-
-print("{:>8s} {:>8s} {:>8s} {:>8s} {:>8s} {:>8s}".format("M", "N", "K", "Torch", "Warp (SIMT)", "Warp (Tiled)"))
-print("--------------------------------------------------------")
-
-for i in range(2, 33):
-    # for i in range(8,9):
-
-    M = i * 128
-    N = M
-    K = N
-
-    # M = TILE_M*21
-    # K = TILE_K*7
-    # N = TILE_M*12
-
-    rng = np.random.default_rng(42)
-
-    A = rng.random((M, K), dtype=np.float32)
-    B = rng.random((K, N), dtype=np.float32)
-    C = np.zeros((M, N), dtype=np.float32)
-
-    results_torch.append(benchmark_torch(A, B, C))
-    results_warp_simt.append(0.0)  # benchmark_warp_simt(A, B, C))
-    results_warp_tiled.append(benchmark_warp_tiled(A, B, C))
-
-    print(
-        "{:>8d} {:>8d} {:>8d} {:>8f} {:>8f} {:>8f}".format(
-            M, N, K, results_torch[-1], results_warp_simt[-1], results_warp_tiled[-1]
-        )
-    )

warp/native/tile_gemm.h

@@ -1,341 +0,0 @@
-/** Copyright (c) 2024 NVIDIA CORPORATION.  All rights reserved.
- * NVIDIA CORPORATION and its licensors retain all intellectual property
- * and proprietary rights in and to this software, related documentation
- * and any modifications thereto.  Any use, reproduction, disclosure or
- * distribution of this software and related documentation without an express
- * license agreement from NVIDIA CORPORATION is strictly prohibited.
- */
-
-#pragma once
-
-#include "builtin.h"
-
-#define USE_CUTE 0
-
-#if USE_CUTE
-#include "cutlass/include/cute/tensor.hpp"
-#include "cutlass/include/cute/algorithm/cooperative_gemm.hpp"
-#endif // USE_CUTE
-
-namespace wp
-{
-
-/*
-// 2D tile zero
-template <typename T, int M, int N, int Index>
-inline CUDA_CALLABLE array_t<T> tile_zeros()
-{
-    const int length = M*N;
-
-    WP_TILE_SHARED __align__(16) T data[length];
-    
-    WP_PRAGMA_UNROLL
-    for (int t=threadIdx.x; t < length; t += blockDim.x)
-    {  
-        data[t] = T(0.0);
-    }
-
-    return array_t<T>(data, M, N, nullptr);
-}
-
-// 2D tile load
-template <typename T, int M, int N, int Index>
-inline CUDA_CALLABLE array_t<T> tile_load(const array_t<T>& src, int i, int j)
-{
-    const int length = M*N;
-
-    WP_TILE_SHARED __align__(16) T data[length];
-    
-    //---------------
-    // naive-synchronous load
-    //
-    // WP_PRAGMA_UNROLL
-    // for (int t=threadIdx.x; t < length; t += blockDim.x)
-    // {  
-    //     data[t] = index(src, i*M + t/N, j*N + t%N);
-    // }
-
-    //---------------
-    // async 128 bit loads (assumes row-major i.e.: stride 1 on y axis and 4-element alignment on dimension)
-    const int s = 4;
-
-    WP_PRAGMA_UNROLL
-    for (int t=threadIdx.x*s; t < length; t += blockDim.x*s)
-    {  
-        __pipeline_memcpy_async(&data[t],
-                                &index(src, i*M + t/N, j*N + t%N),
-                                sizeof(T)*s);
-    }
-
-    __pipeline_commit();
-
-
-    return array_t<T>(data, M, N, nullptr);
-}
-
-// 2D tile store
-template <typename T>
-inline CUDA_CALLABLE void tile_store(array_t<T>& dest, int i, int j, const array_t<T>& src)
-{
-    const int M = src.shape[0];
-    const int N = src.shape[1];
-    
-    const int length = M*N;
-
-    // cooperatively store the tile, using a block-stride iterator
-    WP_PRAGMA_UNROLL
-    for (int t=threadIdx.x; t < length; t += blockDim.x)
-    {  
-        index(dest, i*M + t/N, j*N + t%N) = src.data[t];
-    }
-}
-*/
-
-template <typename T>
-inline CUDA_CALLABLE const T& index(const T* __restrict__ p, int i, int j, int stride)
-{
-    return p[i*stride + j];
-}
-
-template <typename T>
-inline CUDA_CALLABLE T& index(T* __restrict__ p, int i, int j, int stride)
-{
-    return p[i*stride + j];
-}
-
-template <unsigned M, unsigned N, typename T>
-struct partition_t
-{
-    inline partition_t(array_t<T> A)
-    {
-        data = A;
-        
-        // todo: do ceil div for non-multiples of M,N
-        shape[0] = A.shape[0]/M;
-        shape[1] = A.shape[1]/N;
-    }
-
-    // underlying data
-    array_t<T> data;
-    
-    // partition dimensions
-    int shape[2];
-};
-
-template <unsigned M, unsigned N, typename T>
-inline int partition_size(const partition_t<M, N, T>& tile)
-{
-    return tile.shape[0]*tile.shape[1];
-}
-
-// returns the x, y coordinates of a tile given a linear index
-template <unsigned M, unsigned N, typename T>
-inline void partition_coord(const partition_t<M, N, T>& tile, const int t, int& i, int& j)
-{
-    i = t/tile.shape[1];
-    j = t%tile.shape[1];
-}
-
-template <unsigned M, unsigned N, typename T>
-inline mat_t<M, N, T> partition_load(const partition_t<M, N, T>& tile, int i, int j)
-{
-    mat_t<M, N, T> out;
-    
-    const int tile_i = i*M;
-    const int tile_j = j*N;
-
-    WP_PRAGMA_UNROLL
-    for (int i=0; i < M; ++i)
-    {
-        WP_PRAGMA_UNROLL
-        for (int j=0; j < N; ++j)
-        {
-            out.data[i][j] = index(tile.data, tile_i + i, tile_j + j);
-        }
-    }
-
-    return out;
-}
-
-template <unsigned M, unsigned N, typename T>
-inline void partition_store(const partition_t<M, N, T>& tile, int i, int j, const mat_t<M, N, T>& value)
-{
-    mat_t<M, N, T> out;
-
-    const int tile_i = M*i;
-    const int tile_j = N*j;
-
-    WP_PRAGMA_UNROLL
-    for (int i=0; i < M; ++i)
-    {	
-        WP_PRAGMA_UNROLL
-        for (int j=0; j < N; ++j)
-        {
-            index(tile.data, tile_i + i, tile_j + j) = value.data[i][j];
-        }
-    }
-}
-
-
-#if !USE_CUTE
-
-template <typename T>
-inline CUDA_CALLABLE void gemm(const array_t<T>& A, const array_t<T>& B, const array_t<T>& out)
-{   
-    const int TILE_M = 4;
-    const int TILE_N = 4;
-    const int TILE_K = 4;
-
-    partition_t A_tile = partition_t<TILE_M, TILE_K, T>(A);
-    partition_t B_tile = partition_t<TILE_K, TILE_N, T>(B);
-    partition_t C_tile = partition_t<TILE_M, TILE_N, T>(out);
-
-    const int length = partition_size(C_tile);
-
-    __pipeline_wait_prior(0);
-
-    WP_TILE_SYNC();
-
-    for (int t=threadIdx.x; t < length; t += blockDim.x)
-    {  
-        int i, j;
-        partition_coord(C_tile, t, i, j);
-
-        // accumulator
-        mat_t<TILE_M, TILE_N, T> sum = partition_load(C_tile, i, j);
-
-        WP_PRAGMA_UNROLL
-        for (int k=0; k < A_tile.shape[1]; k++)
-        {
-            const mat_t<TILE_M, TILE_K, T> a = partition_load(A_tile, i, k);
-            const mat_t<TILE_K, TILE_N, T> b = partition_load(B_tile, k, j);
-
-            sum += mul(a, b);
-        }
-        
-        partition_store(C_tile, i, j, sum);
-    }
-
-    WP_TILE_SYNC();
-}
-
-
-// 2D gemm accumulate out += A*B
-template <typename TileA, typename TileB, typename TileC>
-inline CUDA_CALLABLE void tile_matmul_scalar(const TileA& A,
-                                             const TileB& B,
-                                             TileC& out)
-{    
-    const int length = tile_size(out);
-
-    WP_TILE_SYNC();
-
-    using T = typename TileA::Type;
-
-    WP_PRAGMA_UNROLL
-    for (int t=threadIdx.x; t < length; t += WP_TILE_BLOCK_DIM)
-    {  
-        // compute output index
-        const int i = t/out.N;
-        const int j = t%out.N;
-
-        T sum(0.0);
-
-        WP_PRAGMA_UNROLL
-        for (int k=0; k < A.N; ++k)
-        {
-            T a = A(i,k);
-            T b = B(k,j);
-
-            sum += a*b; // todo: use fmaf() 
-        }
-        
-        out(i,j) += sum;
-    }
-
-    WP_TILE_SYNC();
-}
-
-#else
-
-
-template <typename T>
-inline CUDA_CALLABLE void tile_matmul(const array_t<T>& A, const array_t<T>& B, const array_t<T>& out)
-{
-	using namespace cute;
-
-    __pipeline_wait_prior(0);
-
-    // ensure smem tile is ready
- 	WP_TILE_SYNC();
-
-	// Define CTA matrix size (static)
-	auto bM = Int<64>{};
-	auto bN = Int<64>{};
-	auto bK = Int<8>{};
-
-	// Define the smem layouts (static)
-	auto sA = make_layout(make_shape(bM, bK), LayoutRight{});  
-	auto sB = make_layout(make_shape(bN, bK));   
-	auto sC = make_layout(make_shape(bM, bN), LayoutRight{});
-
-    Tensor s_a_tensor = make_tensor(make_smem_ptr<float>(A.data), sA);
-    Tensor s_b_tensor = make_tensor(make_smem_ptr<float>(B.data), sB);
-    Tensor s_c_tensor = make_tensor(make_smem_ptr<float>(out.data), sC);
-
-
-    // TiledMMA tiled_mma = make_tiled_mma(UniversalFMA<float,float,float>{},
-    //                              Layout<Shape<_16,_8,_1>>{});  // 16x8x1 UniversalFMA, assumes blockDim=128
-
-
-    // TiledMMA tiled_mma = make_tiled_mma(UniversalFMA<float,float,float>{},
-    //                                     Layout<Shape<_8,_16>,Stride<_16,_1>>{});  // 8x16x1 UniversalFMA, assumes blockDim=128
-
-
-
-    TiledMMA tiled_mma = make_tiled_mma(UniversalFMA<float,float,float>{},
-                                        Layout<Shape<_2,_64>,Stride<_64,_1>>{});  // 8x16x1 UniversalFMA, assumes blockDim=128
-
-
-    cooperative_gemm< AutoVectorizingCopyWithAssumedAlignment<sizeof_bits_v<float>>,
-                      AutoVectorizingCopyWithAssumedAlignment<sizeof_bits_v<float>>, 
-                      AutoVectorizingCopyWithAssumedAlignment<sizeof_bits_v<float>>
-                    >(
-      threadIdx.x, tiled_mma,
-      1.0f, s_a_tensor, s_b_tensor, 1.0f, s_c_tensor,
-      cute::identity(), cute::identity(), cute::identity(), cute::identity()
-    );
-
-    WP_TILE_SYNC();
-
-}
-
-#endif // USE_CUTE
-
-
-#if 0
-
-template <typename TileA, typename TileB, typename TileC>
-void tile_matmul(TileA& a, TileB& b, TileC& c)
-{
-    static_assert(wp::is_same<typename TileA::Type, typename TileB::Type>::value, "Error, tile datatypes must match");
-    static_assert(TileA::N == TileB::M, "Error, inner dimensions must match");
-    static_assert(TileC::M == TileA::M, "Error, first output dimension must match");
-    static_assert(TileC::N == TileB::N, "Error, second output dimension must match");
-   
-    tile_matmul_scalar(a, b, c);
-}
-
-
-template <typename TileA, typename TileB, typename TileC,
-          typename AdjTileA, typename AdjTileB, typename AdjTileC>
-void adj_tile_matmul(TileA& a, TileB& b, TileC& c,
-                     AdjTileA& adj_a, AdjTileB& adj_b, AdjTileC& adj_c)
-{
-    tile_matmul_scalar(adj_c, wp::tile_transpose(b), adj_a);
-    tile_matmul_scalar(wp::tile_transpose(a), adj_c, adj_b);
-}
-
-#endif // 0
-
-} // namespace wp