warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.1__py3-none-manylinux_2_34_aarch64.whl

warp/native/tile_reduce.h

@@ -24,6 +24,20 @@
 namespace wp
 {
 
+
+template <typename T>
+int argmax_tracker(T champion_value, T current_value, int champion_index, int current_index)
+{
+    return current_value > champion_value ? current_index : champion_index;
+}
+
+template <typename T>
+int argmin_tracker(T champion_value, T current_value, int champion_index, int current_index)
+{
+    return current_value < champion_value ? current_index : champion_index;
+}
+    
+
 #if defined(__CUDA_ARCH__)
 
 template <typename T>
@@ -62,6 +76,32 @@ inline CUDA_CALLABLE T warp_shuffle_down(T val, int offset, int mask)
   return output;
 }
 
+// Vector overload
+template <unsigned Length, typename T>
+inline CUDA_CALLABLE wp::vec_t<Length, T> warp_shuffle_down(wp::vec_t<Length, T> val, int offset, int mask)
+{
+    wp::vec_t<Length, T> result;
+
+    for (unsigned i=0; i < Length; ++i)
+        result[i] = __shfl_down_sync(mask, val[i], offset, WP_TILE_WARP_SIZE);
+    
+    return result;
+}
+
+// Matrix overload
+template <unsigned Rows, unsigned Cols, typename T>
+inline CUDA_CALLABLE wp::mat_t<Rows, Cols, T> warp_shuffle_down(wp::mat_t<Rows, Cols, T> val, int offset, int mask)
+{
+    wp::mat_t<Rows, Cols, T> result;
+
+    for (unsigned i=0; i < Rows; ++i)
+        for (unsigned j=0; j < Cols; ++j)
+            result.data[i][j] = __shfl_down_sync(mask, val.data[i][j], offset, WP_TILE_WARP_SIZE);
+    
+    return result;
+}
+
+
 template <typename T, typename Op>
 inline CUDA_CALLABLE T warp_reduce(T val, Op f, unsigned int mask)
 {
@@ -89,6 +129,52 @@ inline CUDA_CALLABLE T warp_reduce(T val, Op f, unsigned int mask)
     return sum;
 }
 
+template <typename T>
+struct ValueAndIndex
+{
+    T value;
+    int index;
+};
+
+template <typename T, typename Op, typename OpTrack>
+inline CUDA_CALLABLE ValueAndIndex<T> warp_reduce_tracked(T val, int idx, Op f, OpTrack track, unsigned int mask)
+{
+    T sum = val;
+    int index = idx;
+
+    if (mask == 0xFFFFFFFF)
+    {
+        // handle case where entire warp is active
+        for (int offset=WP_TILE_WARP_SIZE/2; offset > 0; offset /= 2)
+        {
+            auto shfl_val = warp_shuffle_down(sum, offset, mask);
+            int shfl_idx = warp_shuffle_down(index, offset, mask);            
+            index = track(sum, shfl_val, index, shfl_idx);
+            sum = f(sum, shfl_val);
+        }
+    }
+    else
+    {
+        // handle partial warp case
+        for (int offset=WP_TILE_WARP_SIZE/2; offset > 0; offset /= 2)
+        {            
+            T shfl_val = warp_shuffle_down(sum, offset, mask);
+            int shfl_index = warp_shuffle_down(index, offset, mask);
+            if ((mask & (1 << ((threadIdx.x + offset)%WP_TILE_WARP_SIZE))) != 0)
+            {
+                index = track(sum, shfl_val, index, shfl_index);
+                sum = f(sum, shfl_val);
+            }
+        }
+    }
+
+    ValueAndIndex<T> result;
+    result.value = sum;
+    result.index = index;
+
+    return result;
+}
+
 // non-axis version which computes sum 
 // across the entire tile using the whole block
 template <typename Tile, typename Op>
