Merge pull request #515 from EvernightAurora/master

Feature/Ampere

setup.py

@@ -151,8 +151,8 @@ if disable_jit is not None and disable_jit == "1":
         'build_ext': PCCMBuild,
     }
     from cumm.gemm.main import GemmMainUnitTest
-    from spconv.core import SHUFFLE_SIMT_PARAMS, SHUFFLE_VOLTA_PARAMS, SHUFFLE_TURING_PARAMS
-    from spconv.core import IMPLGEMM_SIMT_PARAMS, IMPLGEMM_VOLTA_PARAMS, IMPLGEMM_TURING_PARAMS
+    from spconv.core import SHUFFLE_SIMT_PARAMS, SHUFFLE_VOLTA_PARAMS, SHUFFLE_TURING_PARAMS, SHUFFLE_AMPERE_PARAMS
+    from spconv.core import IMPLGEMM_SIMT_PARAMS, IMPLGEMM_VOLTA_PARAMS, IMPLGEMM_TURING_PARAMS, IMPLGEMM_AMPERE_PARAMS
     from cumm.conv.main import ConvMainUnitTest
     from cumm.constants import CUMM_CPU_ONLY_BUILD
     from spconv.csrc.sparse.all import SpconvOps
@@ -164,8 +164,8 @@ if disable_jit is not None and disable_jit == "1":
     from spconv.csrc.sparse.convops import ConvTunerSimple, ConvGemmOps
     from spconv.csrc.sparse.inference import InferenceOps
 
-    cu = GemmMainUnitTest(SHUFFLE_SIMT_PARAMS + SHUFFLE_VOLTA_PARAMS + SHUFFLE_TURING_PARAMS)
-    convcu = ConvMainUnitTest(IMPLGEMM_SIMT_PARAMS + IMPLGEMM_VOLTA_PARAMS + IMPLGEMM_TURING_PARAMS)
+    cu = GemmMainUnitTest(SHUFFLE_SIMT_PARAMS + SHUFFLE_VOLTA_PARAMS + SHUFFLE_TURING_PARAMS + SHUFFLE_AMPERE_PARAMS)
+    convcu = ConvMainUnitTest(IMPLGEMM_SIMT_PARAMS + IMPLGEMM_VOLTA_PARAMS + IMPLGEMM_TURING_PARAMS + IMPLGEMM_AMPERE_PARAMS)
     convcu.namespace = "cumm.conv.main"
 
     cu.namespace = "cumm.gemm.main"

spconv/build.py

@@ -22,8 +22,8 @@ from .constants import PACKAGE_NAME, PACKAGE_ROOT, DISABLE_JIT
 
 if project_is_installed(PACKAGE_NAME) and project_is_editable(
         PACKAGE_NAME) and not DISABLE_JIT:
-    from spconv.core import SHUFFLE_SIMT_PARAMS, SHUFFLE_VOLTA_PARAMS, SHUFFLE_TURING_PARAMS
-    from spconv.core import IMPLGEMM_SIMT_PARAMS, IMPLGEMM_VOLTA_PARAMS, IMPLGEMM_TURING_PARAMS
+    from spconv.core import SHUFFLE_SIMT_PARAMS, SHUFFLE_VOLTA_PARAMS, SHUFFLE_TURING_PARAMS, SHUFFLE_AMPERE_PARAMS
+    from spconv.core import IMPLGEMM_SIMT_PARAMS, IMPLGEMM_VOLTA_PARAMS, IMPLGEMM_TURING_PARAMS, IMPLGEMM_AMPERE_PARAMS
 
     from cumm.gemm.main import GemmMainUnitTest
     from cumm.conv.main import ConvMainUnitTest
@@ -38,12 +38,12 @@ if project_is_installed(PACKAGE_NAME) and project_is_editable(
     from spconv.csrc.sparse.convops import SimpleExternalSpconvMatmul
     from spconv.csrc.sparse.inference import InferenceOps
 
-    all_shuffle = SHUFFLE_SIMT_PARAMS + SHUFFLE_VOLTA_PARAMS + SHUFFLE_TURING_PARAMS
+    all_shuffle = SHUFFLE_SIMT_PARAMS + SHUFFLE_VOLTA_PARAMS + SHUFFLE_TURING_PARAMS + SHUFFLE_AMPERE_PARAMS
     all_shuffle = list(filter(lambda x: not x.is_nvrtc, all_shuffle))
     cu = GemmMainUnitTest(all_shuffle)
     cu.namespace = "cumm.gemm.main"
     all_imp = (IMPLGEMM_SIMT_PARAMS + IMPLGEMM_VOLTA_PARAMS +
-               IMPLGEMM_TURING_PARAMS)
+               IMPLGEMM_TURING_PARAMS + IMPLGEMM_AMPERE_PARAMS)
     all_imp = list(filter(lambda x: not x.is_nvrtc, all_imp))
     convcu = ConvMainUnitTest(all_imp)
     convcu.namespace = "cumm.conv.main"

