warp-lang 1.10.0__py3-none-macosx_11_0_arm64.whl

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)
+{
+    __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;
+
+    __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
+    __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
+    __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_a.head, adj_b.head, adj_ret.head);
+    adj_add(a.tail, b.tail, adj_a.tail, adj_b.tail, adj_ret.tail);
+}
+
+} // namespace wp