test_all_algo.py 49.1 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
from cumm.conv.bases import NCHW, NHWC, ConvIterAlgo, ConvOpType
import os
yan.yan's avatar
yan.yan committed
40
41
from cumm.dtypes import get_npdtype_from_tvdtype

yan.yan's avatar
yan.yan committed
42
43
44
45
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
46
from spconv.algo import GEMM, CONV, GEMM_CPP, CONV_CPP, BestAlgoByProfile, BestConvAlgoByProfile, GemmTunerSimple
yan.yan's avatar
yan.yan committed
47
48
49
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
50
from spconv.constants import ALL_WEIGHT_IS_KRSC, SPCONV_CPP_GEMM
yan.yan's avatar
yan.yan committed
51
52
from spconv.core_cc.csrc.sparse.inference import InferenceOps
from spconv.pytorch import functional as Fsp
yan.yan's avatar
yan.yan committed
53
assert ALL_WEIGHT_IS_KRSC is True, "we only support KRSC in spconv >= 2.2"
yan.yan's avatar
yan.yan committed
54
from spconv.pytorch.hash import HashTable
yan.yan's avatar
yan.yan committed
55
56
57
58
59
60
61
62
63
64
65
66
67

# 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:
yan.yan's avatar
yan.yan committed
68
69
70
    def __init__(self, algo: ConvAlgo, subm: bool, shape: List[int], bs: int, dtype: np.dtype, out_dtype: np.dtype, N: int, K: int, C: int, 
        ksize: int, stride: int, padding: int, dilation: int, check_bias: bool = False, check_act: bool = False,
        check_int8_infer: bool = False, dtype_comp: np.dtype = np.dtype(np.float32)) -> None:
yan.yan's avatar
yan.yan committed
71
        ndim = 3
yan.yan's avatar
yan.yan committed
72
        transpose = False
yan.yan's avatar
yan.yan committed
73
74
75
        self.shape = shape 
        self.bs = bs 
        self.dtype = dtype 
yan.yan's avatar
yan.yan committed
76
        self.out_dtype = out_dtype
yan.yan's avatar
yan.yan committed
77
        self.dtype_th = NUMPY_DTYPE_TO_TORCH[dtype]
yan.yan's avatar
yan.yan committed
78
79
        self.out_dtype_th = NUMPY_DTYPE_TO_TORCH[out_dtype]

yan.yan's avatar
yan.yan committed
80
81
        self.K = K 
        self.C = C 
yan.yan's avatar
yan.yan committed
82
83
84
85
        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
86
87
        self.N = N
        self.device = torch.device("cuda:0")
yan.yan's avatar
yan.yan committed
88
        op = expand_nd(ndim, 0)
yan.yan's avatar
yan.yan committed
89
90
        self.kv: int = np.prod(self.ksize)
        self.num_split = 1 if algo == ConvAlgo.MaskImplicitGemm else 2
yan.yan's avatar
yan.yan committed
91
        self.output_scale: float = 3.4
yan.yan's avatar
yan.yan committed
92
93
94
95
96
        self.check_int8_infer = check_int8_infer
        if check_int8_infer:
            assert check_bias and self.dtype == np.int8

        self.dtype_comp = dtype_comp
yan.yan's avatar
yan.yan committed
97
98
99
100
101
102
103
104
105
        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
106
        self.scales = np.random.uniform(0.5, 1.5, size=K).astype(dtype_comp)
yan.yan's avatar
yan.yan committed
107

yan.yan's avatar
yan.yan committed
108
        sparse_dict = generate_sparse_data(shape, [N] * bs, C)
yan.yan's avatar
yan.yan committed
109
110
111
112
113
114
115
116
117

        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
118
119
        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
120
121
122
123
        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
124
        self.use_direct_table = True
yan.yan's avatar
yan.yan committed
125
        self.mask_int_count = div_up(self.kv, 32)
yan.yan's avatar
yan.yan committed
126
        self.out_shape = out_shape
yan.yan's avatar
yan.yan committed
127
128
129
130
131
132
133
134
135
136
137
138
139
        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
140
                                                    self.dilation, op, subm=subm, direct_table=self.use_direct_table)
yan.yan's avatar
yan.yan committed
141
142
143
144
145
146
147
148
149
150
            
            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]
