warp-lang 1.7.2rc1__py3-none-manylinux_2_34_aarch64.whl → 1.8.0__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,44 @@ 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.data[i] = __shfl_down_sync(mask, val.data[i], offset, WP_TILE_WARP_SIZE);
+    
+    return result;
+}
+
+// Quaternion overload
+template <typename T>
+inline CUDA_CALLABLE wp::quat_t<T> warp_shuffle_down(wp::quat_t<T> val, int offset, int mask)
+{
+    wp::quat_t<T> result;
+
+    for (unsigned i=0; i < 4; ++i)
+        result.data[i] = __shfl_down_sync(mask, val.data[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 +141,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>
@@ -159,6 +257,85 @@ 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)));
+
+    // 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)
+    {
+        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 +348,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 +366,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 +405,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>
@@ -261,4 +469,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

warp/native/tuple.h

@@ -0,0 +1,189 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2024 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
+
+namespace wp
+{
+
+template <typename... Types>
+struct tuple_t;
+
+template <>
+struct tuple_t<>
+{
+
+    static constexpr int size() { return 0; }
+
+    // Base case: empty tuple.
+    template <typename Callable>
+    void apply(Callable&&) const { }
+};
+
+template <typename Head, typename... Tail>
+struct tuple_t<Head, Tail...>
+{
+    Head head;
+    tuple_t<Tail...> tail;
+
+    CUDA_CALLABLE inline tuple_t() {}
+    CUDA_CALLABLE inline tuple_t(Head h, Tail... t) : head(h), tail(t...) {}
+
+    static constexpr int size() { return 1 + tuple_t<Tail...>::size(); }
+
+    // Applies a callable to each element.
+    template <typename Callable>
+    void apply(Callable&& func) const
+    {
+        func(head);        // Apply the callable to the current element.
+        tail.apply(func);  // Recursively process the rest of the tuple.
+    }
+};
+
+// Tuple constructor.
+template <typename... Args>
+CUDA_CALLABLE inline tuple_t<Args...>
+tuple(
+    Args... args
+)
+{
+    return tuple_t<Args...>(args...);
+}
+
+// Helper to extract a value from the tuple.
+// Can be replaced with simpler member function version when our CPU compiler
+// backend supports constexpr if statements.
+template <int N, typename Head, typename... Tail>
+struct tuple_get
+{
+    static CUDA_CALLABLE inline const auto&
+    value(
+        const tuple_t<Head, Tail...>& t
+    )
+    {
+        return tuple_get<N - 1, Tail...>::value(t.tail);
+    }
+};
+
+// Specialization for the base case N == 0. Simply return the head of the tuple.
+template <typename Head, typename... Tail>
+struct tuple_get<0, Head, Tail...>
+{
+    static CUDA_CALLABLE inline const auto&
+    value(
+        const tuple_t<Head, Tail...>& t
+    )
+    {
+        return t.head;
+    }
+};
+
+template <int Index, typename... Args>
+CUDA_CALLABLE inline auto
+extract(
+    const tuple_t<Args...>& t
+)
+{
+    return tuple_get<Index, Args...>::value(t);
+}
+
+template <typename... Args>
+CUDA_CALLABLE inline int
+len(
+    const tuple_t<Args...>& t
+)
+{
+    return t.size();
+}
+
+template <typename... Args>
+CUDA_CALLABLE inline void
+adj_len(
+    const tuple_t<Args...>& t,
+    tuple_t<Args...>& adj_t,
+    int adj_ret
+)
+{
+}
+
+template <typename... Args>
+CUDA_CALLABLE inline void
+print(
+    const tuple_t<Args...>& t
+)
+{
+    t.apply([&](auto a) { print(a); });
+}
+
+template <typename... Args>
+CUDA_CALLABLE inline void
+adj_print(
+    const tuple_t<Args...>& t,
+    tuple_t<Args...>& adj_t
+)
+{
+    adj_t.apply([&](auto a) { print(a); });
+}
+
+CUDA_CALLABLE inline tuple_t<>
+add(
+    const tuple_t<>& a,
+    const tuple_t<>& b
+)
+{
+    return tuple_t<>();
+}
+
+template <typename Head, typename... Tail>
+CUDA_CALLABLE inline tuple_t<Head, Tail...>
+add(
+    const tuple_t<Head, Tail...>& a,
+    const tuple_t<Head, Tail...>& b
+)
+{
+    tuple_t<Head, Tail...> out;
+    out.head = add(a.head, b.head);
+    out.tail = add(a.tail, b.tail);
+    return out;
+}
+
+CUDA_CALLABLE inline void
+adj_add(
+    const tuple_t<>& a,
+    const tuple_t<>& b,
+    tuple_t<>& adj_a,
+    tuple_t<>& adj_b,
+    const tuple_t<>& adj_ret
+)
+{
+}
+
+template <typename Head, typename... Tail>
+CUDA_CALLABLE inline void
+adj_add(
+    const tuple_t<Head, Tail...>& a,
+    const tuple_t<Head, Tail...>& b,
+    tuple_t<Head, Tail...>& adj_a,
+    tuple_t<Head, Tail...>& adj_b,
+    const tuple_t<Head, Tail...>& adj_ret
+)
+{
+    adj_add(a.head, b.head, adj_ret.head);
+    adj_add(a.tail, b.tail, adj_ret.tail);
+}
+
+} // namespace wp