OpenI
/
MegEngine
Not watched
1
0
Fork 0
Code
Releases 31 Wiki Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Tree: c42ce93705
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'c42ce93705'
${ noResults }
MegEngine/dnn/test/cuda/cutlass_matmul.cpp

472 lines
20 kB
Raw Blame History 

  1. /**

  2.  * \file dnn/test/cuda/cutlass_matmul.cpp

  3.  * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")

  4.  *

  5.  * Copyright (c) 2014-2020 Megvii Inc. All rights reserved.

  6.  *

  7.  * Unless required by applicable law or agreed to in writing,

  8.  * software distributed under the License is distributed on an

  9.  * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or

  10.  * implied.

  11.  */

  12. #include <cuda.h>

  13. #include "megdnn/oprs/linalg.h"

  14. 

  15. #include "src/common/utils.h"

  16. #include "test/common/checker.h"

  17. #include "test/common/matrix_mul.h"

  18. #include "test/common/tensor.h"

  19. #include "test/common/workspace_wrapper.h"

  20. #include "test/cuda/benchmark.h"

  21. #include "test/cuda/fixture.h"

  22. #include "test/cuda/utils.h"

  23. 

  24. #if CUDA_VERSION >= 9020

  25. namespace megdnn {

  26. namespace test {

  27. namespace {

  28. void test_multibatchsize(

  29.         Handle* handle_cuda, DType A_dtype, DType B_dtype, DType C_dtype,

  30.         const char* algo, const std::vector<matrix_mul::TestArg>& args,

  31.         param::MatrixMul::Format format = param::MatrixMul::Format::DEFAULT,

  32.         const std::function<bool(const matrix_mul::TestArg&)>& filter = {}) {

  33.     Checker<MatrixMulForward> checker(handle_cuda, false);

  34.     if (algo) {

  35.         checker.set_before_exec_callback(AlgoChecker<MatrixMulForward>(algo));

  36.     }

  37.     std::unique_ptr<RNG> rng;

  38.     if (A_dtype.enumv() == DTypeEnum::Float32) {

  39.         rng = std::make_unique<UniformFloatRNG>(-1, 1);

  40.         megdnn_assert(

  41.                 B_dtype.enumv() == DTypeEnum::Float32 &&

  42.                 C_dtype.enumv() == DTypeEnum::Float32);

  43.     }

  44.     megdnn_assert(rng != nullptr);

  45. 

  46.     struct Compare {

  47.         bool is_same(dt_float32 expected, dt_float32 actual) const {

  48.             return expected == actual;

  49.         }

  50.     };

  51. 

  52.     // copy rhs->lhs, lhs is 8 times of rhs

  53.     auto copy = [](SyncedTensor<dt_float32, Compare>& lhs,

  54.                    SyncedTensor<dt_float32, Compare>& rhs) {

  55.         size_t chunk = rhs.layout().span().dist_byte();

  56.         size_t tot = lhs.layout().span().dist_byte();

  57.         megdnn_assert(tot % chunk == 0);

  58.         char* pointer_lhs = reinterpret_cast<char*>(lhs.ptr_mutable_host());

  59.         const char* pointer_rhs = reinterpret_cast<const char*>(rhs.ptr_host());

  60.         for (size_t i = 0; i < tot; i += chunk) {

  61.             std::memcpy(pointer_lhs + i, pointer_rhs, chunk);

  62.         }

  63.     };

  64.     using Param = param::MatrixMul;

  65.     megdnn_assert(format == Param::Format::DEFAULT);

  66.     for (auto&& arg : args) {

  67.         megdnn_assert(arg.mask == 0x0);

  68.         // make m, n, k big enough

  69.         size_t m = arg.m, n = (arg.n << 3), k = (arg.k << 3);

  70.         size_t m_prime = (m << 3);

  71.         if (filter && filter(arg))

  72.             continue;

  73.         TensorShape A{m, k}, B{k, n}, C{m, n};

  74.         TensorShape A_prime{m_prime, k}, C_prime{m_prime, n};

  75.         SyncedTensor<dt_float32, Compare> A_tensor{handle_cuda, {A, A_dtype}},

  76.                 B_tensor{handle_cuda, {B, B_dtype}},

  77.                 C_tensor{handle_cuda, {C, C_dtype}},

  78.                 A_tensor_prime{handle_cuda, {A_prime, A_dtype}},

  79.                 C_tensor_prime{handle_cuda, {C_prime, C_dtype}},

  80.                 C_tensor_batch{handle_cuda, {C_prime, C_dtype}};

  81.         rng->gen(A_tensor.tensornd_host());

  82.         rng->gen(B_tensor.tensornd_host());

  83.         copy(A_tensor_prime, A_tensor);

  84. 

  85.         auto opr_reference = handle_cuda->create_operator<MatrixMulForward>();

  86.         {

  87.             opr_reference->execution_policy().algo.reset();

  88.             for (auto i : opr_reference->get_all_algorithms_info_safe(

  89.                          A_tensor.layout(), B_tensor.layout(), C_tensor.layout())) {

  90.                 if (std::regex_match(

  91.                             i.desc.name.c_str(),

  92.                             std::regex("(" + std::string(algo) + ")(.*)"))) {

  93.                     opr_reference->execution_policy().algo = i.desc;

  94.                     break;

  95.                 }

  96.             }

  97.             megdnn_assert(opr_reference->execution_policy().algo.valid());

  98.             size_t ws_size = opr_reference->get_workspace_in_bytes(

  99.                     A_tensor.layout(), B_tensor.layout(), C_tensor.layout());

  100.             WorkspaceWrapper ws_reference(handle_cuda, ws_size);

  101.             opr_reference->exec(

  102.                     A_tensor.tensornd_dev(), B_tensor.tensornd_dev(),

  103.                     C_tensor.tensornd_dev(), ws_reference.workspace());

  104.         }

  105.         copy(C_tensor_prime, C_tensor);

  106.         checker.set_dtype(0, A_dtype)

  107.                 .set_dtype(1, B_dtype)

  108.                 .set_dtype(2, C_dtype)

  109.                 .set_epsilon(1e-6)

  110.                 .exect({A_tensor_prime.tensornd_host(), B_tensor.tensornd_host(), {}},

  111.                        {{}, {}, C_tensor_prime.tensornd_host()});

  112.         {

  113.             opr_reference->execution_policy().algo.reset();

  114.             for (auto i : opr_reference->get_all_algorithms_info_safe(

  115.                          A_tensor_prime.layout(), B_tensor.layout(),

  116.                          C_tensor_batch.layout())) {

  117.                 if (std::regex_match(

  118.                             i.desc.name.c_str(),

  119.                             std::regex("(" + std::string(algo) + ")(.*)"))) {

  120.                     opr_reference->execution_policy().algo = i.desc;

  121.                     break;

  122.                 }

  123.             }

  124.             megdnn_assert(opr_reference->execution_policy().algo.valid());

  125.             size_t ws_size = opr_reference->get_workspace_in_bytes(

  126.                     A_tensor_prime.layout(), B_tensor.layout(),

  127.                     C_tensor_batch.layout());

  128.             WorkspaceWrapper ws_reference(handle_cuda, ws_size);

  129.             opr_reference->exec(

  130.                     A_tensor_prime.tensornd_dev(), B_tensor.tensornd_dev(),

  131.                     C_tensor_batch.tensornd_dev(), ws_reference.workspace());

  132.         }

  133.         C_tensor_batch.check_with(C_tensor_prime);

  134.     }

  135. }

  136. 

  137. #if MEGDNN_WITH_BENCHMARK

  138. struct BenchArgs {

  139.     size_t m, n, k, mask = 0x0;

  140. };

  141. 

  142. std::vector<BenchArgs> get_square_matmul_args() {

  143.     std::vector<BenchArgs> args;

  144.     args.emplace_back(BenchArgs{128, 128, 128});

  145.     args.emplace_back(BenchArgs{256, 256, 256});

  146.     args.emplace_back(BenchArgs{512, 512, 512});

  147.     args.emplace_back(BenchArgs{1024, 1024, 1024});

  148.     args.emplace_back(BenchArgs{2048, 2048, 2048});

  149.     args.emplace_back(BenchArgs{4096, 4096, 4096});

  150. 

  151.     return args;

  152. }

  153. 

  154. std::vector<BenchArgs> get_feat_model_args() {

  155.     std::vector<BenchArgs> args;

  156. 

  157.     args.emplace_back(BenchArgs{2, 4096, 4096});

  158.     args.emplace_back(BenchArgs{2, 1024, 6912});

  159.     args.emplace_back(BenchArgs{2, 3456, 3456});

  160.     args.emplace_back(BenchArgs{2, 2304, 2304});

  161.     args.emplace_back(BenchArgs{1, 256, 8192});

  162.     args.emplace_back(BenchArgs{2, 864, 864});

  163.     args.emplace_back(BenchArgs{2, 9, 64});

  164. 

  165.     args.emplace_back(BenchArgs{4, 4096, 4096});

  166.     args.emplace_back(BenchArgs{4, 1024, 6912});

  167.     args.emplace_back(BenchArgs{4, 3456, 3456});

  168.     args.emplace_back(BenchArgs{4, 2304, 2304});

  169.     args.emplace_back(BenchArgs{2, 256, 8192});

  170.     args.emplace_back(BenchArgs{4, 864, 864});

  171.     args.emplace_back(BenchArgs{4, 9, 64});

  172. 

  173.     args.emplace_back(BenchArgs{8, 4096, 4096});

  174.     args.emplace_back(BenchArgs{8, 1024, 6912});

  175.     args.emplace_back(BenchArgs{8, 3456, 3456});

  176.     args.emplace_back(BenchArgs{8, 2304, 2304});

  177.     args.emplace_back(BenchArgs{4, 256, 8192});

  178.     args.emplace_back(BenchArgs{8, 864, 864});

  179.     args.emplace_back(BenchArgs{4, 9, 64});

  180. 

  181.     args.emplace_back(BenchArgs{16, 4096, 4096});

  182.     args.emplace_back(BenchArgs{16, 1024, 6912});

  183.     args.emplace_back(BenchArgs{16, 3456, 3456});

  184.     args.emplace_back(BenchArgs{16, 2304, 2304});

  185.     args.emplace_back(BenchArgs{8, 256, 8192});

  186.     args.emplace_back(BenchArgs{16, 864, 864});

  187.     args.emplace_back(BenchArgs{8, 9, 64});

  188. 

  189.     args.emplace_back(BenchArgs{32, 4096, 4096});

  190.     args.emplace_back(BenchArgs{32, 1024, 6912});

  191.     args.emplace_back(BenchArgs{32, 3456, 3456});

  192.     args.emplace_back(BenchArgs{32, 2304, 2304});

  193.     args.emplace_back(BenchArgs{16, 256, 8192});

  194.     args.emplace_back(BenchArgs{32, 864, 864});

  195.     args.emplace_back(BenchArgs{32, 9, 64});

  196. 

  197.     args.emplace_back(BenchArgs{64, 4096, 4096});

  198.     args.emplace_back(BenchArgs{64, 1024, 6912});

  199.     args.emplace_back(BenchArgs{64, 3456, 3456});

  200.     args.emplace_back(BenchArgs{64, 2304, 2304});

  201.     args.emplace_back(BenchArgs{32, 256, 8192});

  202.     args.emplace_back(BenchArgs{64, 864, 864});

  203.     args.emplace_back(BenchArgs{64, 9, 64});

  204. 

  205.     args.emplace_back(BenchArgs{128, 4096, 4096});

  206.     args.emplace_back(BenchArgs{128, 1024, 6912});

  207.     args.emplace_back(BenchArgs{128, 3456, 3456});

  208.     args.emplace_back(BenchArgs{128, 2304, 2304});

  209.     args.emplace_back(BenchArgs{64, 256, 8192});

  210.     args.emplace_back(BenchArgs{128, 864, 864});

  211.     args.emplace_back(BenchArgs{128, 9, 64});

  212. 

  213.     return args;

  214. }

  215. 

  216. #if CUDA_VERSION >= 10010

  217. std::vector<BenchArgs> get_f16_feat_model_args() {

  218.     std::vector<BenchArgs> args;

  219.     args.emplace_back(BenchArgs{128, 9216, 9216});

  220.     args.emplace_back(BenchArgs{128, 6400, 6400});

  221.     args.emplace_back(BenchArgs{128, 5184, 5184});

  222.     return args;

  223. }

  224. #endif

  225. 

  226. void benchmark_matrix_mul(

  227.         Handle* handle, const std::vector<BenchArgs>& args, DType A_dtype,

  228.         DType B_dtype, DType C_dtype, const char* algo = nullptr,

  229.         param::MatrixMul::Format format = param::MatrixMul::Format::DEFAULT) {

  230.     megdnn_assert(A_dtype.enumv() == B_dtype.enumv());

  231.     CUBenchmarker<MatrixMulForward> benchmarker(handle);

  232.     CUBenchmarker<MatrixMulForward> benchmarker_cublas(handle);

  233.     size_t RUNS = 1000;

  234.     benchmarker.set_display(false).set_times(RUNS);

  235.     benchmarker_cublas.set_display(false).set_times(RUNS);

  236.     benchmarker_cublas.set_before_exec_callback(

  237.             AlgoChecker<MatrixMulForward>("CUBLAS"));

  238.     benchmarker.set_dtype(0, A_dtype).set_dtype(1, B_dtype).set_dtype(2, C_dtype);

  239.     benchmarker_cublas.set_dtype(0, A_dtype)

  240.             .set_dtype(1, B_dtype)

  241.             .set_dtype(2, C_dtype);

  242.     using Param = MatrixMul::Param;

  243.     for (auto&& arg : args) {

  244.         size_t m = arg.m, n = arg.n, k = arg.k;

  245.         Param param;

  246.         param.transposeA = arg.mask & 0x1;

  247.         param.transposeB = arg.mask & 0x2;

  248.         param.format = format;

  249.         size_t A0 = m, A1 = k, B0 = k, B1 = n;

  250.         if (param.transposeA) {

  251.             std::swap(A0, A1);

  252.         }

  253.         if (param.transposeB) {

  254.             std::swap(B0, B1);

  255.         }

  256. 

  257.         benchmarker.set_param(param);

  258.         TensorShape A{A0, A1}, B{B0, B1}, C{m, n};

  259.         float time_in_ms = 0.f;

  260.         if (algo) {

  261.             time_in_ms = algo_benchmark<

  262.                                  MatrixMulForward, OprProxy<MatrixMulForward>, CUTimer>(

  263.                                  benchmarker, {A, B, C}, algo) /

  264.                          RUNS;

  265.         } else {

  266.             time_in_ms = benchmarker.execs({A, B, C}) / RUNS;

  267.         }

  268.         benchmarker_cublas.set_param(param);

  269.         auto time_in_ms_cublas = benchmarker_cublas.execs({A, B, C}) / RUNS;

  270.         float flo = 2.0 * m * n * k / (1e12);

  271.         printf("A=%s, B=%s, C=%s, time(algo=%s)=%.2f %.2fTops, "

  272.                "time(cublas)=%.2f %.2fTops, "

  273.                "perf(algo=%s)/perf(cublas)=%.2f\n",

  274.                A.to_string().c_str(), B.to_string().c_str(), C.to_string().c_str(),

  275.                algo, time_in_ms, (flo / (time_in_ms * 1e-3)), time_in_ms_cublas,

  276.                (flo / (time_in_ms_cublas * 1e-3)), algo,

  277.                time_in_ms_cublas / time_in_ms);

  278.     }

  279. }

  280. #endif

  281. }  // namespace

  282. 

  283. TEST_F(CUDA, CUTLASS_GEMM_MULTI_BATCHSIZE) {

  284.     auto args = matrix_mul::get_matmul_args_no_mask();

  285.     test_multibatchsize(

  286.             handle_cuda(), dtype::Float32(), dtype::Float32(), dtype::Float32(),

  287.             "CUTLASS_FLOAT32_SIMT_128X128X8_32X64X8", args,

  288.             param::MatrixMul::Format::DEFAULT);

  289. }

  290. 

  291. TEST_F(CUDA, CUTLASS_GEMM_SPLIT_K_MULTI_BATCHSIZE) {

  292.     auto args = matrix_mul::get_matmul_args_no_mask();

  293.     test_multibatchsize(

  294.             handle_cuda(), dtype::Float32(), dtype::Float32(), dtype::Float32(),

  295.             "CUTLASS_FLOAT32_SIMT_SPLIT_K_128X128X8_32X64X8", args,

  296.             param::MatrixMul::Format::DEFAULT,

  297.             [](const matrix_mul::TestArg& arg) { return arg.k <= arg.n; });

  298. }

  299. 

  300. TEST_F(CUDA, CUTLASS_GEMV_BATCHED_STRIDED_128_MULTI_BATCHSIZE) {

  301.     auto args = matrix_mul::get_matmul_args_no_mask();

  302.     test_multibatchsize(

  303.             handle_cuda(), dtype::Float32(), dtype::Float32(), dtype::Float32(),

  304.             "CUTLASS_FLOAT32_SIMT_GEMV_BATCHED_STRIDED_128", args,

  305.             param::MatrixMul::Format::DEFAULT);

  306. }

  307. 

  308. TEST_F(CUDA, CUTLASS_GEMV_BATCHED_STRIDED_64_MULTI_BATCHSIZE) {

  309.     auto args = matrix_mul::get_matmul_args_no_mask();

  310.     test_multibatchsize(

  311.             handle_cuda(), dtype::Float32(), dtype::Float32(), dtype::Float32(),

  312.             "CUTLASS_FLOAT32_SIMT_GEMV_BATCHED_STRIDED_64", args,

  313.             param::MatrixMul::Format::DEFAULT);

  314. }

  315. 

  316. TEST_F(CUDA, CUTLASS_GEMV_BATCHED_STRIDED_32_MULTI_BATCHSIZE) {

  317.     auto args = matrix_mul::get_matmul_args_no_mask();

  318.     test_multibatchsize(

  319.             handle_cuda(), dtype::Float32(), dtype::Float32(), dtype::Float32(),

  320.             "CUTLASS_FLOAT32_SIMT_GEMV_BATCHED_STRIDED_32", args,

  321.             param::MatrixMul::Format::DEFAULT);

  322. }

  323. 

  324. #define MEGDNN_FOREACH_CUTLASS_KERNEL(cb) \

  325.     cb(1, 64, 256, 8, 32, 64, 8);         \

  326.     cb(2, 256, 64, 8, 64, 32, 8);         \

  327.     cb(3, 32, 256, 8, 16, 64, 8);         \

  328.     cb(4, 256, 32, 8, 64, 16, 8);         \

  329.     cb(5, 128, 128, 8, 32, 64, 8);        \

  330.     cb(6, 128, 64, 8, 64, 32, 8);         \

  331.     cb(7, 64, 128, 8, 32, 64, 8);         \

  332.     cb(8, 128, 32, 8, 64, 32, 8);         \

  333.     cb(9, 32, 128, 8, 32, 64, 8);         \

  334.     cb(10, 64, 64, 8, 32, 64, 8);         \

  335.     cb(11, 32, 64, 8, 32, 64, 8);         \

  336.     cb(12, 64, 32, 8, 64, 32, 8);         \

  337.     cb(13, 32, 32, 8, 32, 32, 8);         \

  338.     cb(14, 8, 32, 8, 8, 32, 8);           \

  339.     cb(15, 16, 32, 8, 16, 32, 8);         \

  340.     cb(16, 16, 64, 8, 16, 64, 8);         \

  341.     cb(17, 16, 128, 8, 16, 64, 8);

  342. 

  343. #define cb(name, tbm, tbn, tbk, wm, wn, wk)                                          \

  344.     TEST_F(CUDA, CUTLASS_GEMM_##name) {                                              \

  345.         matrix_mul::check_matrix_mul<MatrixMulForward>(                              \

  346.                 dtype::Float32(), dtype::Float32(), dtype::Float32(), handle_cuda(), \

  347.                 "CUTLASS_FLOAT32_SIMT_" #tbm "X" #tbn "X" #tbk "_" #wm "X" #wn       \

  348.                 "X" #wk);                                                            \

  349.     }

  350. 

  351. MEGDNN_FOREACH_CUTLASS_KERNEL(cb)

  352. 

  353. #undef cb

  354. 

  355. #define cb(name, tbm, tbn, tbk, wm, wn, wk)                                            \

  356.     TEST_F(CUDA, CUTLASS_GEMM_SPLIT_K_##name) {                                        \

  357.         matrix_mul::check_matrix_mul<MatrixMulForward>(                                \

  358.                 dtype::Float32(), dtype::Float32(), dtype::Float32(), handle_cuda(),   \

  359.                 "CUTLASS_FLOAT32_SIMT_SPLIT_K_" #tbm "X" #tbn "X" #tbk "_" #wm "X" #wn \

  360.                 "X" #wk,                                                               \

  361.                 param::MatrixMul::Format::DEFAULT, 8, 1e-3,                            \

  362.                 matrix_mul::get_matmul_args_split_k());                                \

  363.     }

  364. 

  365. MEGDNN_FOREACH_CUTLASS_KERNEL(cb)

  366. 

  367. #undef cb

  368. #undef MEGDNN_FOREACH_CUTLASS_KERNEL

  369. 

  370. #if CUDA_VERSION >= 10010

  371. #define MEGDNN_FOREACH_CUTLASS_KERNEL(cb)     \

  372.     cb(1, 256, 128, 32, 64, 64, 32, 8, 8, 4); \

  373.     cb(2, 128, 256, 32, 64, 64, 32, 8, 8, 4); \

  374.     cb(3, 128, 128, 32, 64, 64, 32, 8, 8, 4);

  375. 

  376. #define cb(name, tbm, tbn, tbk, wm, wn, wk, im, in, ik)                              \

  377.     TEST_F(CUDA, CUTLASS_F16_884_GEMM_##name) {                                      \

  378.         require_compute_capability(7, 0);                                            \

  379.         matrix_mul::check_matrix_mul<MatrixMulForward>(                              \

  380.                 dtype::Float16(), dtype::Float16(), dtype::Float16(), handle_cuda(), \

  381.                 "CUTLASS_FLOAT16_TENSOR_OP_h" #im #in #ik "_" #tbm "X" #tbn "X" #tbk \

  382.                 "_" #wm "X" #wn "X" #wk,                                             \

  383.                 param::MatrixMul::Format::DEFAULT, 8, 1e-2,                          \

  384.                 matrix_mul::get_matmul_args());                                      \

  385.     }

  386. MEGDNN_FOREACH_CUTLASS_KERNEL(cb)

  387. 

  388. #undef cb

  389. 

  390. #define cb(name, tbm, tbn, tbk, wm, wn, wk, im, in, ik)                              \

  391.     TEST_F(CUDA, CUTLASS_F16_884_GEMM_SPLIT_K_##name) {                              \

  392.         require_compute_capability(7, 0);                                            \

  393.         matrix_mul::check_matrix_mul<MatrixMulForward>(                              \

  394.                 dtype::Float16(), dtype::Float16(), dtype::Float16(), handle_cuda(), \

  395.                 "CUTLASS_FLOAT16_TENSOR_OP_SPLIT_K_h" #im #in #ik "_" #tbm "X" #tbn  \

  396.                 "X" #tbk "_" #wm "X" #wn "X" #wk,                                    \

  397.                 param::MatrixMul::Format::DEFAULT, 8, 1e-3,                          \

  398.                 matrix_mul::get_matmul_args_split_k(), true,                         \

  399.                 param::MatrixMul::ComputeMode::FLOAT32);                             \

  400.     }

  401. MEGDNN_FOREACH_CUTLASS_KERNEL(cb)

  402. 

  403. #undef cb

  404. 

  405. #undef MEGDNN_FOREACH_CUTLASS_KERNEL

  406. #endif

  407. 

  408. #if CUDA_VERSION >= 10020

  409. #define MEGDNN_FOREACH_CUTLASS_KERNEL(cb)      \

  410.     cb(1, 256, 128, 32, 64, 64, 32, 16, 8, 8); \

  411.     cb(2, 128, 256, 32, 64, 64, 32, 16, 8, 8); \

  412.     cb(3, 128, 128, 32, 64, 64, 32, 16, 8, 8);

  413. 

  414. #define cb(name, tbm, tbn, tbk, wm, wn, wk, im, in, ik)                              \

  415.     TEST_F(CUDA, CUTLASS_F16_1688_GEMM_##name) {                                     \

  416.         require_compute_capability(7, 5);                                            \

  417.         matrix_mul::check_matrix_mul<MatrixMulForward>(                              \

  418.                 dtype::Float16(), dtype::Float16(), dtype::Float16(), handle_cuda(), \

  419.                 "CUTLASS_FLOAT16_TENSOR_OP_h" #im #in #ik "_" #tbm "X" #tbn "X" #tbk \

  420.                 "_" #wm "X" #wn "X" #wk,                                             \

  421.                 param::MatrixMul::Format::DEFAULT, 8, 1e-2,                          \

  422.                 matrix_mul::get_matmul_args(), true,                                 \

  423.                 param::MatrixMul::ComputeMode::FLOAT32);                             \

  424.     }

  425. MEGDNN_FOREACH_CUTLASS_KERNEL(cb)

  426. 

  427. #undef cb

  428. 

  429. #define cb(name, tbm, tbn, tbk, wm, wn, wk, im, in, ik)                              \

  430.     TEST_F(CUDA, CUTLASS_F16_1688_GEMM_SPLIT_K_##name) {                             \

  431.         require_compute_capability(7, 5);                                            \

  432.         matrix_mul::check_matrix_mul<MatrixMulForward>(                              \

  433.                 dtype::Float16(), dtype::Float16(), dtype::Float16(), handle_cuda(), \

  434.                 "CUTLASS_FLOAT16_TENSOR_OP_SPLIT_K_h" #im #in #ik "_" #tbm "X" #tbn  \

  435.                 "X" #tbk "_" #wm "X" #wn "X" #wk,                                    \

  436.                 param::MatrixMul::Format::DEFAULT, 8, 1e-3,                          \

  437.                 matrix_mul::get_matmul_args_split_k());                              \

  438.     }

  439. MEGDNN_FOREACH_CUTLASS_KERNEL(cb)

  440. 

  441. #undef cb

  442. 

  443. #undef MEGDNN_FOREACH_CUTLASS_KERNEL

  444. #endif

  445. 

  446. #if MEGDNN_WITH_BENCHMARK

  447. TEST_F(CUDA, BENCHMARK_CUTLASS_MATMUL) {

  448.     benchmark_matrix_mul(

  449.             handle_cuda(), get_square_matmul_args(), dtype::Float32(), dtype::Float32(),

  450.             dtype::Float32(), "CUTLASS_FLOAT32_SIMT");

  451. }

  452. 

  453. TEST_F(CUDA, BENCHMARK_CUTLASS_MATMUL_FEAT) {

  454.     benchmark_matrix_mul(

  455.             handle_cuda(), get_feat_model_args(), dtype::Float32(), dtype::Float32(),

  456.             dtype::Float32(), "CUTLASS_FLOAT32_SIMT");

  457. }

  458. 

  459. #if CUDA_VERSION >= 10010

  460. TEST_F(CUDA, BENCHMARK_CUTLASS_F16_MATMUL_FEAT) {

  461.     benchmark_matrix_mul(

  462.             handle_cuda(), get_f16_feat_model_args(), dtype::Float16(),

  463.             dtype::Float16(), dtype::Float16(), "CUTLASS_FLOAT16_TENSOR_OP");

  464. }

  465. #endif

  466. #endif

  467. }  // namespace test

  468. }  // namespace megdnn

  469. #endif

  470. 

  471. // vim: syntax=cpp.doxygen