test_gemm.cu 40.5 KB
Newer Older
Li Zhang's avatar
Li Zhang committed
1
2
#include <assert.h>
#include <cublas_v2.h>
Chen Xin's avatar
Chen Xin committed
3
#include <math.h>
Li Zhang's avatar
Li Zhang committed
4
5
6
7
8
#include <numeric>
#include <stdexcept>
#include <tuple>
#include <vector>

lvhan028's avatar
lvhan028 committed
9
10
11
12
13
14
15
#include "src/turbomind/layers/DenseWeight.h"
#include "src/turbomind/utils/allocator.h"
#include "src/turbomind/utils/cublasMMWrapper.h"
#include "src/turbomind/utils/cuda_utils.h"
#include "src/turbomind/utils/gemm.h"
#include "src/turbomind/utils/logger.h"
#include "src/turbomind/utils/memory_utils.h"
Li Zhang's avatar
Li Zhang committed
16

lvhan028's avatar
lvhan028 committed
17
using namespace turbomind;
Li Zhang's avatar
Li Zhang committed
18
19
20

// Can be replaced by the function provided by a test framework

Chen Xin's avatar
Chen Xin committed
21
class TestFailureError: public std::exception {
Li Zhang's avatar
Li Zhang committed
22
23
private:
    std::string msg_;
Chen Xin's avatar
Chen Xin committed
24

Li Zhang's avatar
Li Zhang committed
25
26
public:
    explicit TestFailureError() = default;
Chen Xin's avatar
Chen Xin committed
27
28
    explicit TestFailureError(std::string name, std::string msg = "")
    {
Li Zhang's avatar
Li Zhang committed
29
30
        msg_ = fmtstr("TEST FAIL [%s] %s", name.c_str(), msg.c_str());
    }
Chen Xin's avatar
Chen Xin committed
31
32
    const char* what() const throw()
    {
Li Zhang's avatar
Li Zhang committed
33
34
35
36
        return msg_.c_str();
    }
};

Chen Xin's avatar
Chen Xin committed
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
#define EXPECT_TRUE(cond)                                                                                              \
    do {                                                                                                               \
        if (!(cond)) {                                                                                                 \
            TM_LOG_ERROR("TEST FAIL [%s] at %s:%d", __func__, __FILE__, __LINE__);                                     \
            throw TestFailureError(__func__);                                                                          \
        }                                                                                                              \
    } while (false)

#define EXPECT_ALMOST_EQUAL(name, dtype, ctype, out, ref)                                                              \
    do {                                                                                                               \
        bool is_ok = checkResult<dtype, ctype>(name, out, ref);                                                        \
        if (!is_ok) {                                                                                                  \
            TM_LOG_ERROR("TEST FAIL [%s] at %s:%d", __func__, __FILE__, __LINE__);                                     \
            throw TestFailureError(__func__);                                                                          \
        }                                                                                                              \
    } while (false)
Li Zhang's avatar
Li Zhang committed
53
54
55
56
57
58
59
60
61
62
63

////////////////////////////////////////////////////////////////////////////////////

// TensorWrapper is to handle a tensor object as well as its memory buffer,
// because tensor.data is const we cannot set values.
class TensorWrapper {
private:
    IAllocator* allocator;

public:
    std::vector<size_t> shape;
Chen Xin's avatar
Chen Xin committed
64
65
66
    DataType            type;
    Tensor*             tensor;
    void*               data;
Li Zhang's avatar
Li Zhang committed
67
68
69
70

    TensorWrapper(IAllocator* allocator, DataType dtype, std::vector<size_t> shape, bool zero_init = false)
    {
        this->allocator = allocator;
Chen Xin's avatar
Chen Xin committed
71
72
        this->type      = dtype;
        this->shape     = shape;
Li Zhang's avatar
Li Zhang committed
73
74

        size_t tensor_memsize = this->memsize();
Chen Xin's avatar
Chen Xin committed
75
        this->data            = this->allocator->malloc(tensor_memsize, false);
Li Zhang's avatar
Li Zhang committed
76
77
        if (zero_init) {
            check_cuda_error(cudaMemset(data, 0x0, tensor_memsize));
Chen Xin's avatar
Chen Xin committed
78
79
        }
        else {
Li Zhang's avatar
Li Zhang committed
80
81
82
83
84
            setRandomValues();
        }
        this->tensor = new Tensor(MEMORY_GPU, dtype, shape, data);
    }

Chen Xin's avatar
Chen Xin committed
85
86
    TensorWrapper(TensorWrapper const& other):
        allocator(other.allocator), shape(other.shape), type(other.type), data(other.data), tensor(other.tensor)
Li Zhang's avatar
Li Zhang committed
87
    {
lvhan028's avatar
lvhan028 committed
88
        TM_LOG_DEBUG("TensorWrapper copy: this=%p other=%p", data, other.data);
Li Zhang's avatar
Li Zhang committed
89
90
91
92
93
94
95
96
97
    }
    ~TensorWrapper()
    {
        delete tensor;
        allocator->free((void**)(&data));
    }

    void setInvalidValues()
    {
Chen Xin's avatar
Chen Xin committed
98
        size_t type_size   = tensor->type == TYPE_FP32 ? sizeof(float) : sizeof(half);
Li Zhang's avatar
Li Zhang committed
99
100
101
102
103
        size_t tensor_size = type_size * tensor->size();
        // Fill by a random number to guarantee invalid values
        check_cuda_error(cudaMemset(data, 0xdc, tensor_size));
    }

Chen Xin's avatar
Chen Xin committed
104
105
    void setRandomValues()
    {
Li Zhang's avatar
Li Zhang committed
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
        // random initialization
        size_t num_elements = this->size();
        switch (this->type) {
            case TYPE_FP32:
                cudaRandomUniform((float*)data, num_elements);
                break;
            case TYPE_FP16:
                cudaRandomUniform((half*)data, num_elements);
                break;
            default:
                // Will be added more if needed.
                throw std::runtime_error("Not supported data type");
        }
    }

Chen Xin's avatar
Chen Xin committed
121
122
    size_t size()
    {
Li Zhang's avatar
Li Zhang committed
123
124
125
126
127
128
129
        size_t n_elements = 1;
        for (size_t s : this->shape) {
            n_elements *= s;
        }
        return n_elements;
    }

Chen Xin's avatar
Chen Xin committed
130
131
    size_t memsize()
    {
Li Zhang's avatar
Li Zhang committed
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
        size_t type_size = 0;
        switch (this->type) {
            case TYPE_FP32:
                type_size = sizeof(float);
                break;
            case TYPE_FP16:
                type_size = sizeof(half);
                break;
            default:
                throw std::runtime_error("Not supported data type.");
        }
        return type_size * this->size();
    }
};

template<DataType computeType>
Chen Xin's avatar
Chen Xin committed
148
149
void computeReference(GemmOp         transa,
                      GemmOp         transb,
Li Zhang's avatar
Li Zhang committed
150
151
152
                      TensorWrapper& C,
                      TensorWrapper& A,
                      TensorWrapper& B,
Chen Xin's avatar
Chen Xin committed
153
154
                      float          alpha = 1.0f,
                      float          beta  = 0.0f)