yan.yan's avatar
yan.yan committed
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
        
        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
168
169
        self.voxels_np = voxels_np
        self.indices_np = indices_np
yan.yan's avatar
yan.yan committed
170
171
        self.check_bias = check_bias
        self.check_act = check_act
yan.yan's avatar
yan.yan committed
172
173

        self.subm = subm
yan.yan's avatar
yan.yan committed
174
        self.output_add_scale = 1.0
yan.yan's avatar
yan.yan committed
175
        if dtype == np.int8:
yan.yan's avatar
yan.yan committed
176
            self.inp = np.random.randint(-1, 1, size=[voxels_np.shape[0],
yan.yan's avatar
yan.yan committed
177
                                                    C]).astype(np.int8)
yan.yan's avatar
yan.yan committed
178
            self.weight = np.random.randint(-1, 1, size=[K, *self.ksize,
yan.yan's avatar
yan.yan committed
179
                                                    C]).astype(np.int8)
yan.yan's avatar
yan.yan committed
180
181
            
            self.output = np.random.randint(-1, 1, size=[
yan.yan's avatar
yan.yan committed
182
                self.out_inds.shape[0], K
yan.yan's avatar
yan.yan committed
183
184
185
186
187
188
189
190
191
192
193
194
195
            ]).astype(out_dtype)
            self.output_add = np.random.randint(-1, 1, size=[
                self.out_inds.shape[0], K
            ]).astype(out_dtype)
            self.output_add_scale = 14.2
            if check_int8_infer:
                self.bias = np.random.uniform(-5, 5, size=[
                    K
                ]).astype(dtype_comp)
            else:
                self.bias = np.random.randint(-4, 4, size=[
                    K
                ]).astype(dtype)
yan.yan's avatar
yan.yan committed
196
197
198
199
200
201
202
        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
yan.yan's avatar
yan.yan committed
203
204
205
206
            ]).astype(out_dtype)
            self.output_add = np.random.uniform(-1, 1, size=[
                self.out_inds.shape[0], K
            ]).astype(out_dtype)
yan.yan's avatar
yan.yan committed
207
208
209
            self.bias = np.random.uniform(-1, 1, size=[
                K
            ]).astype(dtype)
yan.yan's avatar
yan.yan committed
210
211
212
        # self.bias[:] = 0
        # self.scales[:] = 1
        
yan.yan's avatar
yan.yan committed
213
214
215
216
        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()
        self.dw_ref = np.ascontiguousarray(self.dw_ref.transpose(1, 0, 2).reshape(K, *self.ksize, C))
yan.yan's avatar
yan.yan committed
217
        self.arch = tv.get_compute_capability()
yan.yan's avatar
yan.yan committed
218

yan.yan's avatar
yan.yan committed
219
220
221
    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
222
    def _get_ref_output(self):
yan.yan's avatar
yan.yan committed
223
224
225
226
227
        out_dtype = np.float32 
        if self.dtype == np.int8:
            out_dtype = np.int32 
        output_ref = np.zeros_like(self.output, dtype=out_dtype)

yan.yan's avatar
yan.yan committed
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
        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]
yan.yan's avatar
yan.yan committed
243
244
245
246
247
248
249
250
            if self.dtype == np.int8:
                cc = a.astype(
                    np.int32) @ self.weight_ref[filter_offset].T.astype(
                        np.int32)
            else:
                cc = a.astype(
                    np.float32) @ self.weight_ref[filter_offset].T.astype(
                        np.float32)
yan.yan's avatar
yan.yan committed
251
            output_ref[o_inds] += cc
yan.yan's avatar
yan.yan committed
252
253
            # we use random output as dout here
            a = self.output[self.out_order][o_inds]
yan.yan's avatar
yan.yan committed
254
255
256
257
258
            # 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
259
260
            # 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
261
262
263
264
265
            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
266
267
268
269
270
271
        if not self.check_int8_infer:
            if self.check_bias:
                output_ref += self.bias
                # relu
            if self.check_act:
                output_ref = np.maximum(output_ref, 0)
yan.yan's avatar
yan.yan committed
272
        if self.dtype == np.int8:
yan.yan's avatar
yan.yan committed
273
274
275
            if self.check_int8_infer:
                rescaled = output_ref.astype(self.dtype_comp) * self.scales.astype(self.dtype_comp)
                rescaled += self.bias.astype(self.dtype_comp)
