test_all_algo.py 43.7 KB
Newer Older
yan.yan's avatar
yan.yan committed
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
# Copyright 2021 Yan Yan
#
# 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.

"""Test all gemm/conv kernels.
We can't test all kernels in network because auto-tuner will only find one best kernel.
"""


import sys
from pathlib import Path
from typing import Dict, List, Tuple
import pickle
import sys
import time
from pathlib import Path
from cumm.gemm.algospec.core import GemmAlgo, ShuffleStrideType

import numpy as np
import pccm
import torch
import torch.nn.functional as F
yan.yan's avatar
yan.yan committed
33
from spconv.core_cc.csrc.sparse.convops import GemmTuneResult, ConvTuneResult
yan.yan's avatar
yan.yan committed
34
from spconv.pytorch.core import SparseConvTensor
yan.yan's avatar
yan.yan committed
35
36
from spconv.test_utils import TestCase
from cumm import tensorview as tv
yan.yan's avatar
yan.yan committed
37
from spconv.constants import SPCONV_ALLOW_TF32
yan.yan's avatar
yan.yan committed
38
39
40
41
42
43
from cumm.conv.bases import NCHW, NHWC, ConvIterAlgo, ConvOpType
import os
from cumm.gemm.codeops import div_up
from spconv.core import AlgoHint, ConvAlgo
from spconv.pytorch.conv import expand_nd
from spconv.pytorch import ops
yan.yan's avatar
yan.yan committed
44
from spconv.algo import GEMM, CONV, GEMM_CPP, CONV_CPP, BestAlgoByProfile, BestConvAlgoByProfile, GemmTunerSimple
yan.yan's avatar
yan.yan committed
45
46
47
from spconv.pytorch.cppcore import get_current_stream, torch_tensor_to_tv
from spconv.test_utils import generate_sparse_data, params_grid
import tqdm 
yan.yan's avatar
yan.yan committed
48
from spconv.constants import ALL_WEIGHT_IS_KRSC, SPCONV_CPP_GEMM
yan.yan's avatar
yan.yan committed
49
50
from spconv.core_cc.csrc.sparse.inference import InferenceOps
from spconv.pytorch import functional as Fsp
yan.yan's avatar
yan.yan committed
51
assert ALL_WEIGHT_IS_KRSC is True, "we only support KRSC in spconv >= 2.2"
yan.yan's avatar
yan.yan committed
52
from spconv.pytorch.hash import HashTable
yan.yan's avatar
yan.yan committed
53
54
55
56
57
58
59
60
61
62
63
64
65
66

# TODO remove or release this when tf32 op is ready
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

NUMPY_DTYPE_TO_TORCH = {
    np.float32: torch.float32,
    np.float16: torch.float16,
    np.int8: torch.int8,

}

class SparseConvTester:
    def __init__(self, algo: ConvAlgo, subm: bool, shape: List[int], bs: int, dtype: np.dtype, N: int, K: int, C: int, 
yan.yan's avatar
yan.yan committed
67
        ksize: int, stride: int, padding: int, dilation: int, check_bias: bool = False, check_act: bool = False) -> None:
yan.yan's avatar
yan.yan committed
68
        ndim = 3
yan.yan's avatar
yan.yan committed
69
        transpose = False
yan.yan's avatar
yan.yan committed
70
71
72
73
74
75
        self.shape = shape 
        self.bs = bs 
        self.dtype = dtype 
        self.dtype_th = NUMPY_DTYPE_TO_TORCH[dtype]
        self.K = K 
        self.C = C 
yan.yan's avatar
yan.yan committed
76
77
78
79
        self.ksize = expand_nd(ndim, ksize) 
        self.stride = expand_nd(ndim, stride) 
        self.padding = expand_nd(ndim, padding, ) 
        self.dilation = expand_nd(ndim, dilation) 
yan.yan's avatar
yan.yan committed
80
81
        self.N = N
        self.device = torch.device("cuda:0")
yan.yan's avatar
yan.yan committed
82
        op = expand_nd(ndim, 0)
yan.yan's avatar
yan.yan committed
83
84
        self.kv: int = np.prod(self.ksize)
        self.num_split = 1 if algo == ConvAlgo.MaskImplicitGemm else 2
yan.yan's avatar
yan.yan committed
85
86
87
88
89
90
91
92
93
        if not subm:
            if transpose:
                out_shape = ops.get_deconv_output_size(shape, self.ksize, self.stride,
                                                self.padding, self.dilation, op)
            else:
                out_shape = ops.get_conv_output_size(shape, self.ksize, self.stride,
                                                self.padding, self.dilation)
        else:
            out_shape = shape
yan.yan's avatar
yan.yan committed
94

yan.yan's avatar
yan.yan committed
95
        sparse_dict = generate_sparse_data(shape, [N] * bs, C)
yan.yan's avatar
yan.yan committed
96
97
98
99
100
101
102
103
104

        voxels_np = np.ascontiguousarray(sparse_dict["features"]).astype(
            np.float32)
        indices_np = np.ascontiguousarray(
            sparse_dict["indices"][:, [3, 0, 1, 2]]).astype(np.int32)
        indices_th = torch.from_numpy(indices_np).to(self.device)
        out_inds, pair_ref, indice_num_per_loc = ops.get_indice_pairs(
            indices_th, 1, shape, ConvAlgo.Native, self.ksize, self.stride, self.padding,
            self.dilation, op, subm)
yan.yan's avatar
yan.yan committed
105
106
        self.ref_out_inds = out_inds
        self.ref_out_inds_scalar = Fsp._indice_to_scalar(out_inds.long(), [bs, *out_shape])
yan.yan's avatar
yan.yan committed
107
108
109
110
        self.indice_num_per_loc_np = indice_num_per_loc.cpu().numpy()
        self.indice_pairs_np = pair_ref.cpu().numpy()
        self.pair_native = pair_ref
        self.indice_num_per_loc = indice_num_per_loc
yan.yan's avatar
yan.yan committed
111
112
113
        self.use_direct_table = True
        
        self.out_shape = out_shape