Li Zhang's avatar
Li Zhang committed
155
156
157
158
159
160
161
162
163
{
    size_t m = C.shape[0];
    size_t n = C.shape[1];
    size_t k = A.shape[1];

    size_t lda = (transa == GEMM_OP_N) ? k : m;
    size_t ldb = (transb == GEMM_OP_N) ? n : k;
    size_t ldc = n;

Chen Xin's avatar
Chen Xin committed
164
165
166
    cudaDataType_t atype        = (A.type == TYPE_FP16) ? CUDA_R_16F : CUDA_R_32F;
    cudaDataType_t btype        = (B.type == TYPE_FP16) ? CUDA_R_16F : CUDA_R_32F;
    cudaDataType_t ctype        = (C.type == TYPE_FP16) ? CUDA_R_16F : CUDA_R_32F;
Li Zhang's avatar
Li Zhang committed
167
168
169
170
171
    cudaDataType_t compute_type = (computeType == TYPE_FP16) ? CUDA_R_16F : CUDA_R_32F;

    cublasHandle_t cublas_handle;
    check_cuda_error(cublasCreate(&cublas_handle));

Chen Xin's avatar
Chen Xin committed
172
173
174
175
    half        h_alpha = (half)alpha;
    half        h_beta  = (half)beta;
    const void* _alpha  = (computeType == TYPE_FP16) ? (const void*)&h_alpha : (const void*)&alpha;
    const void* _beta   = (computeType == TYPE_FP16) ? (const void*)&h_beta : (const void*)&beta;
Li Zhang's avatar
Li Zhang committed
176
177
178
179

    check_cuda_error(cublasGemmEx(cublas_handle,
                                  getCublasOperation(transb),
                                  getCublasOperation(transa),
Chen Xin's avatar
Chen Xin committed
180
181
182
                                  n,
                                  m,
                                  k,
Li Zhang's avatar
Li Zhang committed
183
                                  _alpha,
Chen Xin's avatar
Chen Xin committed
184
185
186
187
188
189
                                  (const void*)B.data,
                                  btype,
                                  ldb,
                                  (const void*)A.data,
                                  atype,
                                  lda,
Li Zhang's avatar
Li Zhang committed
190
                                  _beta,
Chen Xin's avatar
Chen Xin committed
191
192
193
                                  (void*)C.data,
                                  ctype,
                                  ldc,
Li Zhang's avatar
Li Zhang committed
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
                                  compute_type,
                                  CUBLAS_GEMM_DEFAULT));
    check_cuda_error(cublasDestroy(cublas_handle));
    cudaDeviceSynchronize();
}

bool almostEqual(float a, float b, float atol = 1e-5, float rtol = 1e-8)
{
    // Params: a = value to compare and b = reference
    // This function follows implementation of numpy.isclose(), which checks
    //   abs(a - b) <= (atol + rtol * abs(b)).
    // Note that the inequality above is asymmetric where b is considered as
    // a reference value. To account into both absolute/relative errors, it
    // uses absolute tolerance and relative tolerance at the same time. The
    // default values of atol and rtol borrowed from numpy.isclose(). For the
    // case of nan value, the result will be true.
    if (isnan(a) && isnan(b)) {
        return true;
    }
    return fabs(a - b) <= (atol + rtol * fabs(b));
}

template<typename T>
Chen Xin's avatar
Chen Xin committed
217
218
bool _checkResult(std::string name, TensorWrapper& out, TensorWrapper& ref, float atol, float rtol)
{
Li Zhang's avatar
Li Zhang committed
219
220
221
222
    assert(out.type == ref.type);

    size_t out_size = out.size();
    size_t ref_size = ref.size();
Chen Xin's avatar
Chen Xin committed
223
224
    T*     h_out    = reinterpret_cast<T*>(malloc(sizeof(T) * out_size));
    T*     h_ref    = reinterpret_cast<T*>(malloc(sizeof(T) * ref_size));
Li Zhang's avatar
Li Zhang committed
225
226
227
228
229
230
231
232
233
234
235
236
237

    cudaMemcpy(h_out, out.data, sizeof(T) * out_size, cudaMemcpyDeviceToHost);
    cudaMemcpy(h_ref, ref.data, sizeof(T) * ref_size, cudaMemcpyDeviceToHost);
    cudaDeviceSynchronize();

    size_t failures = 0;
    for (size_t i = 0; i < out_size; ++i) {
        // The values for the output and the reference.
        float a = (float)h_out[i];
        float b = (float)h_ref[i];

        bool ok = almostEqual(a, b, atol, rtol);
        // Print the error.
Chen Xin's avatar
Chen Xin committed
238
        if (!ok && failures < 4) {
lvhan028's avatar
lvhan028 committed
239
240
241
242
243
            TM_LOG_ERROR(">> invalid result for i=%lu:", i);
            TM_LOG_ERROR(">>    found......: %10.6f", a);
            TM_LOG_ERROR(">>    expected...: %10.6f", b);
            TM_LOG_ERROR(">>    error......: %.6f", fabsf(a - b));
            TM_LOG_ERROR(">>    tol........: %.6f", atol + rtol * fabs(b));
Li Zhang's avatar
Li Zhang committed
244
245
246
247
248
249
250
251
        }

        // Update the number of failures.
        failures += ok ? 0 : 1;
    }

    // Allow not matched up to 1% elements.
    size_t tol_failures = (size_t)(0.01 * out_size);
lvhan028's avatar
lvhan028 committed
252
    TM_LOG_INFO("check....... %30s : %s (failures: %.2f%% atol: %.2e rtol: %.2e)",
Chen Xin's avatar
Chen Xin committed
253
254
255
256
257
                name.c_str(),
                failures <= tol_failures ? "OK" : "FAILED",
                100. * failures / out_size,
                atol,
                rtol);
Li Zhang's avatar
Li Zhang committed
258
259
260
261
    return failures <= tol_failures;
}

template<typename T, DataType computeType>
Chen Xin's avatar
Chen Xin committed
262
263
264
265
266
bool checkResult(std::string name, TensorWrapper& out, TensorWrapper& ref)
{
    float atol  = (computeType == TYPE_FP32) ? 1e-6f : 1e-3f;
    float rtol  = (computeType == TYPE_FP32) ? 1e-4f : 1e-1f;
    bool  is_ok = false;
Li Zhang's avatar
Li Zhang committed
267
268
    if (sizeof(T) == 4) {
        is_ok = _checkResult<float>(name, out, ref, atol, rtol);
Chen Xin's avatar
Chen Xin committed
269
270
    }
    else {
Li Zhang's avatar
Li Zhang committed
271
272
273
274
275
276
        is_ok = _checkResult<half>(name, out, ref, atol, rtol);
    }
    return is_ok;
}

template<typename T, DataType computeType>
Chen Xin's avatar
Chen Xin committed
277
278
bool checkResult(TensorWrapper& out, TensorWrapper& ref)
{
Li Zhang's avatar
Li Zhang committed
279
280
281
282
    return checkResult<T, computeType>("", out, ref);
}