yan.yan's avatar
yan.yan committed
276
277
278
279
                if self.subm:
                    rescaled += self.output_add.astype(self.dtype_comp) * self.output_add_scale
                else:
                    rescaled += self.output_add[self.out_order].astype(self.dtype_comp) * self.output_add_scale
yan.yan's avatar
yan.yan committed
280
281
282
283
284
285
286
287
                if self.check_act:
                    rescaled = np.maximum(rescaled, 0)
                if self.out_dtype == np.int8:
                    output_ref = np.clip(np.round(rescaled), -128, 127).astype(np.int8)
                else:
                    output_ref = rescaled.astype(self.out_dtype)
            else:
                output_ref = np.clip(output_ref, -127, 127)
yan.yan's avatar
yan.yan committed
288
289
290
291
292
293
294
295
296
297
298
299
300
        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:
yan.yan's avatar
yan.yan committed
301
            output_tv = zeros_func(list(self.output.shape), self.out_dtype, 0)
yan.yan's avatar
yan.yan committed
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
        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
324
    np.random.seed(50005)
yan.yan's avatar
yan.yan committed
325
    dtype_to_tol = {
yan.yan's avatar
yan.yan committed
326
        np.float32: (1e-2, 1e-2),
yan.yan's avatar
yan.yan committed
327
328
329
330
331
332
        np.float16: (1e-2, 1e-2),
        np.int8: (1e-4, 1e-4),
    }
    device = torch.device("cuda:0")
    shapes = [[19, 18, 17]]
    batchsizes = [1]
yan.yan's avatar
yan.yan committed
333
    # dtypes = [(np.float32, np.float32), (np.float16, np.float16)]
yan.yan's avatar
yan.yan committed
334
    # dtypes = [np.float16]
yan.yan's avatar
yan.yan committed
335
    # dtypes = [(np.int8, np.int8), (np.int8, np.float32), (np.int8, np.float16)]
yan.yan's avatar
yan.yan committed
336
    dtypes = [(np.int8, np.int8)]
yan.yan's avatar
yan.yan committed
337
    # dtypes = [(np.float16, np.float16)]
yan.yan's avatar
yan.yan committed
338

yan.yan's avatar
yan.yan committed
339
    test_case = TestCase()
yan.yan's avatar
yan.yan committed
340
341
342
343
    # in_channels = [32]
    # out_channels = [32, 48, 64]
    in_channels = [32, 47]
    out_channels = [32, 48, 62]
yan.yan's avatar
yan.yan committed
344
345
346
    in_channels = [16]
    out_channels = [16]

yan.yan's avatar
yan.yan committed
347
348
    # in_channels = [16]
    # out_channels = [16]
yan.yan's avatar
yan.yan committed
349

yan.yan's avatar
yan.yan committed
350
    multiple_base = 16
yan.yan's avatar
yan.yan committed
351
    if subm:
yan.yan's avatar
yan.yan committed
352
        # ksizes = [3, (3, 3, 5), (3, 5, 5), 5]
yan.yan's avatar
yan.yan committed
353
        ksizes = [3, 5]
yan.yan's avatar
yan.yan committed
354
355
356
357
        strides = [1]
        paddings = [0]
        dilations = [1]
    else:
358
        ksizes = [2, 3, (3, 3, 4), 4, (4, 5, 5), 5]
yan.yan's avatar
yan.yan committed
359
        ksizes = [2, 3, 5]
yan.yan's avatar
yan.yan committed
360

yan.yan's avatar
yan.yan committed
361
362
363
        strides = [1, 2, 3]
        paddings = [0, 1]
        dilations = [1, 2]
yan.yan's avatar
yan.yan committed
364

yan.yan's avatar
yan.yan committed
365
    algos = [
yan.yan's avatar
yan.yan committed
366
        # ConvAlgo.MaskSplitImplicitGemm,
yan.yan's avatar
yan.yan committed
367
368
369
        ConvAlgo.MaskImplicitGemm,
    ]
    arch = torch.cuda.get_device_capability()
yan.yan's avatar
yan.yan committed
370
    force_nvrtc = False
