Spaces:
Sleeping
Sleeping
File size: 1,846 Bytes
66c9c8a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | # 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.
import unittest
import warp as wp
from warp.tests.unittest_utils import *
wp.init()
@wp.kernel
def eval_dense_gemm(
m: int,
n: int,
p: int,
t1: int,
t2: int,
A: wp.array(dtype=float),
B: wp.array(dtype=float),
C: wp.array(dtype=float),
):
wp.dense_gemm(m, n, p, t1, t2, A, B, C)
@wp.kernel
def eval_dense_cholesky(n: int, A: wp.array(dtype=float), regularization: float, L: wp.array(dtype=float)):
wp.dense_chol(n, A, regularization, L)
@wp.kernel
def eval_dense_subs(n: int, L: wp.array(dtype=float), b: wp.array(dtype=float), x: wp.array(dtype=float)):
wp.dense_subs(n, L, b, x)
# helper that propagates gradients back to A, treating L as a constant / temporary variable
# allows us to reuse the Cholesky decomposition from the forward pass
@wp.kernel
def eval_dense_solve(
n: int, A: wp.array(dtype=float), L: wp.array(dtype=float), b: wp.array(dtype=float), x: wp.array(dtype=float)
):
wp.dense_solve(n, A, L, b, x)
def test_dense_compilation(test, device):
# just testing compilation of the dense matrix routines
# most are deprecated / WIP
wp.load_module(device=device)
devices = get_test_devices()
class TestDense(unittest.TestCase):
pass
add_function_test(TestDense, "test_dense_compilation", test_dense_compilation, devices=devices)
if __name__ == "__main__":
wp.build.clear_kernel_cache()
unittest.main(verbosity=2)
|