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

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

warp/native/vec.h

@@ -149,26 +149,15 @@ using vec2d = vec_t<2,double>;
 using vec3d = vec_t<3,double>;
 using vec4d = vec_t<4,double>;
 
-//--------------
-// vec<Length, Type> methods
-
-// Should these accept const references as arguments? It's all
-// inlined so maybe it doesn't matter? Even if it does, it
-// probably depends on the Length of the vector...
-
-// negation:
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE vec_t<Length, Type> operator - (vec_t<Length, Type> a)
+inline CUDA_CALLABLE vec_t<Length, Type> operator - (const vec_t<Length, Type>& x)
 {
-    // NB: this constructor will initialize all ret's components to 0, which is
-    // unnecessary... 
     vec_t<Length, Type> ret;
-    for( unsigned i=0; i < Length; ++i )
+    for(unsigned i=0; i < Length; ++i)
     {
-        ret[i] = -a[i];
+        ret[i] = -x[i];
     }
 
-    // Wonder if this does a load of copying when it returns... hopefully not as it's inlined?
     return ret;
 }
 
@@ -843,8 +832,9 @@ inline CUDA_CALLABLE void expect_near(const vec_t<Length, Type>& actual, const v
     if (diff > tolerance)
     {
         printf("Error, expect_near() failed with tolerance "); print(tolerance);
-        printf("\t Expected: "); print(expected); 
-        printf("\t Actual: "); print(actual);
+        printf("    Expected: "); print(expected); 
+        printf("    Actual: "); print(actual);
+        printf("    Max absolute difference: "); print(diff);
     }
 }
 
@@ -979,11 +969,11 @@ template<unsigned Length, typename Type>
 inline CUDA_CALLABLE void adj_div(Type s, vec_t<Length, Type> a, Type& adj_s, vec_t<Length, Type>& adj_a, const vec_t<Length, Type>& adj_ret)
 {
 
-    adj_s -= dot(a , adj_ret)/ (s * s); // - a / s^2
-
-    for( unsigned i=0; i < Length; ++i )
+    for (unsigned i=0; i < Length; ++i)
     {
-        adj_a[i] += s / adj_ret[i];
+        Type inv = Type(1) / a[i];
+        adj_a[i] -= s * adj_ret[i] * inv * inv;
+        adj_s += adj_ret[i] * inv;
     }
 
 #if FP_CHECK

warp/native/warp.cpp

@@ -24,6 +24,11 @@
 #include <stdlib.h>
 #include <string.h>
 
+// MSVC provides _aligned_malloc() instead of the standard aligned_alloc()
+#if defined(_MSC_VER)
+#include <malloc.h>
+#endif
+
 uint16_t float_to_half_bits(float x)
 {
     // adapted from Fabien Giesen's post: https://gist.github.com/rygorous/2156668
@@ -114,7 +119,7 @@ float half_bits_to_float(uint16_t u)
 int init()
 {
 #if WP_ENABLE_CUDA
-    int cuda_init();
+    int cuda_init(void);
     // note: it's safe to proceed even if CUDA initialization failed
     cuda_init();
 #endif
@@ -163,12 +168,28 @@ int is_debug_enabled()
 
 void* alloc_host(size_t s)
 {
-    return malloc(s);
+    // increase CPU array alignment for compatibility with other libs, e.g., JAX, XLA, Eigen.
+    size_t alignment = 64;
+
+    // msvc does not provide the standard aligned_alloc()
+    #if defined(_MSC_VER)
+        return _aligned_malloc(s, alignment);
+    #else
+        // ensure that the size is a multiple of alignment
+        size_t remainder = s % alignment;
+        if (remainder != 0)
+            s += alignment - remainder;
+        return aligned_alloc(alignment, s);
+    #endif
 }
 
 void free_host(void* ptr)
 {
-    free(ptr);
+    #if defined(_MSC_VER)
+        _aligned_free(ptr);
+    #else
+        free(ptr);
+    #endif
 }
 
 bool memcpy_h2h(void* dest, void* src, size_t n)
@@ -990,6 +1011,7 @@ WP_API int cuda_device_get_count() { return 0; }
 WP_API void* cuda_device_get_primary_context(int ordinal) { return NULL; }
 WP_API const char* cuda_device_get_name(int ordinal) { return NULL; }
 WP_API int cuda_device_get_arch(int ordinal) { return 0; }
+WP_API int cuda_device_get_sm_count(int ordinal) { return 0; }
 WP_API void cuda_device_get_uuid(int ordinal, char uuid[16]) {}
 WP_API int cuda_device_get_pci_domain_id(int ordinal) { return -1; }
 WP_API int cuda_device_get_pci_bus_id(int ordinal) { return -1; }
@@ -1050,10 +1072,20 @@ WP_API float cuda_event_elapsed_time(void* start_event, void* end_event) { retur
 
 WP_API bool cuda_graph_begin_capture(void* context, void* stream, int external) { return false; }
 WP_API bool cuda_graph_end_capture(void* context, void* stream, void** graph_ret) { return false; }
+WP_API bool cuda_graph_create_exec(void* context, void* stream, void* graph, void** graph_exec_ret) { return false; }
 WP_API bool cuda_graph_launch(void* graph, void* stream) { return false; }
 WP_API bool cuda_graph_destroy(void* context, void* graph) { return false; }
+WP_API bool cuda_graph_exec_destroy(void* context, void* graph_exec) { return false; }
+WP_API bool capture_debug_dot_print(void* graph, const char *path, uint32_t flags) { return false; }
+
+WP_API bool cuda_graph_insert_if_else(void* context, void* stream, int* condition, void** if_graph_ret, void** else_graph_ret) { return false; }
+WP_API bool cuda_graph_insert_while(void* context, void* stream, int* condition, void** body_graph_ret, uint64_t* handle_ret) { return false; }
+WP_API bool cuda_graph_set_condition(void* context, void* stream, int* condition, uint64_t handle) { return false; }
+WP_API bool cuda_graph_pause_capture(void* context, void* stream, void** graph_ret) { return false; }
+WP_API bool cuda_graph_resume_capture(void* context, void* stream, void* graph) { return false; }
+WP_API bool cuda_graph_insert_child_graph(void* context, void* stream, void* child_graph) { return false; }
 
-WP_API size_t cuda_compile_program(const char* cuda_src, const char* program_name, int arch, const char* include_dir, int num_cuda_include_dirs, const char** cuda_include_dirs, bool debug, bool verbose, bool verify_fp, bool fast_math, bool fuse_fp, bool lineinfo, const char* output_path, size_t num_ltoirs, char** ltoirs, size_t* ltoir_sizes, int* ltoir_input_types) { return 0; }
+WP_API size_t cuda_compile_program(const char* cuda_src, const char* program_name, int arch, const char* include_dir, int num_cuda_include_dirs, const char** cuda_include_dirs, bool debug, bool verbose, bool verify_fp, bool fast_math, bool fuse_fp, bool lineinfo, bool compile_time_trace, const char* output_path, size_t num_ltoirs, char** ltoirs, size_t* ltoir_sizes, int* ltoir_input_types) { return 0; }
 
 WP_API void* cuda_load_module(void* context, const char* ptx) { return NULL; }
 WP_API void cuda_unload_module(void* context, void* module) {}