yan.yan's avatar
yan.yan committed
371
    for shape, bs, C, K, k, s, p, d, algo, dtype_outdtype in tqdm.tqdm(params_grid(
yan.yan's avatar
yan.yan committed
372
373
            shapes, batchsizes, in_channels, out_channels, ksizes,
            strides, paddings, dilations, algos, dtypes)):
yan.yan's avatar
yan.yan committed
374
375
376
377
378
379
380
        dtype, out_dtype = dtype_outdtype
        if (C % 16 != 0 or K % 16 != 0) and dtype == np.int8:
            continue
        dcomp = np.float32
        check_int8_infer = True
        if dtype != np.int8:
            check_int8_infer = False
yan.yan's avatar
yan.yan committed
381
382
        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
383
384
385
386
        # 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
387
        # print(num_batch)
yan.yan's avatar
yan.yan committed
388
389
390
391
        tester = SparseConvTester(algo, subm, shape, bs, dtype, out_dtype, num_batch, K, C, k, s, p, d, 
            check_bias=True, check_act=True, check_int8_infer=check_int8_infer, dtype_comp=np.float32)
        enable_dy_mask = tester.kv > 32
        output_add_cuda = tv.from_numpy(tester.output_add).cuda()
yan.yan's avatar
yan.yan committed
392
        bias = None
yan.yan's avatar
yan.yan committed
393
        scales = None
yan.yan's avatar
yan.yan committed
394
395
        act = tv.gemm.Activation.None_
        if tester.check_bias:
yan.yan's avatar
yan.yan committed
396
397
398
399
400
401
402
            if check_int8_infer:
                bias = tv.from_numpy(tester.bias.astype(dcomp)).cuda()
            else:
                bias = tv.from_numpy(tester.bias).cuda()
            if check_int8_infer:
                scales = tv.from_numpy(tester.scales.astype(dcomp)).cuda()

yan.yan's avatar
yan.yan committed
403
404
405
406
407
408
        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:
yan.yan's avatar
yan.yan committed
409
410
            if tester.dtype == np.int8 and op_type != ConvOpType.kForward:
                continue
yan.yan's avatar
yan.yan committed
411
            inp_tv, weight_tv, output_tv = tester.get_operands(op_type)
yan.yan's avatar
yan.yan committed
412
413
414
            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
415
                    NHWC.layout_type.value, NHWC.interleave, NHWC.interleave, NHWC.interleave, arch, op_type.value, -1, True, False,
yan.yan's avatar
yan.yan committed
416
417
                        use_tf32=SPCONV_ALLOW_TF32, bias=bias if bias is not None else tv.Tensor(), 
                        scale=scales if scales is not None else tv.Tensor())
yan.yan's avatar
yan.yan committed
418
            else:
yan.yan's avatar
yan.yan committed
419
                avail_desps = CONV.get_all_available(inp_tv, weight_tv, output_tv, NHWC, NHWC, NHWC, arch, op_type, -1,
yan.yan's avatar
yan.yan committed
420
421
                        use_tf32=SPCONV_ALLOW_TF32, bias=bias if bias is not None else tv.Tensor(), 
                        scale=scales if scales is not None else tv.Tensor())
yan.yan's avatar
yan.yan committed
422
423
424
425
426
            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
427
            for desp in avail_desps:
yan.yan's avatar
yan.yan committed
428
429
430
431
432
                dcomp = get_npdtype_from_tvdtype(desp.dcomp)
                if enable_dy_mask and not desp.dynamic_mask:
                    continue
                if tester.check_int8_infer and not desp.is_int8_inference:
                    continue
yan.yan's avatar
yan.yan committed
433
434
435
436
437
438
439
                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]
yan.yan's avatar
yan.yan committed
440
441
442
                alpha = 1.0
                if tester.check_int8_infer:
                    alpha = tester.output_scale
yan.yan's avatar
yan.yan committed
443
444
445
                # if mask_width != 32:
                #     continue
                if mask_width not in mask_width_to_mask_out_fwd:
yan.yan's avatar
yan.yan committed
446
                    mask_width_to_mask_out_fwd[mask_width] = torch.zeros([2, div_up(tester.out_inds.shape[0], mask_width), tester.mask_int_count],
yan.yan's avatar
yan.yan committed
447
448
449
                                      dtype=torch.int32,
                                      device=tester.device)
                mask_output_fwd = mask_width_to_mask_out_fwd[mask_width]