@@ -120,6 +206,7 @@ auto tile_reduce_impl(Op f, Tile& t)
 
     // ensure that only threads with at least one valid item participate in the reduction
     unsigned int mask = __ballot_sync(__activemask(), Layout::valid(Layout::linear_from_register(0)));
+    bool warp_is_active = mask != 0;
 
     // warp reduction
     T warp_sum = warp_reduce(thread_sum, f, mask);
@@ -135,7 +222,7 @@ auto tile_reduce_impl(Op f, Tile& t)
     // ensure active_warps is initialized
     WP_TILE_SYNC();
 
-    if (lane_index == 0)
+    if (lane_index == 0 && warp_is_active)
     {
         partials[warp_index] = warp_sum;
         atomicAdd(&active_warps, 1);
@@ -159,6 +246,86 @@ auto tile_reduce_impl(Op f, Tile& t)
     return output;
 }
 
+
+// non-axis version which computes sum 
+// across the entire tile using the whole block
+template <typename Tile, typename Op, typename OpTrack>
+auto tile_arg_reduce_impl(Op f, OpTrack track, Tile& t)
+{
+    using T = typename Tile::Type;
+
+    auto input = t.copy_to_register();
+    auto output = tile_register_t<int, tile_layout_register_t<tile_shape_t<1>>>();
+
+    const int warp_count = (WP_TILE_BLOCK_DIM + WP_TILE_WARP_SIZE - 1)/WP_TILE_WARP_SIZE;
+    const int warp_index = threadIdx.x/WP_TILE_WARP_SIZE;
+    const int lane_index = threadIdx.x%WP_TILE_WARP_SIZE;
+
+    using Layout = typename decltype(input)::Layout;
+
+    int champion_index = Layout::NumRegs > 0 ? Layout::linear_from_register(0) : -1;
+    T thread_sum = input.data[0];
+
+    // thread reduction
+    WP_PRAGMA_UNROLL
+    for (int i=1; i < Layout::NumRegs; ++i)
+    {
+        int linear = Layout::linear_from_register(i);
+        if (!Layout::valid(linear))
+            break;
+
+        champion_index = track(thread_sum, input.data[i], champion_index, linear);
+        thread_sum = f(thread_sum, input.data[i]);        
+    }
+
+    // ensure that only threads with at least one valid item participate in the reduction
+    unsigned int mask = __ballot_sync(__activemask(), Layout::valid(Layout::linear_from_register(0)));
+    bool warp_is_active = mask != 0;
+
+    // warp reduction
+    ValueAndIndex<T> warp_sum = warp_reduce_tracked(thread_sum, champion_index, f, track, mask);
+
+    // fixed size scratch pad for partial results in shared memory
+    WP_TILE_SHARED T partials[warp_count];
+    WP_TILE_SHARED int partials_idx[warp_count];
+
+    // count of active warps
+    WP_TILE_SHARED int active_warps;
+    if (threadIdx.x == 0)
+        active_warps = 0;
+    
+    // ensure active_warps is initialized
+    WP_TILE_SYNC();
+
+    if (lane_index == 0 && warp_is_active)
+    {
+        partials[warp_index] = warp_sum.value;
+        partials_idx[warp_index] = warp_sum.index;
+        atomicAdd(&active_warps, 1);
+    }
+
+    // ensure partials are ready
+    WP_TILE_SYNC();
+
+    // reduce across block, todo: use warp_reduce() here
+    if (threadIdx.x == 0)
+    {
+        T block_sum = partials[0];
+        int block_champion_index = partials_idx[0];
+        
+        WP_PRAGMA_UNROLL
+        for (int i=1; i < active_warps; ++i)
+        {
+            block_champion_index = track(block_sum, partials[i], block_champion_index, partials_idx[i]);
+            block_sum = f(block_sum, partials[i]);
+        }
+
+        output.data[0] = block_champion_index;
+    }
+
+    return output;
+}
+
 #else
 
 // CPU implementation