yan.yan's avatar
yan.yan committed
114
115
116
117
118
119
120
121
122
123
124
125
126
        if algo == ConvAlgo.Native:
            self.out_inds: torch.Tensor = out_inds
            self.num_inds_per_loc: torch.Tensor = indice_num_per_loc
            self.pair_fwd : torch.Tensor = torch.Tensor()
            self.pair_bwd: torch.Tensor = torch.Tensor()
            self.pair_mask_fwd_splits: List[torch.Tensor] = []
            self.pair_mask_bwd_splits: List[torch.Tensor] = []
            self.mask_argsort_fwd_splits: List[torch.Tensor] = []
            self.mask_argsort_bwd_splits: List[torch.Tensor] = []
            self.masks = np.array([])
        else:
            res = ops.get_indice_pairs_implicit_gemm(indices_th, bs, shape,
                                                    algo, self.ksize, self.stride, self.padding,
yan.yan's avatar
yan.yan committed
127
                                                    self.dilation, op, subm=subm, direct_table=self.use_direct_table)
yan.yan's avatar
yan.yan committed
128
129
130
131
132
133
134
135
136
137
            
            self.out_inds = res[0]
            self.num_inds_per_loc = res[1]
            self.pair_fwd = res[2]
            self.pair_bwd = res[3]
            self.pair_mask_fwd_splits = res[4]
            self.pair_mask_bwd_splits = res[5]
            self.mask_argsort_fwd_splits = res[6]
            self.mask_argsort_bwd_splits = res[7]
            self.masks = res[8]
138
            self.mask_int_count = res[9]
yan.yan's avatar
yan.yan committed
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
        
        self.out_inds_scalar = Fsp._indice_to_scalar(self.out_inds.long(), [bs, *out_shape])

        table = HashTable(out_inds.device, torch.int64, torch.int32, self.out_inds.shape[0] * 2)
        # test coords -> test out indexes
        table.insert(self.out_inds_scalar, torch.arange(0, self.out_inds.shape[0], dtype=torch.int32, device=self.device))
        # out_order:  test_order_to_ref, test index for each ref coord
        out_order, is_empty = table.query(self.ref_out_inds_scalar)
        assert is_empty.int().sum().item() == 0, "shouldn't happen"
        self.out_order = out_order.cpu().numpy()

        # inp_table = HashTable(out_inds.device, torch.int64, torch.int32, self.ref_out_inds.shape[0] * 2)
        # inp_table.insert(self.ref_out_inds_scalar, torch.arange(0, self.ref_out_inds.shape[0], dtype=torch.int32, device=self.device))
        # # out_order:  ref index for each out coord
        # out_order, is_empty = inp_table.query(self.out_inds_scalar)


yan.yan's avatar
yan.yan committed
156
157
        self.voxels_np = voxels_np
        self.indices_np = indices_np
yan.yan's avatar
yan.yan committed
158
159
        self.check_bias = check_bias
        self.check_act = check_act
yan.yan's avatar
yan.yan committed
160
161
162
163
164
165
166
167
168
169

        self.subm = subm
        if dtype == np.int8:
            self.inp = np.random.randint(-2, 2, size=[voxels_np.shape[0],
                                                    C]).astype(np.int8)
            self.weight = np.random.randint(-2, 2, size=[K, *self.ksize,
                                                    C]).astype(np.int8)
            self.output = np.random.randint(-2, 2, size=[
                self.out_inds.shape[0], K
            ]).astype(dtype)
yan.yan's avatar
yan.yan committed
170
171
172
173
            self.bias = np.random.randint(-2, 2, size=[
                K
            ]).astype(dtype)

yan.yan's avatar
yan.yan committed
174
175
176
177
178
179
180
181
        else:
            self.inp = np.random.uniform(-1, 1, size=[
                voxels_np.shape[0], C
            ]).astype(dtype)
            self.weight = np.random.uniform(-1, 1, size=[K, *self.ksize, C]).astype(dtype)
            self.output = np.random.uniform(-1, 1, size=[
                self.out_inds.shape[0], K
            ]).astype(dtype)
yan.yan's avatar
yan.yan committed
182
183
184
185
            self.bias = np.random.uniform(-1, 1, size=[
                K
            ]).astype(dtype)

yan.yan's avatar
yan.yan committed
186
187
188
        self.weight_ref = self.weight.transpose(1, 2, 3, 0, 4)
        self.weight_ref = np.ascontiguousarray(self.weight_ref).reshape(-1, K, C)
        self.out_ref, self.din_ref, self.dw_ref = self._get_ref_output()
yan.yan's avatar
yan.yan committed
189
190
191
192
193
        if check_bias:
            self.out_ref += self.bias
            # relu
        if check_act:
            self.out_ref = np.maximum(self.out_ref, 0)
yan.yan's avatar
yan.yan committed
194
        self.dw_ref = np.ascontiguousarray(self.dw_ref.transpose(1, 0, 2).reshape(K, *self.ksize, C))
yan.yan's avatar
yan.yan committed
195
        self.arch = tv.get_compute_capability()
yan.yan's avatar
yan.yan committed
196

yan.yan's avatar
yan.yan committed
197
198
199
200
    def get_output_ref_spt(self):
        return SparseConvTensor(torch.from_numpy(self.out_ref).cuda(), self.ref_out_inds, self.out_shape, self.bs)


yan.yan's avatar
yan.yan committed
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
    def _get_ref_output(self):
        output_ref = np.zeros_like(self.output, dtype=np.float32)
        dinput_ref = np.zeros_like(self.inp, dtype=np.float32)
        dw_ref = np.zeros_like(self.weight_ref,
                                dtype=np.float32)  # KV, K, C

        for filter_offset in range(self.kv):
            if self.subm and filter_offset > self.kv // 2:
                nhot = self.indice_num_per_loc_np[self.kv - 1 - filter_offset]
            elif self.subm and filter_offset == self.kv // 2:
                nhot = self.voxels_np.shape[0]
            else:
                nhot = self.indice_num_per_loc_np[filter_offset]

            i_inds = self.indice_pairs_np[0][filter_offset][:nhot]
            o_inds = self.indice_pairs_np[1][filter_offset][:nhot]
            a = self.inp[i_inds]
            cc = a.astype(
                np.float32) @ self.weight_ref[filter_offset].T.astype(
                    np.float32)
            output_ref[o_inds] += cc
yan.yan's avatar
yan.yan committed
222
223
            # we use random output as dout here
            a = self.output[self.out_order][o_inds]
yan.yan's avatar
yan.yan committed
224
225
226
227
228
            # NK @ KC
            cc = a.astype(
                np.float32) @ self.weight_ref[filter_offset].astype(
                    np.float32)
            dinput_ref[i_inds] += cc
yan.yan's avatar
yan.yan committed
229
230
            # use random output and random inp as dout and inp
            out_gather = self.output[self.out_order][o_inds]  # [N, K]