yan.yan's avatar
yan.yan committed
450
451
452
453
                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
454
455
456
457
458
                output_add_cur_tv = tv.Tensor()
                output_add_cur = None 
                if is_fwd and tester.check_int8_infer:
                    output_add_cur = output_add_cuda
                    output_add_cur_tv = output_add_cur
yan.yan's avatar
yan.yan committed
459
                if subm:
yan.yan's avatar
yan.yan committed
460
                    if desp.op_type.value == ConvOpType.kForward.value:
yan.yan's avatar
yan.yan committed
461
                        indice_pairs = tester.pair_fwd
yan.yan's avatar
yan.yan committed
462
                    elif desp.op_type.value == ConvOpType.kBackwardInput.value:
yan.yan's avatar
yan.yan committed
463
464
465
466
467
468
                        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
469
                        beta = 1 if j > 0 else 0
yan.yan's avatar
yan.yan committed
470
471
                        if bias_cur is not None and not tester.check_int8_infer:
                            # this beta is used for C-beta (use C as bias, not standalone bias)
yan.yan's avatar
yan.yan committed
472
473
474
                            beta = 1
                        if j > 0:
                            bias_cur = None
yan.yan's avatar
yan.yan committed
475
476
                        if output_add_cur is not None and tester.check_int8_infer:
                            beta = tester.output_add_scale
yan.yan's avatar
yan.yan committed
477
478
                        mask_filter = tester.masks[j].item()
                        reverse_mask = False
yan.yan's avatar
yan.yan committed
479
                        if desp.op_type.value == ConvOpType.kBackwardWeight.value:
yan.yan's avatar
yan.yan committed
480
481
482
                            mask_op = mask_output[j]
                        else:
                            mask_op = tester.pair_mask_fwd_splits[j]
yan.yan's avatar
yan.yan committed
483
                        if desp.op_type.value == ConvOpType.kBackwardInput.value:
yan.yan's avatar
yan.yan committed
484
485
                            reverse_mask = True
                        mask_output_run = torch_tensor_to_tv(mask_output[j], dtype=tv.uint32)
yan.yan's avatar
yan.yan committed
486
                        if desp.op_type.value == ConvOpType.kBackwardWeight.value:
yan.yan's avatar
yan.yan committed
487
                            mask_output_run = tv.Tensor()
yan.yan's avatar
yan.yan committed
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
                        # 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,
yan.yan's avatar
yan.yan committed
506
                                alpha=alpha,
yan.yan's avatar
yan.yan committed
507
508
509
                                beta=beta,
                                verbose=False,
                                force_nvrtc=force_nvrtc,
yan.yan's avatar
yan.yan committed
510
                                bias=bias_cur if is_fwd and bias_cur is not None else tv.Tensor(),
yan.yan's avatar
yan.yan committed
511
                                scale=scales if is_fwd and scales is not None else tv.Tensor(),
yan.yan's avatar
yan.yan committed
512
                                act_type=act,
yan.yan's avatar
yan.yan committed
513
                                output_add=output_add_cur_tv)
yan.yan's avatar
yan.yan committed
514
515
516
517
518
519
520
521
522
523
524
525
526
527
                        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,
yan.yan's avatar
yan.yan committed
528
                                alpha=alpha,
yan.yan's avatar
yan.yan committed
529
530
531
                                beta=beta,
                                verbose=False,
                                force_nvrtc=force_nvrtc,
yan.yan's avatar
yan.yan committed
532
                                bias=bias_cur if is_fwd else None,
yan.yan's avatar
yan.yan committed
533
                                scale=scales if is_fwd else None,
yan.yan's avatar
yan.yan committed
534
                                act_type=act,
yan.yan's avatar
yan.yan committed
535
                                output_add=output_add_cur,
yan.yan's avatar
yan.yan committed
536
537
                            )

yan.yan's avatar
yan.yan committed
538
539
540
541
542
543
544
                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
545
                    if desp.op_type.value == ConvOpType.kForward.value:
yan.yan's avatar
yan.yan committed
546
547
548
549
                        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
550
                    elif desp.op_type.value == ConvOpType.kBackwardInput.value:
yan.yan's avatar
yan.yan committed
551
552
553
554
555
556
557
558
559
560
561
                        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