@@ -171,9 +338,9 @@ auto tile_reduce_impl(Op f, Tile& t)
     auto input = t.copy_to_register();
     auto output = tile_register_t<T, tile_layout_register_t<tile_shape_t<1>>>();
 
-   using Layout = typename decltype(input)::Layout;
+    using Layout = typename decltype(input)::Layout;
 
-   T sum = input.data[0];
+    T sum = input.data[0];
 
     WP_PRAGMA_UNROLL
     for (int i=1; i < Layout::NumRegs; ++i)
@@ -189,6 +356,34 @@ auto tile_reduce_impl(Op f, Tile& t)
     return output;
 }
 
+template <typename Tile, typename Op, typename OpTrack>
+auto tile_arg_reduce_impl(Op f, OpTrack track, Tile& t)
+{
+    using T = typename Tile::Type;
+
+    auto input = t.copy_to_register();
+    auto output = tile_register_t<int, tile_layout_register_t<tile_shape_t<1>>>();
+
+    using Layout = typename decltype(input)::Layout;
+
+    int champion_index = Layout::NumRegs > 0 ? Layout::linear_from_register(0) : -1;
+    T sum = input.data[0];
+
+    WP_PRAGMA_UNROLL
+    for (int i=1; i < Layout::NumRegs; ++i)
+    {
+        int linear = Layout::linear_from_register(i);
+        if (!Layout::valid(linear))
+            break;
+
+        champion_index = track(sum, input.data[i], champion_index, linear);
+        sum = f(sum, input.data[i]);
+    }
+
+    output.data[0] = champion_index;
+    return output;
+}
+
 #endif // !defined(__CUDA_ARCH__)
 
 inline void adj_tile_reduce_impl() 
@@ -200,6 +395,9 @@ inline void adj_tile_reduce_impl()
 #define tile_reduce(op, t) tile_reduce_impl([](auto x, auto y) { return op(x, y);}, t)
 #define adj_tile_reduce(op, a, adj_op, adj_a, adj_ret) adj_tile_reduce_impl()
 
+#define tile_arg_reduce(op, opTrack, t) tile_arg_reduce_impl([](auto x, auto y) { return op(x, y);}, [](auto a, auto b, auto c, auto d) { return opTrack(a, b, c, d); }, t)
+#define adj_tile_arg_reduce(op, a, adj_op, adj_a, adj_ret) adj_tile_arg_reduce_impl()
+
 // convenience methods for specific reductions
 
 template <typename Tile>
@@ -214,25 +412,26 @@ void adj_tile_sum(Tile& t, Tile& adj_t, AdjTile& adj_ret)
 {
     using T = typename Tile::Type;
 
-#if !defined(__CUDA_ARCH__)
-
-    for (int i=0; i < Tile::Layout::Size; ++i)
-    {
-        adj_t(i) += adj_ret.data[0];
+    auto adj_reg = adj_ret.grad_to_register();
 
-    }
+#if !defined(__CUDA_ARCH__)
+    T scratch = adj_reg.data[0];
 #else
     // broadcast incoming adjoint to block
     WP_TILE_SHARED T scratch;
     if (WP_TILE_THREAD_IDX == 0)
-        scratch = adj_ret.data[0];
+        scratch = adj_reg.data[0];
 
     WP_TILE_SYNC();
+#endif
 
-    // broadcast scalar across input dimensions (note zero strides)
-    auto adj_ret_reg = tile_shared_t<T, tile_layout_strided_t<typename Tile::Layout::Shape, tile_stride_t<0, 0>>, false>(&scratch, nullptr).copy_to_register();
+    auto adj_ret_reg = tile_register_like<Tile>();
+    using Layout = typename decltype(adj_ret_reg)::Layout;
+    for (int i=0; i < Layout::NumRegs; ++i)
+    {
+        adj_ret_reg.data[i] += scratch;
+    }
     adj_t.grad_add(adj_ret_reg);
-#endif 
 }
 
 template <typename Tile>
