megenginelite-sys 1.8.2

A safe megenginelite wrapper in Rust
Documentation 
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
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

/**
 * \file dnn/src/cuda/batched_matrix_mul/int8x8x32.cuh
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */
#pragma once
#include "src/cuda/utils.cuh"

namespace megdnn {
namespace cuda {

template <typename ThreadConfig_, int m_, int n_, int k_tot, int k_>
struct UnrollConfig {
    typedef ThreadConfig_ ThreadConfig;
    static int const unroll_m = m_;
    static int const unroll_n = n_;
    static int const block_m = ThreadConfig::thread_x * m_;
    static int const block_n = ThreadConfig::thread_y * n_;
    static int const unroll_k = k_tot;
    static int const unroll = k_;
    static int const thread_k = k_tot / k_;
    static int const load_m = (m_ / 4) / (ThreadConfig::thread_y / thread_k) * 4;
    static int const load_n = (n_ / 4) / (ThreadConfig::thread_x / thread_k) * 4;
};

template <int x_, int y_>
struct ThreadConfig {
    static int const thread_x = x_;
    static int const thread_y = y_;
};

template <int row, int col>
struct SmemConfig {
    static int const smem_row = row;
    static int const smem_col = col;
};

template <typename SmemConfig_>
struct Global2SharedMem {
    typedef SmemConfig_ SmemConfig;
    const int8_t* g_ptr;
    int32_t* smem;
    int smem_off;
    int smem_bound;
    int32_t reg[SmemConfig::smem_col][SmemConfig::smem_row / 4];
    int ld_src;
    int ld_dst;
    int check_bound_row;
    int check_bound_col;
    int step;
    bool tr;
    bool aligned;

    __device__ __forceinline__ Global2SharedMem(
            int32_t* smem_, int s_off, int s_bound, int ld_src_, int ld_dst_, int b_r_,
            int b_c_, int step_, bool tr_, bool al_)
            : smem(smem_),
              smem_off(s_off),
              smem_bound(s_bound),
              ld_src(ld_src_),
              ld_dst(ld_dst_),
              check_bound_row(b_r_),
              check_bound_col(b_c_),
              step(step_),
              tr(tr_),
              aligned(al_) {}

    __device__ __forceinline__ void gmem2reg_cpy();
    __device__ __forceinline__ void reg2smem_cpy();
    __device__ __forceinline__ void iter_forward();
};

template <typename UnrollConfig, typename ThreadConfig>
__global__ void batched_8x8x32_kern(
        const int8_t* a, int lda, int sta, bool tra, const int8_t* b, int ldb, int stb,
        bool trb, int32_t* c, int ldc, int stc, int m, int n, int k);

void exec_igemm_8x8x32(
        const int8_t* A, const int8_t* B, int32_t* C, const int batch_count,
        const int m, const int n, const int k, int ldA, int ldB, int ldC, int stA,
        int stB, int stC, bool transA, bool transB, cudaStream_t stream);

}  // namespace cuda
}  // namespace megdnn

// vim: syntax=cuda.doxygen