562
563
                        # beta = 1 if j == 1 else 0
                        beta = 1 if j > 0 else 0
yan.yan's avatar
yan.yan committed
564
565
                        if bias_cur is not None and not tester.check_int8_infer:
                            # this beta is used for C-beta (use C as bias, not standalone bias)
yan.yan's avatar
yan.yan committed
566
567
568
                            beta = 1
                        if j > 0:
                            bias_cur = None
yan.yan's avatar
yan.yan committed
569
570
                        if output_add_cur is not None and tester.check_int8_infer:
                            beta = tester.output_add_scale
yan.yan's avatar
yan.yan committed
571
572
                        mask_filter = tester.masks[j].item()
                        reverse_mask = False
yan.yan's avatar
yan.yan committed
573
                        if desp.op_type.value == ConvOpType.kBackwardWeight.value:
yan.yan's avatar
yan.yan committed
574
575
576
                            mask_op = mask_output[j]
                        else:
                            mask_op = mask_ops[j]
yan.yan's avatar
yan.yan committed
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
                        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
593
                                force_nvrtc=force_nvrtc,
yan.yan's avatar
yan.yan committed
594
595
                                bias=bias_cur if is_fwd and bias_cur is not None else tv.Tensor(),
                                scale=scales if is_fwd and scales is not None else tv.Tensor(),
yan.yan's avatar
yan.yan committed
596
                                act_type=act,
yan.yan's avatar
yan.yan committed
597
                                output_add=output_add_cur_tv,
yan.yan's avatar
yan.yan committed
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
                            )
                        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
615
                                force_nvrtc=force_nvrtc,
yan.yan's avatar
yan.yan committed
616
617
                                bias=bias_cur if is_fwd else None,
                                scale=scales if is_fwd else None,
yan.yan's avatar
yan.yan committed
618
                                act_type=act,
yan.yan's avatar
yan.yan committed
619
                                output_add=output_add_cur,
yan.yan's avatar
yan.yan committed
620
                            )
yan.yan's avatar
yan.yan committed
621
622
623
624
625
626

                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
627
                    out_my = out_my[tester.out_order]
yan.yan's avatar
yan.yan committed
628
                    if dtype != np.float16:
yan.yan's avatar
yan.yan committed
629
630
                        if dtype == np.int8:
                            print("max int8 diff", np.abs(out_ref - out_my).max())
yan.yan's avatar
yan.yan committed
631
632
633
                        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
634
635
                        if (error_norm > 5):
                            print(f"{desp}, Error={error_norm}")
yan.yan's avatar
yan.yan committed
636
                        assert error_norm < 10 * multipler
yan.yan's avatar
yan.yan committed
637
638
639
640
641
642
                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
643
                        assert error_norm < 10 * multipler, f"{desp}, {error_norm}, {k}, {s}, {p}, {d}"
yan.yan's avatar
yan.yan committed
644
645
646
647
648
649
650
651
652
653
        if not tester.check_int8_infer:
            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():
                    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, 
                            ConvOpType.kBackwardWeight.value, mask_width, True, False,
                            use_tf32=SPCONV_ALLOW_TF32)
yan.yan's avatar
yan.yan committed
654
                    else:
yan.yan's avatar
yan.yan committed
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
                        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)
                    for desp in avail_desps:
                        if enable_dy_mask and not desp.dynamic_mask:
                            continue
                        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(
                                    BestConvAlgoByProfile(desp, tester.arch, spk),
                                    desp.op_type.value,
                                    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,
                                )
                        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]
yan.yan's avatar
yan.yan committed
697

yan.yan's avatar
yan.yan committed
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
                                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,
                                )
                        dw_ref = tester.dw_ref
                        dw_my = weight_tv.cpu().numpy()
                        if dtype != np.float16:
                            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)
                            if (error_norm > 5):
                                print(f"{desp}, Error={error_norm}, {spk}")
                            assert error_norm < 10 * multipler
yan.yan's avatar
yan.yan committed
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748

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
749
750
    multiple_base = 128

yan.yan's avatar
yan.yan committed
751
752
    arch = torch.cuda.get_device_capability()
    stream = get_current_stream()
yan.yan's avatar
yan.yan committed
753
    force_nvrtc = False