yan.yan's avatar
yan.yan committed
231
232
233
234
235
            inp_gather = self.inp[i_inds]  # [N, C]
            # KN @ NC
            dw_res = out_gather.astype(
                np.float32).T @ inp_gather.astype(np.float32)
            dw_ref[filter_offset] = dw_res
yan.yan's avatar
yan.yan committed
236
237
        if self.dtype == np.int8:
            output_ref = np.clip(output_ref, -127, 127)
yan.yan's avatar
yan.yan committed
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
        return output_ref, dinput_ref, dw_ref

    def get_operands(self, op_type: ConvOpType):
        zeros_func = tv.zeros if not self.subm else tv.empty
        if op_type == ConvOpType.kBackwardInput:
            inp_tv = zeros_func(list(self.inp.shape), self.dtype, 0)
        else:
            inp_tv = tv.from_numpy(self.inp).cuda()
        if op_type == ConvOpType.kBackwardWeight:
            weight_tv = zeros_func(list(self.weight.shape), self.dtype, 0)
        else:
            weight_tv = tv.from_numpy(self.weight).cuda()
        if op_type == ConvOpType.kForward:
            output_tv = zeros_func(list(self.output.shape), self.dtype, 0)
        else:
            output_tv = tv.from_numpy(self.output).cuda()
        return inp_tv, weight_tv, output_tv

    def get_operands_torch(self, op_type: ConvOpType):
        zeros_func = torch.zeros if not self.subm else torch.empty
        if op_type == ConvOpType.kBackwardInput:
            inp_tv = zeros_func(list(self.inp.shape), dtype=self.dtype_th, device=self.device)
        else:
            inp_tv = torch.from_numpy(self.inp).cuda()
        if op_type == ConvOpType.kBackwardWeight:
            weight_tv = zeros_func(list(self.weight.shape), dtype=self.dtype_th, device=self.device)
        else:
            weight_tv = torch.from_numpy(self.weight).cuda()
        if op_type == ConvOpType.kForward:
            output_tv = zeros_func(list(self.output.shape), dtype=self.dtype_th, device=self.device)
        else:
            output_tv = torch.from_numpy(self.output).cuda()
        return inp_tv, weight_tv, output_tv

def _test_impgemm_conv_cuda(subm: bool):
    ndim = 3
yan.yan's avatar
yan.yan committed
274
    np.random.seed(50005)
yan.yan's avatar
yan.yan committed
275
    dtype_to_tol = {
yan.yan's avatar
yan.yan committed
276
        np.float32: (1e-2, 1e-2),
yan.yan's avatar
yan.yan committed
277
278
279
280
281
282
283
        np.float16: (1e-2, 1e-2),
        np.int8: (1e-4, 1e-4),
    }
    device = torch.device("cuda:0")
    shapes = [[19, 18, 17]]
    batchsizes = [1]
    dtypes = [np.float32, np.float16]
yan.yan's avatar
yan.yan committed
284
285
    # dtypes = [np.float16]

yan.yan's avatar
yan.yan committed
286
    # dtypes = [np.int8]
yan.yan's avatar
yan.yan committed
287
    test_case = TestCase()
yan.yan's avatar
yan.yan committed
288
289
290
291
    # in_channels = [32]
    # out_channels = [32, 48, 64]
    in_channels = [32, 47]
    out_channels = [32, 48, 62]
yan.yan's avatar
yan.yan committed
292
293
    # in_channels = [32]
    # out_channels = [32]
yan.yan's avatar
yan.yan committed
294

yan.yan's avatar
yan.yan committed
295
    multiple_base = 16
yan.yan's avatar
yan.yan committed
296
    if subm:
297
        ksizes = [3, (3, 3, 5), (3, 5, 5), 5]
yan.yan's avatar
yan.yan committed
298
299
300
301
        strides = [1]
        paddings = [0]
        dilations = [1]
    else:
302
        ksizes = [2, 3, (3, 3, 4), 4, (4, 5, 5), 5]
yan.yan's avatar
yan.yan committed
303
304
305
        strides = [1, 2, 3]
        paddings = [0, 1]
        dilations = [1, 2]
yan.yan's avatar
yan.yan committed
306

yan.yan's avatar
yan.yan committed
307
    algos = [
yan.yan's avatar
yan.yan committed
308
        # ConvAlgo.MaskSplitImplicitGemm,
yan.yan's avatar
yan.yan committed
309
310
311
        ConvAlgo.MaskImplicitGemm,
    ]
    arch = torch.cuda.get_device_capability()
yan.yan's avatar
yan.yan committed
312
    force_nvrtc = False
yan.yan's avatar
yan.yan committed
313
314
315
    for shape, bs, C, K, k, s, p, d, algo, dtype in tqdm.tqdm(params_grid(
            shapes, batchsizes, in_channels, out_channels, ksizes,
            strides, paddings, dilations, algos, dtypes)):
yan.yan's avatar
yan.yan committed
316
317
        shape_prod = np.prod(shape)
        num_batch = np.random.randint(int(0.2 * shape_prod), int(0.7 * shape_prod))
yan.yan's avatar
yan.yan committed
318
319
320
321
        # C = np.random.randint(int(0.3 * C), int(0.7 * C))
        # K = np.random.randint(int(0.3 * K), int(0.7 * K))
        multipler = max(C, K) / multiple_base
        multipler = max(multipler, 1.0)
yan.yan's avatar
yan.yan committed
322
        # print(num_batch)
yan.yan's avatar
yan.yan committed
323
324
325
326
327
        tester = SparseConvTester(algo, subm, shape, bs, dtype, num_batch, K, C, k, s, p, d, check_bias=True, check_act=True)
        bias = None
        act = tv.gemm.Activation.None_
        if tester.check_bias:
            bias = tv.from_numpy(tester.bias).cuda()
yan.yan's avatar
yan.yan committed
328
329
330
331
332
333
334
        atol, rtol = dtype_to_tol[dtype]
        mask_width_to_mask_out_fwd: Dict[int, torch.Tensor] = {}
        mask_width_to_mask_out_bwd: Dict[int, torch.Tensor] = {}
        op_types = [ConvOpType.kForward, ConvOpType.kBackwardInput]
        spk = 1
        for op_type in op_types:
            inp_tv, weight_tv, output_tv = tester.get_operands(op_type)