spconv/gencode/__main__.py

@@ -5,9 +5,9 @@ from cumm.common import CompileInfo
 from cumm.conv.main import ConvMainUnitTest
 from cumm.gemm.main import GemmMainUnitTest
 from pccm.builder.pybind import gen_cmake
-from spconv.core import (IMPLGEMM_SIMT_PARAMS, IMPLGEMM_TURING_PARAMS,
+from spconv.core import (IMPLGEMM_SIMT_PARAMS, IMPLGEMM_TURING_PARAMS, IMPLGEMM_AMPERE_PARAMS,
                          IMPLGEMM_VOLTA_PARAMS, SHUFFLE_SIMT_PARAMS,
-                         SHUFFLE_TURING_PARAMS, SHUFFLE_VOLTA_PARAMS)
+                         SHUFFLE_TURING_PARAMS, SHUFFLE_VOLTA_PARAMS, SHUFFLE_AMPERE_PARAMS)
 from spconv.csrc.hash.core import HashTable
 from spconv.csrc.sparse.all import SpconvOps
 from spconv.csrc.sparse.alloc import ExternalAllocator, StaticAllocator
@@ -25,7 +25,7 @@ def main(include: str,
          libname: str = "spconv",
          prefix: str = "spconvlib",
          inference_only: bool = False):
-    all_shuffle = SHUFFLE_SIMT_PARAMS + SHUFFLE_VOLTA_PARAMS + SHUFFLE_TURING_PARAMS
+    all_shuffle = SHUFFLE_SIMT_PARAMS + SHUFFLE_VOLTA_PARAMS + SHUFFLE_TURING_PARAMS + SHUFFLE_AMPERE_PARAMS
     all_shuffle = list(filter(lambda x: not x.is_nvrtc, all_shuffle))
     if inference_only:
         all_shuffle = list(filter(lambda x: x.shuffle_stride != ShuffleStrideType.ShuffleAB, all_shuffle))
@@ -33,7 +33,7 @@ def main(include: str,
     cu = GemmMainUnitTest(all_shuffle)
     cu.namespace = "cumm.gemm.main"
     all_imp = (IMPLGEMM_SIMT_PARAMS + IMPLGEMM_VOLTA_PARAMS +
-               IMPLGEMM_TURING_PARAMS)
+               IMPLGEMM_TURING_PARAMS + IMPLGEMM_AMPERE_PARAMS)
     # all_imp = IMPLGEMM_SIMT_PARAMS
     all_imp = list(filter(lambda x: not x.is_nvrtc, all_imp))
     if inference_only:

test/benchmark.py

@@ -290,6 +290,49 @@ class Net2(nn.Module):
         return self.net(x)
 
 
+
+class NetSm(nn.Module):
+    def __init__(self, shape, algo):
+        super().__init__()
+        self.net = spconv.SparseSequential(
+            spconv.SubMConv3d(3,
+                              8,
+                              3,
+                              bias=False,
+                              indice_key="c0",
+                              algo=algo),
+            spconv.SubMConv3d(8,
+                              16,
+                              3,
+                              bias=False,
+                              indice_key="c0",
+                              algo=algo),
+            spconv.SubMConv3d(16,
+                              32,
+                              3,
+                              bias=False,
+                              indice_key="c0",
+                              algo=algo),
+            spconv.SubMConv3d(32,
+                              64,
+                              3,
+                              bias=False,
+                              indice_key="c0",
+                              algo=algo),
+            
+        )
+        max_batch_size = 1
+        # grid (dense map) is used for indice generation. use pre-allocated grid can run faster.
+        self.grid = torch.full([max_batch_size, *shape], -1,
+                               dtype=torch.int32).cuda()
+        # self.grid = None
+        self.shape = shape
+
+    def forward(self, features, coors, batch_size, enable_timer: bool = False):
+        x = spconv.SparseConvTensor(features, coors, self.shape, batch_size,
+                                    self.grid, enable_timer=enable_timer)
+        return self.net(x)
+
 import numpy as np
 from cumm import tensorview as tv
 from spconv.core_cc.csrc.sparse.all import SpconvOps
@@ -359,7 +402,7 @@ def main():
     # MaskImpGemm: 51.0ms
     # MaskSplitImpGemm: 41.1ms
     # algo = None
-    net = Net(spatial_shape, algo).to(device).eval().to(dtype)# .train()
+    net = NetSm(spatial_shape, algo).to(device).eval().to(dtype)# .train()
     # net.load_state_dict(net.state_dict())
     spconv.assign_name_for_sparse_modules(net)
     print(coors_th.shape)