add env to control splitk in bwd

CHANGELOG.md

@@ -4,6 +4,7 @@
 ### Added
 - add sm_37
 - add fp16 kernels witl fp32 accumulator (run slower, but can avoid nan if channel size is too large)
+- add SPCONV_BWD_SPLITK env to control splitk candidates.
 
 ## [2.1.20] - 2021-12-6
 ### Added

spconv/algo.py

@@ -20,7 +20,7 @@ from spconv.core_cc.cumm.conv.main import ConvAlgoDesp, ConvMainUnitTest, ConvPa
 from cumm.conv.bases import ConvLayout, ConvLayoutType, ConvOpType
 from cumm.gemm.algospec.core import GemmAlgo, ShuffleStrideType, get_min_arch_of_algo_str, get_available_algo_str_from_arch
 from cumm.gemm.codeops import group_by, div_up
-from spconv.constants import NDIM_DONT_CARE
+from spconv.constants import NDIM_DONT_CARE, SPCONV_BWD_SPLITK
 from typing import Optional
 import time
 from threading import Lock
@@ -345,7 +345,7 @@ class SimpleGemm:
             params.beta = beta
             params.stream = stream
             if desp.split_k_serial and hint & AlgoHint.BackwardWeight.value:
-                splitk_tests = [1, 2, 4, 8, 16, 32, 64]
+                splitk_tests = SPCONV_BWD_SPLITK
             else:
                 splitk_tests = [1]
             spk_speeds = []
@@ -646,7 +646,7 @@ class SimpleConv:
                 params.reverse_mask = reverse_mask
             params.mask_filter = mask_filter
             if desp.split_k_serial and op_type == ConvOpType.kBackwardWeight:
-                splitk_tests = [1, 2, 4, 8, 16, 32, 64]
+                splitk_tests = SPCONV_BWD_SPLITK
                 # splitk_tests = [1]
             else:
                 splitk_tests = [1]

spconv/constants.py

@@ -39,3 +39,7 @@ else:
     BOOST_ROOT = Path(_BOOST_ROOT)
     assert BOOST_ROOT.exists(), "you provide BOOST_ROOT, but it not exists"
     assert (BOOST_ROOT / "boost" / "geometry").exists(), "you provide BOOST_ROOT, but BOOST_ROOT/boost/geometry not exists"
+
+# for f16 backward weight, larger splitk, larger compute error.
+# so we use this env to control maximum splitk.
+SPCONV_BWD_SPLITK = list(map(int, os.getenv("SPCONV_BWD_SPLITK", "1,2,4,8,16,32,64").split(",")))