yan.yan's avatar
yan.yan committed
335
336
337
            if SPCONV_CPP_GEMM:
                avail_desps = CONV_CPP.get_all_available(inp_tv, weight_tv, output_tv, 
                    NHWC.layout_type.value, NHWC.layout_type.value, 
yan.yan's avatar
yan.yan committed
338
339
                    NHWC.layout_type.value, NHWC.interleave, NHWC.interleave, NHWC.interleave, arch, op_type.value, -1, True, False,
                        use_tf32=SPCONV_ALLOW_TF32)
yan.yan's avatar
yan.yan committed
340
            else:
yan.yan's avatar
yan.yan committed
341
342
                avail_desps = CONV.get_all_available(inp_tv, weight_tv, output_tv, NHWC, NHWC, NHWC, arch, op_type, -1,
                        use_tf32=SPCONV_ALLOW_TF32)
yan.yan's avatar
yan.yan committed
343
344
345
346
347
            if op_type == ConvOpType.kForward and tester.check_act:
                act = tv.gemm.Activation.ReLU
            else:
                act = tv.gemm.Activation.None_
            assert avail_desps
yan.yan's avatar
yan.yan committed
348
349
350
351
352
353
354
355
356
357
            for desp in avail_desps:
                if not subm:
                    if op_type == ConvOpType.kForward:
                        output_tv.zero_()
                    else:
                        inp_tv.zero_()
                # this algo must success
                mask_width = desp.tile_shape[0]
                # if mask_width != 32:
                #     continue
358
                
yan.yan's avatar
yan.yan committed
359
                if mask_width not in mask_width_to_mask_out_fwd:
360
                    mask_width_to_mask_out_fwd[mask_width] = torch.zeros([2, tester.mask_int_count * div_up(tester.out_inds.shape[0], mask_width)],
yan.yan's avatar
yan.yan committed
361
362
363
                                      dtype=torch.int32,
                                      device=tester.device)
                mask_output_fwd = mask_width_to_mask_out_fwd[mask_width]
yan.yan's avatar
yan.yan committed
364
365
366
367
                is_fwd = desp.op_type.value == ConvOpType.kForward.value
                bias_cur = bias 
                if op_type != ConvOpType.kForward:
                    bias_cur = None
yan.yan's avatar
yan.yan committed
368
                if subm:
yan.yan's avatar
yan.yan committed
369
                    if desp.op_type.value == ConvOpType.kForward.value:
yan.yan's avatar
yan.yan committed
370
                        indice_pairs = tester.pair_fwd
yan.yan's avatar
yan.yan committed
371
                    elif desp.op_type.value == ConvOpType.kBackwardInput.value:
yan.yan's avatar
yan.yan committed
372
373
374
375
376
377
                        indice_pairs = tester.pair_bwd
                    else:
                        indice_pairs = tester.pair_fwd
                    mask_output = mask_output_fwd
                    # print([bin(x.item()) for x in masks])
                    for j in range(tester.num_split):
yan.yan's avatar
yan.yan committed
378
379
380
381
382
                        beta = 1 if j > 0 else 0
                        if bias_cur is not None:
                            beta = 1
                        if j > 0:
                            bias_cur = None
yan.yan's avatar
yan.yan committed
383
384
                        mask_filter = tester.masks[j].item()
                        reverse_mask = False
yan.yan's avatar
yan.yan committed
385
                        if desp.op_type.value == ConvOpType.kBackwardWeight.value:
yan.yan's avatar
yan.yan committed
386
387
388
                            mask_op = mask_output[j]
                        else:
                            mask_op = tester.pair_mask_fwd_splits[j]
yan.yan's avatar
yan.yan committed
389
                        if desp.op_type.value == ConvOpType.kBackwardInput.value:
yan.yan's avatar
yan.yan committed
390
391
                            reverse_mask = True
                        mask_output_run = torch_tensor_to_tv(mask_output[j], dtype=tv.uint32)
yan.yan's avatar
yan.yan committed
392
                        if desp.op_type.value == ConvOpType.kBackwardWeight.value:
yan.yan's avatar
yan.yan committed
393
                            mask_output_run = tv.Tensor()
yan.yan's avatar
yan.yan committed
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
                        # force_nvrtc = desp.op_type.value == ConvOpType.kBackwardInput.value
                        # if force_nvrtc:
                        #     desp.is_nvrtc = True
                        # print(force_nvrtc, desp.op_type, op_type)
                        if SPCONV_CPP_GEMM:

                            CONV_CPP.run_with_tuned_result(
                                ConvTuneResult(desp, tester.arch, spk),
                                desp.op_type.value,
                                inp_tv,
                                weight_tv,
                                output_tv,
                                torch_tensor_to_tv(mask_op, dtype=tv.uint32),
                                torch_tensor_to_tv(tester.mask_argsort_fwd_splits[j]),
                                mask_output_run,
                                torch_tensor_to_tv(indice_pairs),
                                reverse_mask,
                                mask_filter=mask_filter,
                                mask_width=mask_width,
                                beta=beta,
                                verbose=False,
                                force_nvrtc=force_nvrtc,
yan.yan's avatar
yan.yan committed
416
417
                                bias=bias_cur if is_fwd and bias_cur is not None else tv.Tensor(),
                                act_type=act,
418
                                mask_int_count=tester.mask_int_count,
yan.yan's avatar
yan.yan committed
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
                            )
                        else:
                            CONV.run_with_tuned_result(
                                BestConvAlgoByProfile(desp, tester.arch, spk),
                                desp.op_type.value,
                                inp_tv,
                                weight_tv,
                                output_tv,
                                torch_tensor_to_tv(mask_op, dtype=tv.uint32),
                                torch_tensor_to_tv(tester.mask_argsort_fwd_splits[j]),
                                mask_output_run,
                                torch_tensor_to_tv(indice_pairs),
                                reverse_mask,
                                mask_filter=mask_filter,
                                mask_width=mask_width,
                                beta=beta,
                                verbose=False,
                                force_nvrtc=force_nvrtc,
yan.yan's avatar
yan.yan committed
437
438
                                bias=bias_cur if is_fwd else None,
                                act_type=act,
439
                                mask_int_count=tester.mask_int_count
yan.yan's avatar
yan.yan committed
440
441
                            )

yan.yan's avatar
yan.yan committed
442
443
444
445
446
447
448
                else:
                    if mask_width not in mask_width_to_mask_out_bwd:
                        mask_width_to_mask_out_bwd[mask_width] = torch.zeros([2, div_up(tester.indices_np.shape[0], mask_width)],
                                        dtype=torch.int32,
                                        device=tester.device)
                    mask_output_bwd = mask_width_to_mask_out_bwd[mask_width]