template<typename T>
Chen Xin's avatar
Chen Xin committed
283
284
std::string toString()
{
Li Zhang's avatar
Li Zhang committed
285
286
287
288
289
290
    std::string str = "dtype=";
    str += std::is_same<T, float>::value ? "FP32" : "FP16";
    return str;
}

template<typename T, DataType ctype>
Chen Xin's avatar
Chen Xin committed
291
292
std::string toString()
{
Li Zhang's avatar
Li Zhang committed
293
294
295
296
297
298
299
    std::string str = "dtype=";
    str += std::is_same<T, float>::value ? "FP32" : "FP16";
    str += ", compute_type=";
    str += (ctype == TYPE_FP32) ? "FP32" : "FP16";
    return str;
}

Chen Xin's avatar
Chen Xin committed
300
301
std::string toString(GemmOp op)
{
Li Zhang's avatar
Li Zhang committed
302
303
304
305
306
307
308
309
    return op == GEMM_OP_N ? "N" : "T";
}

struct GemmOpPair {
    GemmOp transa;
    GemmOp transb;
};

Chen Xin's avatar
Chen Xin committed
310
311
static const std::vector<GemmOpPair> op_pairs{
    {GEMM_OP_N, GEMM_OP_N}, {GEMM_OP_N, GEMM_OP_T}, {GEMM_OP_T, GEMM_OP_N}, {GEMM_OP_T, GEMM_OP_T}};
Li Zhang's avatar
Li Zhang committed
312

Chen Xin's avatar
Chen Xin committed
313
static inline std::string getTestName(const char* func_name, GemmOp transa, GemmOp transb, size_t m, size_t n, size_t k)
Li Zhang's avatar
Li Zhang committed
314
315
{
    return fmtstr("%s [opA=%s, opB=%s, m=%ld, n=%ld, k=%ld]",
Chen Xin's avatar
Chen Xin committed
316
317
318
319
320
321
                  func_name,
                  getGemmOpString(transa).c_str(),
                  getGemmOpString(transb).c_str(),
                  m,
                  n,
                  k);
Li Zhang's avatar
Li Zhang committed
322
323
}

Chen Xin's avatar
Chen Xin committed
324
static inline std::string getTestName(const char* func_name, GemmOpPair op_pairs, size_t m, size_t n, size_t k)
Li Zhang's avatar
Li Zhang committed
325
326
327
328
329
330
331
{
    return getTestName(func_name, op_pairs.transa, op_pairs.transb, m, n, k);
}

/////////////////////////////////// Unittests //////////////////////////////////////////