@@ -261,4 +460,31 @@ void adj_tile_min(Tile& t, Tile& adj_t, AdjTile& adj_ret)
 
 
 
+template <typename Tile>
+auto tile_argmax(Tile& t)
+{
+    return tile_arg_reduce(max, argmax_tracker, t);
+}
+
+template <typename Tile, typename AdjTile>
+void adj_tile_argmax(Tile& t, Tile& adj_t, AdjTile& adj_ret)
+{
+    // todo: not implemented
+}
+
+template <typename Tile>
+auto tile_argmin(Tile& t)
+{
+    return tile_arg_reduce(min, argmin_tracker, t);
+}
+
+template <typename Tile, typename AdjTile>
+void adj_tile_argmin(Tile& t, Tile& adj_t, AdjTile& adj_ret)
+{
+    // todo: not implemented
+}
+
+
+
+
 } // namespace wp

warp/native/tile_scan.h

@@ -0,0 +1,240 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include "tile.h"
+
+#if defined(__clang__)
+// disable warnings related to C++17 extensions on CPU JIT builds
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wc++17-extensions"
+#endif
+
+namespace wp
+{    
+
+#if defined(__CUDA_ARCH__)
+
+
+template<typename T>
+inline CUDA_CALLABLE T scan_warp_inclusive(int lane, T value)
+{
+//Computes an inclusive cumulative sum
+#pragma unroll
+    for (int i = 1; i <= 32; i *= 2)
+    {
+        auto n = __shfl_up_sync(0xffffffffu, value, i, 32);
+
+        if (lane >= i)
+            value = value + n;
+    }
+    return value;
+}
+
+
+template<typename T>
+inline CUDA_CALLABLE T thread_block_scan_inclusive(int lane, int warp_index, int num_warps, T value)
+{
+    WP_TILE_SHARED T sums[1024 / WP_TILE_WARP_SIZE]; // 1024 is the maximum number of threads per block
+
+    value = scan_warp_inclusive(lane, value);
+
+    if (lane == 31)
+    {
+        sums[warp_index] = value;
+    }
+
+    WP_TILE_SYNC();
+
+    if (warp_index == 0)
+    {
+        T v = lane < num_warps ? sums[lane] : T(0);
+        v = scan_warp_inclusive(lane, v);
+        if (lane < num_warps)
+            sums[lane] = v;
+    }
+
+    WP_TILE_SYNC();
+
+    if (warp_index > 0)
+    {
+        value += sums[warp_index - 1];
+    }
+
+    return value;
+}
+
+template<typename T, bool exclusive>
+inline CUDA_CALLABLE void thread_block_scan(T* values, int num_elements)
+{
+    const int num_threads_in_block = blockDim.x;
+    const int num_iterations = (num_elements + num_threads_in_block - 1) / num_threads_in_block;
+
+    WP_TILE_SHARED T offset;
+    if (threadIdx.x == 0)
+        offset = T(0);
+
+    WP_TILE_SYNC();
+
+    const int lane = WP_TILE_THREAD_IDX % WP_TILE_WARP_SIZE;
+    const int warp_index = WP_TILE_THREAD_IDX / WP_TILE_WARP_SIZE;
+    const int num_warps = num_threads_in_block / WP_TILE_WARP_SIZE;
+
+    for (int i = 0; i < num_iterations; ++i)
+    {
+        int element_index = WP_TILE_THREAD_IDX + i * num_threads_in_block;
+        T orig_value = element_index < num_elements ? values[element_index] : T(0);
+        T value = thread_block_scan_inclusive(lane, warp_index, num_warps, orig_value);
+        if (element_index < num_elements)
+        {
+            T new_value = value + offset;
+            if constexpr (exclusive)
+                new_value -= orig_value;
+            values[element_index] = new_value;
+        }
+
+        WP_TILE_SYNC();
+
+        if (threadIdx.x == num_threads_in_block - 1)        
+            offset += value;        
+
+        WP_TILE_SYNC();
+    }
+}
+
+template<typename Tile>
+inline CUDA_CALLABLE auto tile_scan_inclusive_impl(Tile& t)
+{
+    using T = typename Tile::Type;
+    constexpr int num_elements_to_scan = Tile::Layout::Shape::size(); 
+
+    // create a temporary shared tile to hold the input values
+    WP_TILE_SHARED T smem[num_elements_to_scan];
+    tile_shared_t<T, tile_layout_strided_t<typename Tile::Layout::Shape>, false> scratch(smem, nullptr);
+
+    // copy input values to scratch space
+    scratch.assign(t);
+
+    T* values = &scratch.data(0);
+    thread_block_scan<T, false>(values, num_elements_to_scan);
+
+    auto result =  scratch.copy_to_register();
+    
+    WP_TILE_SYNC();
+
+    return result;
+}
+
+template<typename Tile>
+inline CUDA_CALLABLE auto tile_scan_exclusive_impl(Tile& t)
+{
+    using T = typename Tile::Type;
+    constexpr int num_elements_to_scan = Tile::Layout::Shape::size(); 
+
+    // create a temporary shared tile to hold the input values
+    WP_TILE_SHARED T smem[num_elements_to_scan];
+    tile_shared_t<T, tile_layout_strided_t<typename Tile::Layout::Shape>, false> scratch(smem, nullptr);
+
+    // copy input values to scratch space
+    scratch.assign(t);
+
+    T* values = &scratch.data(0);
+    thread_block_scan<T, true>(values, num_elements_to_scan);
+
+    auto result = scratch.copy_to_register();
+
+    WP_TILE_SYNC();
+
+    return result;
+}
+
+#else
+
+template<typename Tile>
+inline auto tile_scan_inclusive_impl(Tile& t)
+{
+    using T = typename Tile::Type;
+    constexpr int num_elements_to_scan = Tile::Layout::Shape::size(); 
+    
+    auto input = t.copy_to_register();
+    auto output = tile_register_like<Tile>();
+
+    using Layout = typename decltype(input)::Layout;
+
+    T sum = T(0);
+    for (int i = 0; i < num_elements_to_scan; ++i)
+    {
+        sum += input.data[i];
+        output.data[i] = sum;
+    }
+
+    return output;
+}
+
+template<typename Tile>
+inline auto tile_scan_exclusive_impl(Tile& t)
+{
+    using T = typename Tile::Type;
+    constexpr int num_elements_to_scan = Tile::Layout::Shape::size(); 
+    
+    auto input = t.copy_to_register();
+    auto output = tile_register_like<Tile>();
+
+    using Layout = typename decltype(input)::Layout;
+
+    T sum = T(0);
+    for (int i = 0; i < num_elements_to_scan; ++i)
+    {
+        output.data[i] = sum;
+        sum += input.data[i];
+    }
+
+    return output;
+}
+
+#endif // !defined(__CUDA_ARCH__)
+
+template <typename Tile>
+auto tile_scan_inclusive(Tile& t)
+{
+    return tile_scan_inclusive_impl(t);
+}
+
+template <typename Tile, typename AdjTile>
+void adj_tile_scan_inclusive(Tile& t, Tile& adj_t, AdjTile& adj_ret)
+{
+    // todo: not implemented
+}
+
+template <typename Tile>
+auto tile_scan_exclusive(Tile& t)
+{
+    return tile_scan_exclusive_impl(t);
+}
+
+template <typename Tile, typename AdjTile>
+void adj_tile_scan_exclusive(Tile& t, Tile& adj_t, AdjTile& adj_ret)
+{
+    // todo: not implemented
+}
+
+} // namespace wp
+
+#if defined(__clang__)
+#pragma clang diagnostic pop
+#endif