yan.yan's avatar
yan.yan committed
449
                    if desp.op_type.value == ConvOpType.kForward.value:
yan.yan's avatar
yan.yan committed
450
451
452
453
                        indice_pairs = tester.pair_fwd  # inp -> out
                        mask_ops = tester.pair_mask_fwd_splits
                        mask_argsorts = tester.mask_argsort_fwd_splits
                        mask_output = mask_output_fwd
yan.yan's avatar
yan.yan committed
454
                    elif desp.op_type.value == ConvOpType.kBackwardInput.value:
yan.yan's avatar
yan.yan committed
455
456
457
458
459
460
461
462
463
464
465
                        indice_pairs = tester.pair_bwd  # out -> inp
                        mask_ops = tester.pair_mask_bwd_splits
                        mask_argsorts = tester.mask_argsort_bwd_splits
                        mask_output = mask_output_bwd
                    else:
                        indice_pairs = tester.pair_fwd  # inp -> out
                        mask_ops = tester.pair_mask_fwd_splits
                        mask_argsorts = tester.mask_argsort_fwd_splits
                        mask_output = mask_output_fwd

                    for j in range(tester.num_split):
yan.yan's avatar
yan.yan committed
466
467
468
469
470
471
                        # beta = 1 if j == 1 else 0
                        beta = 1 if j > 0 else 0
                        if bias_cur is not None:
                            beta = 1
                        if j > 0:
                            bias_cur = None
yan.yan's avatar
yan.yan committed
472
473
                        mask_filter = tester.masks[j].item()
                        reverse_mask = False
yan.yan's avatar
yan.yan committed
474
                        if desp.op_type.value == ConvOpType.kBackwardWeight.value:
yan.yan's avatar
yan.yan committed
475
476
477
                            mask_op = mask_output[j]
                        else:
                            mask_op = mask_ops[j]
yan.yan's avatar
yan.yan committed
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
                        if SPCONV_CPP_GEMM:

                            CONV_CPP.run_with_tuned_result(
                                ConvTuneResult(desp, tester.arch, spk),
                                desp.op_type.value,
                                inp_tv,
                                weight_tv,
                                output_tv,
                                torch_tensor_to_tv(mask_op, dtype=tv.uint32),
                                torch_tensor_to_tv(mask_argsorts[j]),
                                torch_tensor_to_tv(mask_output[j], dtype=tv.uint32),
                                torch_tensor_to_tv(indice_pairs),
                                reverse_mask,
                                mask_filter=mask_filter,
                                mask_width=mask_width,
                                beta=beta,
                                verbose=False,
yan.yan's avatar
yan.yan committed
495
496
497
                                force_nvrtc=force_nvrtc,
                                bias=bias if is_fwd and bias is not None else tv.Tensor(),
                                act_type=act,
498
                                mask_int_count=tester.mask_int_count,
yan.yan's avatar
yan.yan committed
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
                            )
                        else:
                            CONV.run_with_tuned_result(
                                BestConvAlgoByProfile(desp, tester.arch, spk),
                                desp.op_type.value,
                                inp_tv,
                                weight_tv,
                                output_tv,
                                torch_tensor_to_tv(mask_op, dtype=tv.uint32),
                                torch_tensor_to_tv(mask_argsorts[j]),
                                torch_tensor_to_tv(mask_output[j], dtype=tv.uint32),
                                torch_tensor_to_tv(indice_pairs),
                                reverse_mask,
                                mask_filter=mask_filter,
                                mask_width=mask_width,
                                beta=beta,
                                verbose=False,
yan.yan's avatar
yan.yan committed
516
517
518
                                force_nvrtc=force_nvrtc,
                                bias=bias if is_fwd else None,
                                act_type=act,
519
                                mask_int_count=tester.mask_int_count,
yan.yan's avatar
yan.yan committed
520
                            )
yan.yan's avatar
yan.yan committed
521
522
523
524
525
526

                out_ref = tester.out_ref
                din_ref = tester.din_ref
                dw_ref = tester.dw_ref
                if op_type == ConvOpType.kForward:
                    out_my = output_tv.cpu().numpy()
yan.yan's avatar
yan.yan committed
527
                    out_my = out_my[tester.out_order]
yan.yan's avatar
yan.yan committed
528
529
530
531
                    if dtype != np.float16:
                        test_case.assertAllClose(out_ref, out_my, atol=atol, rtol=rtol)
                    else:
                        error_norm = np.linalg.norm(out_ref.reshape(-1) - out_my.reshape(-1))
yan.yan's avatar
yan.yan committed
532
533
                        if (error_norm > 5):
                            print(f"{desp}, Error={error_norm}")
yan.yan's avatar
yan.yan committed
534
                        assert error_norm < 10 * multipler
yan.yan's avatar
yan.yan committed
535
536
537
538
539
540
                else:
                    din_my = inp_tv.cpu().numpy()
                    if dtype != np.float16:
                        test_case.assertAllClose(din_ref, din_my, atol=atol, rtol=rtol)
                    else:
                        error_norm = np.linalg.norm(din_ref.reshape(-1) - din_my.reshape(-1))
yan.yan's avatar
yan.yan committed
541
                        assert error_norm < 10 * multipler, f"{desp}, {error_norm}, {k}, {s}, {p}, {d}"
yan.yan's avatar
yan.yan committed
542
543
544
        inp_tv, weight_tv, output_tv = tester.get_operands(ConvOpType.kBackwardWeight)
        for spk in [1, 4, 16, 64]:
            for mask_width, mask_output in mask_width_to_mask_out_fwd.items():
yan.yan's avatar
yan.yan committed
545
546
547
548
                if SPCONV_CPP_GEMM:
                    avail_desps = CONV_CPP.get_all_available(inp_tv, weight_tv, output_tv, 
                        NHWC.layout_type.value, NHWC.layout_type.value, 
                        NHWC.layout_type.value, NHWC.interleave, NHWC.interleave, NHWC.interleave, arch, 
yan.yan's avatar
yan.yan committed
549
550
                        ConvOpType.kBackwardWeight.value, mask_width, True, False,
                        use_tf32=SPCONV_ALLOW_TF32)
yan.yan's avatar
yan.yan committed
551
                else:
yan.yan's avatar
yan.yan committed
552
553
                    avail_desps = CONV.get_all_available(inp_tv, weight_tv, output_tv, NHWC, NHWC, NHWC, arch, ConvOpType.kBackwardWeight, mask_width,
                        use_tf32=SPCONV_ALLOW_TF32)