template<typename T, DataType computeType>
Chen Xin's avatar
Chen Xin committed
332
333
334
335
void testGemmCorrectnessMatmul(size_t m, size_t n, size_t k)
{
    TM_LOG_INFO(
        "Matmul function correctness test [m=%ld, n=%ld, k=%ld, %s]", m, n, k, toString<T, computeType>().c_str());
Li Zhang's avatar
Li Zhang committed
336
337
338
339
340
    cudaStream_t stream;
    check_cuda_error(cudaStreamCreate(&stream));

    Allocator<AllocatorType::CUDA> allocator(getDevice());

Chen Xin's avatar
Chen Xin committed
341
    DataType      dtype = getTensorType<T>();
Li Zhang's avatar
Li Zhang committed
342
343
344
345
346
347
348
349
    TensorWrapper a_tensor(&allocator, dtype, {m, k}, false);
    TensorWrapper b_tensor(&allocator, dtype, {k, n}, false);
    TensorWrapper c_tensor(&allocator, dtype, {m, n}, true);
    TensorWrapper expected(&allocator, dtype, {m, n}, true);

    std::shared_ptr<Gemm> gemm = createGemm(&allocator, stream, false, false);
    gemm->setTypes(a_tensor.type, b_tensor.type, c_tensor.type, computeType);

Chen Xin's avatar
Chen Xin committed
350
    for (auto& op_pair : op_pairs) {
Li Zhang's avatar
Li Zhang committed
351
        std::string tc_name = getTestName(__func__, op_pair, m, n, k);
lvhan028's avatar
lvhan028 committed
352
        TM_LOG_DEBUG(tc_name);
Chen Xin's avatar
Chen Xin committed
353
        computeReference<computeType>(op_pair.transa, op_pair.transb, expected, a_tensor, b_tensor);
Li Zhang's avatar
Li Zhang committed
354
355
356
357
358

        size_t lda = (op_pair.transa == GEMM_OP_N) ? k : m;
        size_t ldb = (op_pair.transb == GEMM_OP_N) ? n : k;
        size_t ldc = n;

Chen Xin's avatar
Chen Xin committed
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
        c_tensor.setInvalidValues();  // to guarantee C has invalid data
        gemm->gemm(op_pair.transa,
                   op_pair.transb,
                   m,
                   n,
                   k,
                   a_tensor.data,
                   a_tensor.type,
                   lda,
                   b_tensor.data,
                   b_tensor.type,
                   ldb,
                   c_tensor.data,
                   c_tensor.type,
                   ldc);
Li Zhang's avatar
Li Zhang committed
374
375
376
        EXPECT_ALMOST_EQUAL(tc_name + " api1", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
377
        gemm->gemm(op_pair.transa, op_pair.transb, m, n, k, a_tensor.data, lda, b_tensor.data, ldb, c_tensor.data, ldc);
Li Zhang's avatar
Li Zhang committed
378
379
380
        EXPECT_ALMOST_EQUAL(tc_name + " api2", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
381
        gemm->gemm(op_pair.transa, op_pair.transb, m, n, k, a_tensor.data, b_tensor.data, c_tensor.data);
Li Zhang's avatar
Li Zhang committed
382
383
384
        EXPECT_ALMOST_EQUAL(tc_name + " api3", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
385
386
387
388
389
390
391
392
        gemm->gemm(op_pair.transa,
                   op_pair.transb,
                   m,
                   n,
                   k,
                   a_tensor.data,
                   DenseWeight<T>{(const T*)b_tensor.data, nullptr, nullptr},
                   c_tensor.data);
Li Zhang's avatar
Li Zhang committed
393
394
395
396
397
398
        EXPECT_ALMOST_EQUAL(tc_name + " api4", T, computeType, c_tensor, expected);
    }
    check_cuda_error(cudaStreamDestroy(stream));
}

template<typename T, DataType computeType>
Chen Xin's avatar
Chen Xin committed
399
400
void testGemmConsistencyMatmul(size_t m, size_t n, size_t k)
{
Li Zhang's avatar
Li Zhang committed
401
    // Test if Gemm is consistent with cublasWrapper
Chen Xin's avatar
Chen Xin committed
402
403
    TM_LOG_INFO(
        "Matmul function consistency test [m=%ld, n=%ld, k=%ld, %s]", m, n, k, toString<T, computeType>().c_str());
Li Zhang's avatar
Li Zhang committed
404
405

    Allocator<AllocatorType::CUDA> allocator(getDevice());
Chen Xin's avatar
Chen Xin committed
406
    cudaStream_t                   stream;
Li Zhang's avatar
Li Zhang committed
407
408
    check_cuda_error(cudaStreamCreate(&stream));

Chen Xin's avatar
Chen Xin committed
409
    DataType      dtype = getTensorType<T>();
Li Zhang's avatar
Li Zhang committed
410
411
412
413
414
    TensorWrapper a_tensor(&allocator, dtype, {m, k}, false);
    TensorWrapper b_tensor(&allocator, dtype, {k, n}, false);
    TensorWrapper c_tensor(&allocator, dtype, {m, n}, true);
    TensorWrapper expected(&allocator, dtype, {m, n}, true);

Chen Xin's avatar
Chen Xin committed
415
    cublasHandle_t   cublas_handle;
Li Zhang's avatar
Li Zhang committed
416
417
418
419
    cublasLtHandle_t cublaslt_handle;
    check_cuda_error(cublasCreate(&cublas_handle));
    check_cuda_error(cublasLtCreate(&cublaslt_handle));
    check_cuda_error(cublasSetStream(cublas_handle, stream));
Chen Xin's avatar
Chen Xin committed
420
421
422
423
    cublasAlgoMap   cublas_algo_map(GEMM_CONFIG);
    std::mutex*     cublas_wrapper_mutex = new std::mutex();
    cublasMMWrapper cublas_wrapper(
        cublas_handle, cublaslt_handle, stream, &cublas_algo_map, cublas_wrapper_mutex, &allocator);
Li Zhang's avatar
Li Zhang committed
424
425
426
427
428
429
430
431

    cudaDataType_t cuda_dtype = std::is_same<float, T>::value ? CUDA_R_32F : CUDA_R_16F;
    cudaDataType_t cuda_ctype = (DataType::TYPE_FP32 == computeType) ? CUDA_R_32F : CUDA_R_16F;
    cublas_wrapper.setGemmConfig(cuda_dtype, cuda_dtype, cuda_dtype, cuda_ctype);

    std::shared_ptr<Gemm> gemm = createGemm(&allocator, stream, false, false);
    gemm->setTypes(a_tensor.type, b_tensor.type, c_tensor.type, computeType);

Chen Xin's avatar
Chen Xin committed
432
    for (auto& op_pair : op_pairs) {
Li Zhang's avatar
Li Zhang committed
433
434
435
436
437
438
439
440
        std::string tc_name = getTestName(__func__, op_pair, m, n, k);

        // Switch A/B because Gemm expects column major layout as cublas does.
        size_t lda = (op_pair.transa == GEMM_OP_N) ? k : m;
        size_t ldb = (op_pair.transb == GEMM_OP_N) ? n : k;
        size_t ldc = n;
        cublas_wrapper.Gemm(getCublasOperation(op_pair.transb),
                            getCublasOperation(op_pair.transa),
Chen Xin's avatar
Chen Xin committed
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
                            n,
                            m,
                            k,
                            b_tensor.data,
                            ldb,
                            a_tensor.data,
                            lda,
                            expected.data,
                            ldc);

        c_tensor.setInvalidValues();  // to guarantee C has invalid data
        gemm->gemm(op_pair.transa,
                   op_pair.transb,
                   m,
                   n,
                   k,
                   a_tensor.data,
                   a_tensor.type,
                   lda,
                   b_tensor.data,
                   b_tensor.type,
                   ldb,
                   c_tensor.data,
                   c_tensor.type,
                   ldc);
Li Zhang's avatar
Li Zhang committed
466
467
468
        EXPECT_ALMOST_EQUAL(tc_name + " api1", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
469
        gemm->gemm(op_pair.transa, op_pair.transb, m, n, k, a_tensor.data, lda, b_tensor.data, ldb, c_tensor.data, ldc);
Li Zhang's avatar
Li Zhang committed
470
471
472
        EXPECT_ALMOST_EQUAL(tc_name + " api2", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
473
        gemm->gemm(op_pair.transa, op_pair.transb, m, n, k, a_tensor.data, b_tensor.data, c_tensor.data);
Li Zhang's avatar
Li Zhang committed
474
475
476
        EXPECT_ALMOST_EQUAL(tc_name + " api3", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
477
478
479
480
481
482
483
484
        gemm->gemm(op_pair.transa,
                   op_pair.transb,
                   m,
                   n,
                   k,
                   a_tensor.data,
                   DenseWeight<T>{(const T*)b_tensor.data, nullptr, nullptr},
                   c_tensor.data);
Li Zhang's avatar
Li Zhang committed
485
486
487
488
489
490
491
492
493
494
        EXPECT_ALMOST_EQUAL(tc_name + " api4", T, computeType, c_tensor, expected);
    }

    delete cublas_wrapper_mutex;
    check_cuda_error(cublasLtDestroy(cublaslt_handle));
    check_cuda_error(cublasDestroy(cublas_handle));
    check_cuda_error(cudaStreamDestroy(stream));
}

template<typename T, DataType computeType>
Chen Xin's avatar
Chen Xin committed
495
496
void testGemmConsistencyBatchedMatmul(size_t m, size_t n, size_t k)
{
Li Zhang's avatar
Li Zhang committed
497
    // Test if Gemm is consistent with cublasWrapper
lvhan028's avatar
lvhan028 committed
498
    TM_LOG_INFO("Batched gemm function consistency test [m=%ld, n=%ld, k=%ld, %s]",
Chen Xin's avatar
Chen Xin committed
499
500
501
502
                m,
                n,
                k,
                toString<T, computeType>().c_str());
Li Zhang's avatar
Li Zhang committed
503
504

    Allocator<AllocatorType::CUDA> allocator(getDevice());
Chen Xin's avatar
Chen Xin committed
505
    cudaStream_t                   stream;
Li Zhang's avatar
Li Zhang committed
506
507
508
    check_cuda_error(cudaStreamCreate(&stream));

    // batch of in/out tensors
Chen Xin's avatar
Chen Xin committed
509
510
511
    DataType                    a_type = getTensorType<T>();
    DataType                    b_type = getTensorType<T>();
    DataType                    c_type = getTensorType<T>();
Li Zhang's avatar
Li Zhang committed
512
513
514
515
    std::vector<TensorWrapper*> a_tensors;
    std::vector<TensorWrapper*> b_tensors;
    std::vector<TensorWrapper*> c_tensors;
    std::vector<TensorWrapper*> expecteds;
Chen Xin's avatar
Chen Xin committed
516
    const size_t                batch_size = 3;
Li Zhang's avatar
Li Zhang committed
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
    for (size_t i = 0; i < batch_size; ++i) {
        a_tensors.push_back(new TensorWrapper(&allocator, a_type, {m, k}, false));
        b_tensors.push_back(new TensorWrapper(&allocator, b_type, {k, n}, false));
        c_tensors.push_back(new TensorWrapper(&allocator, c_type, {m, n}, true));
        expecteds.push_back(new TensorWrapper(&allocator, c_type, {m, n}, true));
    }

    const T* hA[]{(const T*)a_tensors[0]->data,
                  (const T*)a_tensors[1]->data,
                  (const T*)a_tensors[2]->data,
                  nullptr,  // for memory alignment.
                  (const T*)b_tensors[0]->data,
                  (const T*)b_tensors[1]->data,
                  (const T*)b_tensors[2]->data,
                  nullptr,  // for memory alignment.
                  (const T*)c_tensors[0]->data,
                  (const T*)c_tensors[1]->data,
                  (const T*)c_tensors[2]->data,
                  nullptr,  // for memory alignment.
                  (const T*)expecteds[0]->data,
                  (const T*)expecteds[1]->data,
                  (const T*)expecteds[2]->data};

    T** batch_tensor_ptrs = reinterpret_cast<T**>(allocator.malloc(sizeof(T*) * 16, false));
Chen Xin's avatar
Chen Xin committed
541
542
543
544
545
    check_cuda_error(cudaMemcpyAsync((void*)batch_tensor_ptrs, hA, sizeof(T*) * 16, cudaMemcpyHostToDevice, stream));
    const void* const* batch_a        = reinterpret_cast<const void* const*>(batch_tensor_ptrs);
    const void* const* batch_b        = reinterpret_cast<const void* const*>(batch_tensor_ptrs + 4);
    void* const*       batch_c        = reinterpret_cast<void* const*>(batch_tensor_ptrs + 8);
    void* const*       batch_expected = reinterpret_cast<void* const*>(batch_tensor_ptrs + 12);
Li Zhang's avatar
Li Zhang committed
546

Chen Xin's avatar
Chen Xin committed
547
    cublasHandle_t   cublas_handle;
Li Zhang's avatar
Li Zhang committed
548
549
550
551
    cublasLtHandle_t cublaslt_handle;
    check_cuda_error(cublasCreate(&cublas_handle));
    check_cuda_error(cublasLtCreate(&cublaslt_handle));
    check_cuda_error(cublasSetStream(cublas_handle, stream));
Chen Xin's avatar
Chen Xin committed
552
553
554
555
    cublasAlgoMap   cublas_algo_map(GEMM_CONFIG);
    std::mutex*     cublas_wrapper_mutex = new std::mutex();
    cublasMMWrapper cublas_wrapper(
        cublas_handle, cublaslt_handle, stream, &cublas_algo_map, cublas_wrapper_mutex, &allocator);
Li Zhang's avatar
Li Zhang committed
556
557
558
559
560
561
562
563

    cudaDataType_t dtype = std::is_same<float, T>::value ? CUDA_R_32F : CUDA_R_16F;
    cudaDataType_t ctype = (computeType == DataType::TYPE_FP32) ? CUDA_R_32F : CUDA_R_16F;
    cublas_wrapper.setGemmConfig(dtype, dtype, dtype, ctype);

    std::shared_ptr<Gemm> gemm = createGemm(&allocator, stream, false, false);
    gemm->setTypes(a_type, b_type, c_type, computeType);

Chen Xin's avatar
Chen Xin committed
564
    for (auto& op_pair : op_pairs) {
Li Zhang's avatar
Li Zhang committed
565
        std::string tc_name = getTestName(__func__, op_pair, m, n, k);
lvhan028's avatar
lvhan028 committed
566
        TM_LOG_DEBUG(tc_name);
Li Zhang's avatar
Li Zhang committed
567
568
569
570
571
572
573
574
575
576
577

        size_t lda = (op_pair.transa == GEMM_OP_N) ? k : m;
        size_t ldb = (op_pair.transb == GEMM_OP_N) ? n : k;
        size_t ldc = n;

        // Switch A/B because Gemm expects column major layout as cublas does.
        cublas_wrapper.batchedGemm(getCublasOperation(op_pair.transb),  // N
                                   getCublasOperation(op_pair.transa),  // T
                                   n,
                                   m,
                                   k,
Chen Xin's avatar
Chen Xin committed
578
579
580
581
582
583
                                   (const void* const*)batch_b,
                                   ldb,
                                   (const void* const*)batch_a,
                                   lda,
                                   (void* const*)batch_expected,
                                   ldc,
Li Zhang's avatar
Li Zhang committed
584
585
                                   batch_size);

Chen Xin's avatar
Chen Xin committed
586
587
588
589
590
591
592
593
594
595
596
597
598
599
        gemm->batchedGemm(op_pair.transa,
                          op_pair.transb,
                          m,
                          n,
                          k,
                          batch_a,
                          a_type,
                          lda,
                          batch_b,
                          b_type,
                          ldb,
                          batch_c,
                          c_type,
                          ldc,
Li Zhang's avatar
Li Zhang committed
600
601
                          batch_size);
        for (size_t i = 0; i < batch_size; ++i) {
Chen Xin's avatar
Chen Xin committed
602
603
            EXPECT_ALMOST_EQUAL(
                tc_name + " api1 batch" + std::to_string(i), T, computeType, *c_tensors[i], *expecteds[i]);
Li Zhang's avatar
Li Zhang committed
604
605
606
607
608
        }

        for (size_t i = 0; i < batch_size; ++i) {
            c_tensors[i]->setInvalidValues();
        }
Chen Xin's avatar
Chen Xin committed
609
610
        gemm->batchedGemm(
            op_pair.transa, op_pair.transb, m, n, k, batch_a, lda, batch_b, ldb, batch_c, ldc, batch_size);
Li Zhang's avatar
Li Zhang committed
611
        for (size_t i = 0; i < batch_size; ++i) {
Chen Xin's avatar
Chen Xin committed
612
613
            EXPECT_ALMOST_EQUAL(
                tc_name + " api2 batch" + std::to_string(i), T, computeType, *c_tensors[i], *expecteds[i]);
Li Zhang's avatar
Li Zhang committed
614
615
616
617
618
        }

        for (size_t i = 0; i < batch_size; ++i) {
            c_tensors[i]->setInvalidValues();
        }
Chen Xin's avatar
Chen Xin committed
619
        gemm->batchedGemm(op_pair.transa, op_pair.transb, m, n, k, batch_a, batch_b, batch_c, batch_size);
Li Zhang's avatar
Li Zhang committed
620
        for (size_t i = 0; i < batch_size; ++i) {
Chen Xin's avatar
Chen Xin committed
621
622
            EXPECT_ALMOST_EQUAL(
                tc_name + " api3 batch" + std::to_string(i), T, computeType, *c_tensors[i], *expecteds[i]);
Li Zhang's avatar
Li Zhang committed
623
624
625
626
627
628
629
630
631
632
633
634
635
        }
    }
    a_tensors.clear();
    b_tensors.clear();
    c_tensors.clear();
    expecteds.clear();
    delete cublas_wrapper_mutex;
    check_cuda_error(cublasLtDestroy(cublaslt_handle));
    check_cuda_error(cublasDestroy(cublas_handle));
    check_cuda_error(cudaStreamDestroy(stream));
}

template<typename T, DataType computeType>
Chen Xin's avatar
Chen Xin committed
636
637
void testGemmConsistencyStridedBatchedMatmul(size_t batch_size, size_t m, size_t n, size_t k)
{
Li Zhang's avatar
Li Zhang committed
638
    // Test if Gemm is consistent with cublasWrapper
lvhan028's avatar
lvhan028 committed
639
    TM_LOG_INFO("Strided batched gemm function consistency test [bsz=%ld, m=%ld, n=%ld, k=%ld, %s]",
Chen Xin's avatar
Chen Xin committed
640
641
642
643
644
                batch_size,
                m,
                n,
                k,
                toString<T, computeType>().c_str());
Li Zhang's avatar
Li Zhang committed
645
646

    Allocator<AllocatorType::CUDA> allocator(getDevice());
Chen Xin's avatar
Chen Xin committed
647
    cudaStream_t                   stream;
Li Zhang's avatar
Li Zhang committed
648
649
    check_cuda_error(cudaStreamCreate(&stream));

Chen Xin's avatar
Chen Xin committed
650
    DataType      data_type = getTensorType<T>();
Li Zhang's avatar
Li Zhang committed
651
652
653
654
655
    TensorWrapper a_tensor(&allocator, data_type, {batch_size, m, k}, false);
    TensorWrapper b_tensor(&allocator, data_type, {batch_size, k, n}, false);
    TensorWrapper c_tensor(&allocator, data_type, {batch_size, m, n}, true);
    TensorWrapper expected(&allocator, data_type, {batch_size, m, n}, true);

Chen Xin's avatar
Chen Xin committed
656
    cublasHandle_t   cublas_handle;
Li Zhang's avatar
Li Zhang committed
657
658
659
660
    cublasLtHandle_t cublaslt_handle;
    check_cuda_error(cublasCreate(&cublas_handle));
    check_cuda_error(cublasLtCreate(&cublaslt_handle));
    check_cuda_error(cublasSetStream(cublas_handle, stream));
Chen Xin's avatar
Chen Xin committed
661
662
663
664
    cublasAlgoMap   cublas_algo_map(GEMM_CONFIG);
    std::mutex*     cublas_wrapper_mutex = new std::mutex();
    cublasMMWrapper cublas_wrapper(
        cublas_handle, cublaslt_handle, stream, &cublas_algo_map, cublas_wrapper_mutex, &allocator);
Li Zhang's avatar
Li Zhang committed
665
666
667
668
669
670
671
672

    cudaDataType_t dtype = std::is_same<float, T>::value ? CUDA_R_32F : CUDA_R_16F;
    cudaDataType_t ctype = (computeType == DataType::TYPE_FP32) ? CUDA_R_32F : CUDA_R_16F;
    cublas_wrapper.setGemmConfig(dtype, dtype, dtype, ctype);

    std::shared_ptr<Gemm> gemm = createGemm(&allocator, stream, false, false);
    gemm->setTypes(a_tensor.type, b_tensor.type, c_tensor.type, computeType);

Chen Xin's avatar
Chen Xin committed
673
    for (auto& op_pair : op_pairs) {
Li Zhang's avatar
Li Zhang committed
674
675
676
677
678
679
680
681
682
683
684
685
        std::string tc_name = getTestName(__func__, op_pair, m, n, k);

        // Switch A/B because Gemm expects column major layout as cublas does.
        size_t lda = (op_pair.transa == GEMM_OP_N) ? k : m;
        size_t ldb = (op_pair.transb == GEMM_OP_N) ? n : k;
        size_t ldc = n;

        int64_t stridea = m * k;
        int64_t strideb = k * n;
        int64_t stridec = m * n;

        float alpha = 1.0f;
Chen Xin's avatar
Chen Xin committed
686
        float beta  = 0.0f;
Li Zhang's avatar
Li Zhang committed
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710

        cublas_wrapper.stridedBatchedGemm(getCublasOperation(op_pair.transb),
                                          getCublasOperation(op_pair.transa),
                                          n,
                                          m,
                                          k,
                                          alpha,
                                          b_tensor.data,
                                          getCublasDataType(b_tensor.type),
                                          ldb,
                                          strideb,
                                          a_tensor.data,
                                          getCublasDataType(a_tensor.type),
                                          lda,
                                          stridea,
                                          beta,
                                          expected.data,
                                          getCublasDataType(expected.type),
                                          ldc,
                                          stridec,
                                          batch_size,
                                          getCublasDataType(computeType));

        c_tensor.setInvalidValues();  // to guarantee C has invalid data
Chen Xin's avatar
Chen Xin committed
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
        gemm->stridedBatchedGemm(op_pair.transa,
                                 op_pair.transb,
                                 m,
                                 n,
                                 k,
                                 a_tensor.data,
                                 a_tensor.type,
                                 lda,
                                 stridea,
                                 b_tensor.data,
                                 b_tensor.type,
                                 ldb,
                                 strideb,
                                 c_tensor.data,
                                 c_tensor.type,
                                 ldc,
                                 stridec,
                                 batch_size,
                                 computeType,
                                 alpha,
                                 beta);
Li Zhang's avatar
Li Zhang committed
732
733
734
        EXPECT_ALMOST_EQUAL(tc_name + " api1", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
        gemm->stridedBatchedGemm(op_pair.transa,
                                 op_pair.transb,
                                 m,
                                 n,
                                 k,
                                 a_tensor.data,
                                 lda,
                                 stridea,
                                 b_tensor.data,
                                 ldb,
                                 strideb,
                                 c_tensor.data,
                                 ldc,
                                 stridec,
                                 batch_size,
                                 alpha,
                                 beta);
Li Zhang's avatar
Li Zhang committed
752
753
754
        EXPECT_ALMOST_EQUAL(tc_name + " api2", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
755
756
757
758
759
760
761
762
763
764
765
766
767
768
        gemm->stridedBatchedGemm(op_pair.transa,
                                 op_pair.transb,
                                 m,
                                 n,
                                 k,
                                 a_tensor.data,
                                 stridea,
                                 b_tensor.data,
                                 strideb,
                                 c_tensor.data,
                                 stridec,
                                 batch_size,
                                 alpha,
                                 beta);
Li Zhang's avatar
Li Zhang committed
769
770
771
        EXPECT_ALMOST_EQUAL(tc_name + " api3", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
772
773
774
775
776
        gemm->stridedBatchedGemm(op_pair.transa,
                                 op_pair.transb,
                                 m,
                                 n,
                                 k,
Li Zhang's avatar
Li Zhang committed
777
778
779
                                 a_tensor.data,
                                 b_tensor.data,
                                 c_tensor.data,
Chen Xin's avatar
Chen Xin committed
780
781
782
                                 batch_size,
                                 alpha,
                                 beta);
Li Zhang's avatar
Li Zhang committed
783
784
785
786
787
788
789
790
791
792
793
794
795
        EXPECT_ALMOST_EQUAL(tc_name + " api4", T, computeType, c_tensor, expected);
    }

    delete cublas_wrapper_mutex;
    check_cuda_error(cublasLtDestroy(cublaslt_handle));
    check_cuda_error(cublasDestroy(cublas_handle));
    check_cuda_error(cudaStreamDestroy(stream));
}

#ifdef SPARSITY_ENABLED
// The current SpGemm only supports TYPE_FP16 for T, computeType,
// but let us keep these template variables for later use.
template<typename T, DataType computeType>
Chen Xin's avatar
Chen Xin committed
796
797
798
799
void testSpGemmCorrectnessMatmul(size_t m, size_t n, size_t k)
{
    TM_LOG_INFO(
        "Sparse gemm function correctness test [m=%ld, n=%ld, k=%ld, %s]", m, n, k, toString<T, computeType>().c_str());
Li Zhang's avatar
Li Zhang committed
800
801
802
803
804
    cudaStream_t stream;
    check_cuda_error(cudaStreamCreate(&stream));

    Allocator<AllocatorType::CUDA> allocator(getDevice());

Chen Xin's avatar
Chen Xin committed
805
    DataType      dtype = getTensorType<T>();
Li Zhang's avatar
Li Zhang committed
806
807
808
809
810
811
812
813
    TensorWrapper a_tensor(&allocator, dtype, {m, k}, false);
    TensorWrapper b_tensor(&allocator, dtype, {k, n}, false);
    TensorWrapper c_tensor(&allocator, dtype, {m, n}, true);
    TensorWrapper expected(&allocator, dtype, {m, n}, true);

    std::shared_ptr<Gemm> gemm = createGemm(&allocator, stream, true, false);
    gemm->setTypes(a_tensor.type, b_tensor.type, c_tensor.type, computeType);

Chen Xin's avatar
Chen Xin committed
814
    for (auto& op_pair : op_pairs) {
Li Zhang's avatar
Li Zhang committed
815
816
        // A/B will be switched in SpGemm.
        std::string tc_name = getTestName(__func__, op_pair, m, n, k);
lvhan028's avatar
lvhan028 committed
817
        TM_LOG_DEBUG(tc_name);
Li Zhang's avatar
Li Zhang committed
818
819

        b_tensor.setRandomValues();
Chen Xin's avatar
Chen Xin committed
820
821
        pruneMatrixB(b_tensor.data, stream, b_tensor.shape[0], b_tensor.shape[1], op_pair.transb);
        computeReference<computeType>(op_pair.transa, op_pair.transb, expected, a_tensor, b_tensor);
Li Zhang's avatar
Li Zhang committed
822
823

        void* b_compressed;
Chen Xin's avatar
Chen Xin committed
824
825
        compressMatrixB(
            &b_compressed, allocator, stream, b_tensor.data, b_tensor.shape[0], b_tensor.shape[1], op_pair.transb);
Li Zhang's avatar
Li Zhang committed
826
827
828
829
830

        size_t lda = (op_pair.transa == GEMM_OP_N) ? k : m;
        size_t ldb = (op_pair.transb == GEMM_OP_N) ? n : k;
        size_t ldc = n;

Chen Xin's avatar
Chen Xin committed
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
        c_tensor.setInvalidValues();  // to guarantee C has invalid data
        gemm->gemm(op_pair.transa,
                   op_pair.transb,
                   m,
                   n,
                   k,
                   a_tensor.data,
                   a_tensor.type,
                   lda,
                   b_compressed,
                   b_tensor.type,
                   ldb,
                   c_tensor.data,
                   c_tensor.type,
                   ldc);
Li Zhang's avatar
Li Zhang committed
846
847
848
        EXPECT_ALMOST_EQUAL(tc_name + " api1", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
849
        gemm->gemm(op_pair.transa, op_pair.transb, m, n, k, a_tensor.data, lda, b_compressed, ldb, c_tensor.data, ldc);
Li Zhang's avatar
Li Zhang committed
850
851
852
        EXPECT_ALMOST_EQUAL(tc_name + " api2", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
853
        gemm->gemm(op_pair.transa, op_pair.transb, m, n, k, a_tensor.data, b_compressed, c_tensor.data);
Li Zhang's avatar
Li Zhang committed
854
855
856
        EXPECT_ALMOST_EQUAL(tc_name + " api3", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
857
858
859
860
861
        gemm->gemm(op_pair.transa,
                   op_pair.transb,
                   m,
                   n,
                   k,
Li Zhang's avatar
Li Zhang committed
862
863
864
865
866
867
868
869
870
871
872
                   a_tensor.data,
                   DenseWeight<T>{(const T*)b_tensor.data, nullptr, (const T*)b_compressed},
                   c_tensor.data);
        EXPECT_ALMOST_EQUAL(tc_name + " api4", T, computeType, c_tensor, expected);

        allocator.free((void**)(&b_compressed));
    }
    check_cuda_error(cudaStreamDestroy(stream));
}

template<typename T, DataType computeType>
Chen Xin's avatar
Chen Xin committed
873
874
void testSpGemmConsistencyMatmul(size_t m, size_t n, size_t k)
{
Li Zhang's avatar
Li Zhang committed
875
    // Test if Gemm is consistent with cublasWrapper
lvhan028's avatar
lvhan028 committed
876
    TM_LOG_INFO("Sparse Matmul function consistency test [m=%ld, n=%ld, k=%ld, %s]",
Chen Xin's avatar
Chen Xin committed
877
878
879
880
                m,
                n,
                k,
                toString<T, computeType>().c_str());
Li Zhang's avatar
Li Zhang committed
881
882

    Allocator<AllocatorType::CUDA> allocator(getDevice());
Chen Xin's avatar
Chen Xin committed
883
    cudaStream_t                   stream;
Li Zhang's avatar
Li Zhang committed
884
885
    check_cuda_error(cudaStreamCreate(&stream));

Chen Xin's avatar
Chen Xin committed
886
    DataType      dtype = getTensorType<T>();
Li Zhang's avatar
Li Zhang committed
887
888
889
890
891
    TensorWrapper a_tensor(&allocator, dtype, {m, k}, false);
    TensorWrapper b_tensor(&allocator, dtype, {k, n}, false);
    TensorWrapper c_tensor(&allocator, dtype, {m, n}, true);
    TensorWrapper expected(&allocator, dtype, {m, n}, true);

Chen Xin's avatar
Chen Xin committed
892
    cublasHandle_t   cublas_handle;
Li Zhang's avatar
Li Zhang committed
893
894
895
896
    cublasLtHandle_t cublaslt_handle;
    check_cuda_error(cublasCreate(&cublas_handle));
    check_cuda_error(cublasLtCreate(&cublaslt_handle));
    check_cuda_error(cublasSetStream(cublas_handle, stream));
Chen Xin's avatar
Chen Xin committed
897
898
899
900
    cublasAlgoMap   cublas_algo_map(GEMM_CONFIG);
    std::mutex*     cublas_wrapper_mutex = new std::mutex();
    cublasMMWrapper cublas_wrapper(
        cublas_handle, cublaslt_handle, stream, &cublas_algo_map, cublas_wrapper_mutex, &allocator);
Li Zhang's avatar
Li Zhang committed
901
902
903
904
905
906
907
908

    cudaDataType_t cu_dtype = std::is_same<float, T>::value ? CUDA_R_32F : CUDA_R_16F;
    cudaDataType_t cu_ctype = (DataType::TYPE_FP32 == computeType) ? CUDA_R_32F : CUDA_R_16F;
    cublas_wrapper.setGemmConfig(cu_dtype, cu_dtype, cu_dtype, cu_ctype);

    std::shared_ptr<Gemm> gemm = createGemm(&allocator, stream, true, false);
    gemm->setTypes(a_tensor.type, b_tensor.type, c_tensor.type, computeType);

Chen Xin's avatar
Chen Xin committed
909
    for (auto& op_pair : op_pairs) {
Li Zhang's avatar
Li Zhang committed
910
        std::string tc_name = getTestName(__func__, op_pair, m, n, k);
lvhan028's avatar
lvhan028 committed
911
        TM_LOG_DEBUG(tc_name);
Li Zhang's avatar
Li Zhang committed
912
913

        b_tensor.setRandomValues();
Chen Xin's avatar
Chen Xin committed
914
        pruneMatrixB(b_tensor.data, stream, b_tensor.shape[0], b_tensor.shape[1], op_pair.transb);
Li Zhang's avatar
Li Zhang committed
915
916
917
918
919
920
921
922
923
924

        // Switch A/B because Gemm expects column major layout as cublas does.
        size_t lda = (op_pair.transa == GEMM_OP_N) ? k : m;
        size_t ldb = (op_pair.transb == GEMM_OP_N) ? n : k;
        size_t ldc = n;
        cublas_wrapper.Gemm(getCublasOperation(op_pair.transb),
                            getCublasOperation(op_pair.transa),
                            n,
                            m,
                            k,
Chen Xin's avatar
Chen Xin committed
925
926
927
928
929
930
                            b_tensor.data,
                            ldb,
                            a_tensor.data,
                            lda,
                            expected.data,
                            ldc);
Li Zhang's avatar
Li Zhang committed
931
932

        void* b_compressed;
Chen Xin's avatar
Chen Xin committed
933
934
        compressMatrixB(
            &b_compressed, allocator, stream, b_tensor.data, b_tensor.shape[0], b_tensor.shape[1], op_pair.transb);
Li Zhang's avatar
Li Zhang committed
935
936

        c_tensor.setInvalidValues();  // to guarantee C has invalid data
Chen Xin's avatar
Chen Xin committed
937
938
939
940
941
942
943
944
945
946
947
948
949
950
        gemm->gemm(op_pair.transa,
                   op_pair.transb,
                   m,
                   n,
                   k,
                   a_tensor.data,
                   a_tensor.type,
                   lda,
                   b_compressed,
                   b_tensor.type,
                   ldb,
                   c_tensor.data,
                   c_tensor.type,
                   ldc);
Li Zhang's avatar
Li Zhang committed
951
952
953
        EXPECT_ALMOST_EQUAL(tc_name + " api1", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
954
        gemm->gemm(op_pair.transa, op_pair.transb, m, n, k, a_tensor.data, lda, b_compressed, ldb, c_tensor.data, ldc);
Li Zhang's avatar
Li Zhang committed
955
956
957
        EXPECT_ALMOST_EQUAL(tc_name + " api1", T, computeType, c_tensor, expected);

        c_tensor.setInvalidValues();
Chen Xin's avatar
Chen Xin committed
958
        gemm->gemm(op_pair.transa, op_pair.transb, m, n, k, a_tensor.data, b_compressed, c_tensor.data);
Li Zhang's avatar
Li Zhang committed
959
960
961
962
963
964
965
966
967
968
        EXPECT_ALMOST_EQUAL(tc_name + " api3", T, computeType, c_tensor, expected);
    }

    delete cublas_wrapper_mutex;
    check_cuda_error(cublasLtDestroy(cublaslt_handle));
    check_cuda_error(cublasDestroy(cublas_handle));
    check_cuda_error(cudaStreamDestroy(stream));
}
#endif

Chen Xin's avatar
Chen Xin committed
969
970
int main(int argc, char* argv[])
{
Li Zhang's avatar
Li Zhang committed
971
972
973
    // testGemmCreate();
    using testcase_t = std::tuple<size_t, size_t, size_t>;

Chen Xin's avatar
Chen Xin committed
974
975
    std::vector<testcase_t> testcases = {
        {16, 32, 64}, {255, 255, 255}, {1041, 2047, 9999}, {1041, 1, 9999}, {1041, 999, 1}};
Li Zhang's avatar
Li Zhang committed
976
977

    // Computation correctness tests
Chen Xin's avatar
Chen Xin committed
978
    for (testcase_t& tc : testcases) {
Li Zhang's avatar
Li Zhang committed
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
        size_t m = std::get<0>(tc);
        size_t n = std::get<1>(tc);
        size_t k = std::get<2>(tc);

        testGemmCorrectnessMatmul<float, TYPE_FP32>(m, n, k);
        testGemmCorrectnessMatmul<half, TYPE_FP32>(m, n, k);
        testGemmCorrectnessMatmul<half, TYPE_FP16>(m, n, k);

        testGemmConsistencyMatmul<float, TYPE_FP32>(m, n, k);
        testGemmConsistencyMatmul<half, TYPE_FP32>(m, n, k);
        testGemmConsistencyMatmul<half, TYPE_FP16>(m, n, k);

        testGemmConsistencyBatchedMatmul<float, TYPE_FP32>(m, n, k);
        testGemmConsistencyBatchedMatmul<half, TYPE_FP32>(m, n, k);
        testGemmConsistencyBatchedMatmul<half, TYPE_FP16>(m, n, k);

        testGemmConsistencyStridedBatchedMatmul<float, TYPE_FP32>(7, m, n, k);
        testGemmConsistencyStridedBatchedMatmul<half, TYPE_FP32>(7, m, n, k);
        testGemmConsistencyStridedBatchedMatmul<half, TYPE_FP16>(7, m, n, k);
    }

#ifdef SPARSITY_ENABLED
    // Reset for SpGemm test.
    testcases.clear();
    testcases.insert(testcases.end(),
Chen Xin's avatar
Chen Xin committed
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
                     {{8, 32, 32},  // minimum possible example.
                      {8, 32, 64},
                      {64, 64, 64},
                      {16, 32, 64},
                      {1024, 32, 1024},
                      {1024, 1024, 32},
                      {16, 1024, 1024},
                      {1024, 1024, 1024}});

    for (testcase_t& tc : testcases) {
Li Zhang's avatar
Li Zhang committed
1014
1015
1016
1017
1018
1019
1020
        size_t m = std::get<0>(tc);
        size_t n = std::get<1>(tc);
        size_t k = std::get<2>(tc);
        testSpGemmCorrectnessMatmul<half, TYPE_FP16>(m, n, k);
        testSpGemmConsistencyMatmul<half, TYPE_FP16>(m, n, k);
    }
#endif
lvhan028's avatar
lvhan028 committed
1021
    TM_LOG_INFO("Test done");
Li Zhang's avatar
Li Zhang committed
1022
1023
    return 0;
}