CUDA batch matrix multiplication and ops (#23)

* update continuous integration * cublas-based sparse LU solver class * batched sparse cuda matrix operations * update cuda installs in ci * add test to ci * add missing files * add missing files * fix install of torch tools in ci * add missing new line * add C++ extensions to gitignore * clear lingering printf * license, move tests and utils * license, move tests * fix tests in ci * fix comparisons * rename util to random_sparse_binary_matrix, rename init_batch_size for clarity * rename util to random_sparse_binary_matrix * restore looping idiom, random_sparse_binary_matrix in toplevel __init__ * re-remove from toplevel init * cols/rows -> num_cols/num_rows * cols/rows -> num_cols/num_rows Co-authored-by: Maurizio Monge <[email protected]>
facebookresearch · Dec 21, 2021 · dbb3943 · dbb3943
1 parent ec9546c
commit dbb3943
Show file tree

Hide file tree

Showing 4 changed files with 540 additions and 0 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -117,6 +117,7 @@ run_tests: &run_tests
       working_directory: ~/project
       command: |
         pytest -s theseus/tests/test_theseus_layer.py
+        pytest -s theseus/extlib/tests/test_mat_mult.py
         pytest -s theseus/extlib/tests/test_cusolver_lu_solver.py
 
 # -------------------------------------------------------------------------------------

diff --git a/setup.py b/setup.py
@@ -29,6 +29,9 @@ def parse_requirements_file(path):
 
 if "CUDA_HOME" in os.environ:
     ext_modules = [
+        torch_cpp_ext.CUDAExtension(
+            name="theseus.extlib.mat_mult", sources=["theseus/extlib/mat_mult.cu"]
+        ),
         torch_cpp_ext.CUDAExtension(
             name="theseus.extlib.cusolver_lu_solver",
             sources=[