warp-lang 1.6.1__py3-none-macosx_10_13_universal2.whl → 1.7.0__py3-none-macosx_10_13_universal2.whl

warp/native/tile_reduce.h

@@ -1,9 +1,18 @@
-/** 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.
+/*
+ * 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
@@ -15,6 +24,8 @@
 namespace wp
 {
 
+#if defined(__CUDA_ARCH__)
+
 template <typename T>
 inline CUDA_CALLABLE T warp_shuffle_down(T val, int offset, int mask)
 {
@@ -148,7 +159,39 @@ auto tile_reduce_impl(Op f, Tile& t)
     return output;
 }
 
-void adj_tile_reduce_impl() 
+#else
+
+// CPU implementation
+
+template <typename Tile, typename Op>
+auto tile_reduce_impl(Op f, Tile& t)
+{
+   using T = typename Tile::Type;
+
+    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;
+
+   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;
+
+        sum = f(sum, input.data[i]);
+    }
+
+    output.data[0] = sum;
+    return output;
+}
+
+#endif // !defined(__CUDA_ARCH__)
+
+inline void adj_tile_reduce_impl() 
 {
     // todo: general purpose reduction gradients not implemented 
 }
@@ -171,16 +214,25 @@ 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];
+
+    }
+#else
     // broadcast incoming adjoint to block
     WP_TILE_SHARED T scratch;
-    if (threadIdx.x == 0)
+    if (WP_TILE_THREAD_IDX == 0)
         scratch = adj_ret.data[0];
 
     WP_TILE_SYNC();
 
     // 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>>>(&scratch, NULL).copy_to_register();
+    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();
     adj_t.grad_add(adj_ret_reg);
+#endif 
 }
 
 template <typename Tile>

warp/native/vec.h

@@ -1,9 +1,18 @@
-/** Copyright (c) 2022 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.
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2022 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
@@ -497,37 +506,98 @@ inline CUDA_CALLABLE void adj_indexref(vec_t<Length, Type>* v, int idx,
 
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE void augassign_add(vec_t<Length, Type>& v, int idx, Type value)
+inline CUDA_CALLABLE void add_inplace(vec_t<Length, Type>& v, int idx, Type value)
 {
+#ifndef NDEBUG
+    if (idx < 0 || idx >= Length)
+    {
+        printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
     v[idx] += value;
 }
 
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE void adj_augassign_add(vec_t<Length, Type>& v, int idx, Type value,
+inline CUDA_CALLABLE void adj_add_inplace(vec_t<Length, Type>& v, int idx, Type value,
                                         vec_t<Length, Type>& adj_v, int adj_idx, Type& adj_value)
 {
+#ifndef NDEBUG
+    if (idx < 0 || idx >= Length)
+    {
+        printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
     adj_value += adj_v[idx];
 }
 
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE void augassign_sub(vec_t<Length, Type>& v, int idx, Type value)
+inline CUDA_CALLABLE void sub_inplace(vec_t<Length, Type>& v, int idx, Type value)
 {
+#ifndef NDEBUG
+    if (idx < 0 || idx >= Length)
+    {
+        printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
     v[idx] -= value;
 }
 
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE void adj_augassign_sub(vec_t<Length, Type>& v, int idx, Type value,
+inline CUDA_CALLABLE void adj_sub_inplace(vec_t<Length, Type>& v, int idx, Type value,
                                         vec_t<Length, Type>& adj_v, int adj_idx, Type& adj_value)
 {
+#ifndef NDEBUG
+    if (idx < 0 || idx >= Length)
+    {
+        printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
     adj_value -= adj_v[idx];
 }
 
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE vec_t<Length, Type> assign(vec_t<Length, Type>& v, int idx, Type value)
+inline CUDA_CALLABLE void assign_inplace(vec_t<Length, Type>& v, int idx, Type value)
+{
+#ifndef NDEBUG
+    if (idx < 0 || idx >= Length)
+    {
+        printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
+    v[idx] = value;
+}
+
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE void adj_assign_inplace(vec_t<Length, Type>& v, int idx, Type value, vec_t<Length, Type>& adj_v, int& adj_idx, Type& adj_value)
+{
+#ifndef NDEBUG
+    if (idx < 0 || idx >= Length)
+    {
+        printf("vec index %d out of bounds at %s %d\n", idx, __FILE__, __LINE__);
+        assert(0);
+    }
+#endif
+
+    adj_value += adj_v[idx];
+}
+
+
+template<unsigned Length, typename Type>
+inline CUDA_CALLABLE vec_t<Length, Type> assign_copy(vec_t<Length, Type>& v, int idx, Type value)
 {
 #ifndef NDEBUG
     if (idx < 0 || idx >= Length)
@@ -543,7 +613,7 @@ inline CUDA_CALLABLE vec_t<Length, Type> assign(vec_t<Length, Type>& v, int idx,
 }
 
 template<unsigned Length, typename Type>
-inline CUDA_CALLABLE void adj_assign(vec_t<Length, Type>& v, int idx, Type value, vec_t<Length, Type>& adj_v, int& adj_idx, Type& adj_value, const vec_t<Length, Type>& adj_ret)
+inline CUDA_CALLABLE void adj_assign_copy(vec_t<Length, Type>& v, int idx, Type value, vec_t<Length, Type>& adj_v, int& adj_idx, Type& adj_value, const vec_t<Length, Type>& adj_ret)
 {
 #ifndef NDEBUG
     if (idx < 0 || idx >= Length)
@@ -765,7 +835,7 @@ inline CUDA_CALLABLE vec_t<Length,Type> sign(vec_t<Length,Type> v)
 template<unsigned Length, typename Type>
 inline CUDA_CALLABLE void expect_near(const vec_t<Length, Type>& actual, const vec_t<Length, Type>& expected, const Type& tolerance)
 {
-    const Type diff(0);
+    Type diff(0);
     for(size_t i=0; i<Length; ++i)
     {
         diff = max(diff,abs(actual[i] - expected[i]));