yan.yan's avatar
yan.yan committed
754
755
756
    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
757
758
759
760
        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
761
        atol, rtol = dtype_to_tol[dtype]
yan.yan's avatar
yan.yan committed
762
763
        multipler = max(C, K) / multiple_base
        multipler = max(multipler, 1.0)
yan.yan's avatar
yan.yan committed
764
765
766
767
768
769
770

        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
771
772
        # op_types = [ConvOpType.kForward]

yan.yan's avatar
yan.yan committed
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
        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
790
791
792
793
                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
794
795
796
797

                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
798
                            # output_th += bias_th
yan.yan's avatar
yan.yan committed
799
800
801
                    else:
                        output_tv.zero_()
                    inited = subm
yan.yan's avatar
yan.yan committed
802
                    # determine last valid subm indices, then apply 
yan.yan's avatar
yan.yan committed
803
804
805
806
807
808
809
810
811
812
813
814
                    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
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
                        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
832
833
                                beta=beta,
                                force_nvrtc=force_nvrtc)
yan.yan's avatar
yan.yan committed
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
                        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
850
851
                                beta=beta,
                                force_nvrtc=force_nvrtc)
yan.yan's avatar
yan.yan committed
852
                        inited = True
yan.yan's avatar
yan.yan committed
853
854
855
856
857
858
859
                    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
860
861
862
863
864
865
866
867
868
                    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
869
                        assert error_norm < 10 * multipler
yan.yan's avatar
yan.yan committed
870
871
872
873
874

            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
875
                if SPCONV_CPP_GEMM:
yan.yan's avatar
yan.yan committed
876
877
                    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
878
                else:
yan.yan's avatar
yan.yan committed
879
880
                    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
881

yan.yan's avatar
yan.yan committed
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
                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
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
                        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
917
918
                                beta=beta,
                                force_nvrtc=force_nvrtc)
yan.yan's avatar
yan.yan committed
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
                        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
935
936
                                beta=beta,
                                force_nvrtc=force_nvrtc)
yan.yan's avatar
yan.yan committed
937

yan.yan's avatar
yan.yan committed
938
939
940
941
942
943
944
945
946
947
                        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
948
                        assert error_norm < 10 * multipler
yan.yan's avatar
yan.yan committed
949
950
951
952
            else:
                a = output_tv
                b = inp_tv
                c = weight_tv.select(1, tester.kv // 2)
yan.yan's avatar
yan.yan committed
953
                if SPCONV_CPP_GEMM:
yan.yan's avatar
yan.yan committed
954
955
                    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
956
                else:
yan.yan's avatar
yan.yan committed
957
958
                    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
959

yan.yan's avatar
yan.yan committed
960
                for desp in avail_desps:
yan.yan's avatar
yan.yan committed
961
962
                    # print(desp, C, K, k, s, p, d)
                    # desp.is_nvrtc = True
yan.yan's avatar
yan.yan committed
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
                    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
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
                        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
997
998
                                beta=beta,
                                force_nvrtc=force_nvrtc)
yan.yan's avatar
yan.yan committed
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014

                        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
1015
1016
                                                    beta=beta,
                                                    force_nvrtc=force_nvrtc)
yan.yan's avatar
yan.yan committed
1017
1018
1019
1020

                    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
1021
                        assert error_norm < 1 * multipler, f"{desp}, {error_norm}"
yan.yan's avatar
yan.yan committed
1022
1023
                    else:
                        error_norm = np.linalg.norm(dw_ref.reshape(-1) - dw_my.reshape(-1))
yan.yan's avatar
yan.yan committed
1024
                        assert error_norm < 10 * multipler, f"{desp}, {error_norm}"
yan.yan's avatar
yan.yan committed
1025
1026
1027


def test_all_algo_unit():
yan.yan's avatar
yan.yan committed
1028
    # for i in range(5):
yan.yan's avatar
yan.yan committed
1029
    _test_impgemm_conv_cuda(True)
yan.yan's avatar
yan.yan committed
1030
    _test_impgemm_conv_cuda(False)
1031
1032
    # _test_native_conv_cuda(True)
    # _test_native_conv_cuda(False)
yan.yan's avatar
yan.yan committed
1033
1034
1035
1036


if __name__ == "__main__":
    test_all_algo_unit()