yan.yan's avatar
yan.yan committed
554
555
556
557
558
559
560
561
562
563
564
565
566
                for desp in avail_desps:
                    weight_tv.zero_()
                    if subm:
                        indice_pairs = tester.pair_fwd
                        for j in range(tester.num_split):
                            beta = 0
                            mask_filter = tester.masks[j].item()
                            mask_op = mask_output[j]
                            mask_op_tv = torch_tensor_to_tv(mask_op, dtype=tv.uint32)
                            # mask_op_np = mask_op_tv.cpu().numpy()
                            # bit_ref = np.bitwise_or.reduce(mask_op_np, axis=0)
                            # bit_my = mask_filter
                            CONV.run_with_tuned_result(
yan.yan's avatar
yan.yan committed
567
568
                                BestConvAlgoByProfile(desp, tester.arch, spk),
                                desp.op_type.value,
yan.yan's avatar
yan.yan committed
569
570
571
572
573
574
575
576
577
578
579
580
                                inp_tv,
                                weight_tv,
                                output_tv,
                                mask_op_tv,
                                torch_tensor_to_tv(tester.mask_argsort_fwd_splits[j]),
                                tv.Tensor(),
                                torch_tensor_to_tv(indice_pairs),
                                reverse_mask=False,
                                mask_filter=mask_filter,
                                mask_width=mask_width,
                                beta=beta,
                                verbose=False,
581
                                mask_int_count=tester.mask_int_count,
yan.yan's avatar
yan.yan committed
582
583
584
585
586
587
588
589
590
591
592
593
594
                            )
                    else:
                        indice_pairs = tester.pair_fwd  # inp -> out
                        mask_ops = tester.pair_mask_fwd_splits
                        mask_argsorts = tester.mask_argsort_fwd_splits
                        for j in range(tester.num_split):
                            # beta = 1 if j == 1 else 0
                            beta = 0
                            mask_filter = tester.masks[j].item()
                            reverse_mask = False
                            mask_op = mask_output[j]

                            CONV.run_with_tuned_result(
yan.yan's avatar
yan.yan committed
595
596
                                BestConvAlgoByProfile(desp, tester.arch, spk),
                                desp.op_type.value,
yan.yan's avatar
yan.yan committed
597
598
599
600
601
602
603
604
605
606
607
608
                                inp_tv,
                                weight_tv,
                                output_tv,
                                torch_tensor_to_tv(mask_op, dtype=tv.uint32),
                                torch_tensor_to_tv(mask_argsorts[j]),
                                torch_tensor_to_tv(mask_output[j], dtype=tv.uint32),
                                torch_tensor_to_tv(indice_pairs),
                                reverse_mask,
                                mask_filter=mask_filter,
                                mask_width=mask_width,
                                beta=beta,
                                verbose=False,
609
                                mask_int_count=tester.mask_int_count,
yan.yan's avatar
yan.yan committed
610
611
612
613
614
615
616
617
618
                            )
                    dw_ref = tester.dw_ref
                    dw_my = weight_tv.cpu().numpy()
                    if dtype != np.float16:
                        # print(desp, spk, K, C, mask_width, algo)
                        test_case.assertAllClose(dw_ref, dw_my, atol=atol, rtol=rtol)
                    else:
                        error_norm = np.linalg.norm(dw_ref.reshape(-1) - dw_my.reshape(-1))
                        # print(desp, error_norm)
yan.yan's avatar
yan.yan committed
619
                        if (error_norm > 5):
yan.yan's avatar
yan.yan committed
620
621
                            print(f"{desp}, Error={error_norm}, {spk}")
                        assert error_norm < 10 * multipler
yan.yan's avatar
yan.yan committed
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646

def _test_native_conv_cuda(subm: bool):
    ndim = 3
    dtype_to_tol = {
        np.float32: (1e-4, 1e-4),
        np.float16: (1e-2, 1e-2),
        np.int8: (1e-4, 1e-4),
    }
    device = torch.device("cuda:0")
    shapes = [[19, 18, 17]]
    batchsizes = [1]
    dtypes = [np.float32, np.float16]
    test_case = TestCase()
    in_channels = [32, 47]
    out_channels = [32, 48, 62]
    if subm:
        ksizes = [3, 5]
        strides = [1]
        paddings = [0]
        dilations = [1]
    else:
        ksizes = [2, 3]
        strides = [1, 2, 3]
        paddings = [0, 1]
        dilations = [1, 2]
yan.yan's avatar
yan.yan committed
647
648
    multiple_base = 128

yan.yan's avatar
yan.yan committed
649
650
    arch = torch.cuda.get_device_capability()
    stream = get_current_stream()
yan.yan's avatar
yan.yan committed
651
    force_nvrtc = False
yan.yan's avatar
yan.yan committed
652
653
654
    for shape, bs, C, K, k, s, p, d, dtype in tqdm.tqdm(params_grid(
            shapes, batchsizes, in_channels, out_channels, ksizes,
            strides, paddings, dilations, dtypes)):
yan.yan's avatar
yan.yan committed
655
656
657
658
        tester = SparseConvTester(ConvAlgo.Native, subm, shape, bs, dtype, 1500, K, C, k, s, p, d, check_bias=True, check_act=True)
        bias = None
        if tester.check_bias:
            bias = tv.from_numpy(tester.bias).cuda()
yan.yan's avatar
yan.yan committed
659
        atol, rtol = dtype_to_tol[dtype]
yan.yan's avatar
yan.yan committed
660
661
        multipler = max(C, K) / multiple_base
        multipler = max(multipler, 1.0)
yan.yan's avatar
yan.yan committed
662
663
664
665
666
667
668

        kv_center = tester.kv // 2
        kv = tester.kv
        pair_in = torch_tensor_to_tv(tester.pair_native)[0]
        pair_out = torch_tensor_to_tv(tester.pair_native)[1]

        op_types = [ConvOpType.kForward, ConvOpType.kBackwardInput, ConvOpType.kBackwardWeight]
yan.yan's avatar
yan.yan committed
669
670
        # op_types = [ConvOpType.kForward]

yan.yan's avatar
yan.yan committed
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
        indice_pair_num_cpu = tester.indice_num_per_loc_np
        spk = 1

        out_ref = tester.out_ref
        din_ref = tester.din_ref
        dw_ref = tester.dw_ref.reshape(K, -1, C)

        for op_type in op_types:
            inp_th, weight_th, output_th = tester.get_operands_torch(op_type)
            weight_th = weight_th.view(K, -1, C)
            inp_tv = torch_tensor_to_tv(inp_th)
            weight_tv = torch_tensor_to_tv(weight_th)
            output_tv = torch_tensor_to_tv(output_th)
            if op_type == ConvOpType.kForward:
                a = inp_tv
                c = output_tv
                b = weight_tv.select(1, tester.kv // 2)
yan.yan's avatar
yan.yan committed
688
689
690
691
                if SPCONV_CPP_GEMM:
                    avail_desps = GEMM_CPP.get_all_available(a, b, c, False, True, False, arch, ShuffleStrideType.ShuffleAC.value)
                else:
                    avail_desps = GEMM.get_all_available(a, b, c, False, True, False, arch, ShuffleStrideType.ShuffleAC)
yan.yan's avatar
yan.yan committed
692
693
694
695

                for desp in avail_desps:
                    if subm:
                        torch.mm(inp_th, weight_th[:, tester.kv // 2].T, out=output_th)
yan.yan's avatar
yan.yan committed
696
                            # output_th += bias_th
yan.yan's avatar
yan.yan committed
697
698
699
                    else:
                        output_tv.zero_()
                    inited = subm
yan.yan's avatar
yan.yan committed
700
                    # determine last valid subm indices, then apply 
yan.yan's avatar
yan.yan committed
701
702
703
704
705
706
707
708
709
710
711
712
                    for i, nhot in enumerate(indice_pair_num_cpu):
                        if subm and i == kv_center:
                            continue
                        if subm and i > kv_center:
                            nhot = indice_pair_num_cpu[kv - i - 1]
                        if nhot <= 0:
                            continue
                        inp_indices = pair_in[i].slice_first_axis(0, nhot)
                        out_indices = pair_out[i].slice_first_axis(0, nhot)
                        b = weight_tv.select(1, i)
                        # inp @ filter.T, NC @ KC
                        beta = 1.0 if inited else 0.0
yan.yan's avatar
yan.yan committed
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
                        if SPCONV_CPP_GEMM:
                            GEMM_CPP.run_with_tuned_result(
                                GemmTuneResult(desp, tester.arch, 1),
                                a,
                                b,
                                c,
                                False,
                                True,
                                False,
                                arch=arch,
                                stream_int=stream,
                                shuffle_type=ShuffleStrideType.ShuffleAC.value,
                                a_inds=inp_indices,
                                b_inds=tv.Tensor(),
                                c_inds=out_indices,
                                hint=AlgoHint.Fowrard.value,
                                alpha=1.0,
yan.yan's avatar
yan.yan committed
730
731
                                beta=beta,
                                force_nvrtc=force_nvrtc)
yan.yan's avatar
yan.yan committed
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
                        else:
                            GEMM.run_with_tuned_result(
                                BestAlgoByProfile(desp, tester.arch, 1),
                                a,
                                b,
                                c,
                                False,
                                True,
                                False,
                                arch=arch,
                                stream=stream,
                                shuffle_type=ShuffleStrideType.ShuffleAC,
                                a_inds=inp_indices,
                                c_inds=out_indices,
                                hint=AlgoHint.Fowrard.value,
                                alpha=1.0,
yan.yan's avatar
yan.yan committed
748
749
                                beta=beta,
                                force_nvrtc=force_nvrtc)
yan.yan's avatar
yan.yan committed
750
                        inited = True
yan.yan's avatar
yan.yan committed
751
752
753
754
755
756
757
                    if bias is not None and tester.check_act:
                        InferenceOps.bias_add_act_inplace(output_tv, bias, tv.gemm.Activation.ReLU, 0, 0)
                    else:
                        if bias is not None:
                            InferenceOps.bias_add_inplace(output_tv, bias, 0)
                        if tester.check_act:
                            InferenceOps.activation_inplace(output_tv, tv.gemm.Activation.ReLU, 0, 0)
yan.yan's avatar
yan.yan committed
758
759
760
761
762
763
764
765
766
                    out_my = output_tv.cpu().numpy()
                    if dtype != np.float16:
                        # error_norm = np.linalg.norm(out_ref.reshape(-1) - out_my.reshape(-1))
                        # assert error_norm < 1
                        # print(desp, K, C, k, error_norm)

                        test_case.assertAllClose(out_ref, out_my, atol=atol, rtol=rtol)
                    else:
                        error_norm = np.linalg.norm(out_ref.reshape(-1) - out_my.reshape(-1))
yan.yan's avatar
yan.yan committed
767
                        assert error_norm < 10 * multipler
yan.yan's avatar
yan.yan committed
768
769
770
771
772

            elif op_type == ConvOpType.kBackwardInput:
                a = output_tv
                b = weight_tv.select(1, tester.kv // 2)
                c = inp_tv
yan.yan's avatar
yan.yan committed
773
                if SPCONV_CPP_GEMM:
yan.yan's avatar
yan.yan committed
774
775
                    avail_desps = GEMM_CPP.get_all_available(a, b, c, False, False, False, arch, ShuffleStrideType.ShuffleAC.value,
                        use_tf32=SPCONV_ALLOW_TF32)
yan.yan's avatar
yan.yan committed
776
                else:
yan.yan's avatar
yan.yan committed
777
778
                    avail_desps = GEMM.get_all_available(a, b, c, False, False, False, arch, ShuffleStrideType.ShuffleAC,
                        use_tf32=SPCONV_ALLOW_TF32)
yan.yan's avatar
yan.yan committed
779

yan.yan's avatar
yan.yan committed
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
                for desp in avail_desps:
                    if subm:
                        torch.mm(output_th, weight_th[:, tester.kv // 2], out=inp_th)
                    else:
                        inp_tv.zero_()
                    inited = subm
                    for i, nhot in enumerate(indice_pair_num_cpu):
                        if subm and i == kv_center:
                            continue
                        if subm and i > kv_center:
                            nhot = indice_pair_num_cpu[kv - i - 1]
                        if nhot <= 0:
                            continue
                        inp_indices = pair_in[i].slice_first_axis(0, nhot)
                        out_indices = pair_out[i].slice_first_axis(0, nhot)
                        b = weight_tv.select(1, i)
                        # inp @ filter.T, NC @ KC
                        beta = 1.0 if inited else 0.0
yan.yan's avatar
yan.yan committed
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
                        if SPCONV_CPP_GEMM:
                            GEMM_CPP.run_with_tuned_result(
                                GemmTuneResult(desp, tester.arch, 1),
                                a,
                                b,
                                c,
                                False,
                                False,
                                False,
                                arch=arch,
                                stream_int=stream,
                                shuffle_type=ShuffleStrideType.ShuffleAC.value,
                                a_inds=out_indices,
                                b_inds=tv.Tensor(),
                                c_inds=inp_indices,
                                hint=AlgoHint.Fowrard.value,
                                alpha=1.0,
yan.yan's avatar
yan.yan committed
815
816
                                beta=beta,
                                force_nvrtc=force_nvrtc)
yan.yan's avatar
yan.yan committed
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
                        else:
                            GEMM.run_with_tuned_result(
                                BestAlgoByProfile(desp, tester.arch, 1),
                                a,
                                b,
                                c,
                                False,
                                False,
                                False,
                                arch=arch,
                                stream=stream,
                                shuffle_type=ShuffleStrideType.ShuffleAC,
                                a_inds=out_indices,
                                c_inds=inp_indices,
                                hint=AlgoHint.Fowrard.value,
                                alpha=1.0,
yan.yan's avatar
yan.yan committed
833
834
                                beta=beta,
                                force_nvrtc=force_nvrtc)
yan.yan's avatar
yan.yan committed
835

yan.yan's avatar
yan.yan committed
836
837
838
839
840
841
842
843
844
845
                        inited = True
                    din_my = inp_tv.cpu().numpy()
                    if dtype != np.float16:
                        # error_norm = np.linalg.norm(din_ref.reshape(-1) - din_my.reshape(-1))
                        # print(desp, K, C, k, error_norm)
                        test_case.assertAllClose(din_ref, din_my, atol=atol, rtol=rtol)
                        # assert error_norm < 1

                    else:
                        error_norm = np.linalg.norm(din_ref.reshape(-1) - din_my.reshape(-1))
yan.yan's avatar
yan.yan committed
846
                        assert error_norm < 10 * multipler
yan.yan's avatar
yan.yan committed
847
848
849
850
            else:
                a = output_tv
                b = inp_tv
                c = weight_tv.select(1, tester.kv // 2)
yan.yan's avatar
yan.yan committed
851
                if SPCONV_CPP_GEMM:
yan.yan's avatar
yan.yan committed
852
853
                    avail_desps = GEMM_CPP.get_all_available(a, b, c, True, False, False, arch, ShuffleStrideType.ShuffleAB.value,
                        use_tf32=SPCONV_ALLOW_TF32)
yan.yan's avatar
yan.yan committed
854
                else:
yan.yan's avatar
yan.yan committed
855
856
                    avail_desps = GEMM.get_all_available(a, b, c, True, False, False, arch, ShuffleStrideType.ShuffleAB,
                        use_tf32=SPCONV_ALLOW_TF32)
yan.yan's avatar
yan.yan committed
857

yan.yan's avatar
yan.yan committed
858
                for desp in avail_desps:
yan.yan's avatar
yan.yan committed
859
860
                    # print(desp, C, K, k, s, p, d)
                    # desp.is_nvrtc = True
yan.yan's avatar
yan.yan committed
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
                    inited = subm
                    weight_tv.zero_()
                    if subm:
                        torch.mm(output_th.T, inp_th, out=weight_th[:, kv_center])

                    for i, nhot in enumerate(indice_pair_num_cpu):
                        if subm and i == kv_center:
                            continue
                        if subm and i > kv_center:
                            nhot = indice_pair_num_cpu[kv - i - 1]
                        if nhot <= 0:
                            continue
                        beta = 1.0 if inited else 0.0
                        inp_indices = pair_in[i].slice_first_axis(0, nhot)
                        out_indices = pair_out[i].slice_first_axis(0, nhot)
                        a_inds = out_indices
                        b_inds = inp_indices
yan.yan's avatar
yan.yan committed
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
                        if SPCONV_CPP_GEMM:
                            GEMM_CPP.run_with_tuned_result(
                                GemmTuneResult(desp, tester.arch, 32),
                                a,
                                b,
                                weight_tv.select(1, i),
                                True,
                                False,
                                False,
                                arch=arch,
                                stream_int=stream,
                                shuffle_type=ShuffleStrideType.ShuffleAB.value,
                                a_inds=a_inds,
                                b_inds=b_inds,
                                c_inds=tv.Tensor(),
                                hint=AlgoHint.BackwardWeight.value,
                                alpha=1.0,
yan.yan's avatar
yan.yan committed
895
896
                                beta=beta,
                                force_nvrtc=force_nvrtc)
yan.yan's avatar
yan.yan committed
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912

                        else:
                            GEMM.run_with_tuned_result(BestAlgoByProfile(desp, tester.arch, 32),
                                                    a,
                                                    b,
                                                    weight_tv.select(1, i),
                                                    True,
                                                    False,
                                                    False,
                                                    arch=arch,
                                                    stream=stream,
                                                    shuffle_type=ShuffleStrideType.ShuffleAB,
                                                    a_inds=a_inds,
                                                    b_inds=b_inds,
                                                    hint=AlgoHint.BackwardWeight.value,
                                                    alpha=1.0,
yan.yan's avatar
yan.yan committed
913
914
                                                    beta=beta,
                                                    force_nvrtc=force_nvrtc)
yan.yan's avatar
yan.yan committed
915
916
917
918

                    dw_my = weight_tv.cpu().numpy()
                    if dtype != np.float16:
                        error_norm = np.linalg.norm(dw_ref.reshape(-1) - dw_my.reshape(-1))
yan.yan's avatar
yan.yan committed
919
                        assert error_norm < 1 * multipler, f"{desp}, {error_norm}"
yan.yan's avatar
yan.yan committed
920
921
                    else:
                        error_norm = np.linalg.norm(dw_ref.reshape(-1) - dw_my.reshape(-1))
yan.yan's avatar
yan.yan committed
922
                        assert error_norm < 10 * multipler, f"{desp}, {error_norm}"
yan.yan's avatar
yan.yan committed
923
924
925


def test_all_algo_unit():
yan.yan's avatar
yan.yan committed
926
    # for i in range(5):
yan.yan's avatar
yan.yan committed
927
928
    _test_impgemm_conv_cuda(True)
    _test_impgemm_conv_cuda(False)
929
930
    # _test_native_conv_cuda(True)
    # _test_native_conv_cuda(False)
yan.yan's avatar
yan.yan committed
931
932
933
934


if __name__ == "__main__":
    test_all_algo_unit()