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
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154

/**
 * \file dnn/src/cuda/batched_matrix_mul/cublas_lt.cpp
 * 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.
 */
#include "./algo.h"
#include "src/cuda/handle.h"
#include "src/cuda/matrix_mul/cublasLt_wrapper.h"
#include "src/cuda/utils.h"

using namespace megdnn;
using namespace cuda;

#if CUDA_VERSION >= 10010
static inline CUBLASLTMatmulDesc::SizeArgs from_local_size_args(
        const BatchedMatrixMulForwardImpl::AlgoBase::SizeArgs& args) {
    auto&& param = args.opr->param();
    auto&& handle = concrete_handle(args.opr->handle());
    bool transA = param.transposeA;
    bool transB = param.transposeB;
    return {handle, transA, transB, args.layout_a, args.layout_b, args.layout_c};
}

bool BatchedMatrixMulForwardImpl::AlgoCublasLt::is_available(
        const SizeArgs& args) const {
    auto cublasLt_args = from_local_size_args(args);
    auto&& dev_prop = current_device_prop();
    bool is_dev_support = dev_prop.major >= 7;
    bool res = is_dev_support && CUBLASLTMatmulDesc(cublasLt_args, true)
                                         .is_available(cublasLt_args, INT_MAX);
    return res;
}

size_t BatchedMatrixMulForwardImpl::AlgoCublasLt::get_workspace_in_bytes(
        const SizeArgs& args) const {
    auto cublasLt_args = from_local_size_args(args);
    cublasLtMatmulAlgo_t algo;
    CUBLASLTMatmulDesc desc(cublasLt_args, true);
    desc.get_algorithm_heuristic(cublasLt_args, INT_MAX, algo);
    return desc.get_workspace_bundle(cublasLt_args, algo).total_size_in_bytes();
}

void BatchedMatrixMulForwardImpl::AlgoCublasLt::exec(const ExecArgs& args) const {
    auto cublasLt_args = from_local_size_args(args);
    cublasLtMatmulAlgo_t algo;
    CUBLASLTMatmulDesc desc(cublasLt_args, true);
    desc.get_algorithm_heuristic(cublasLt_args, INT_MAX, algo);
    auto ws_bundle = desc.get_workspace_bundle(cublasLt_args, algo);
    auto&& handle = concrete_handle(args.opr->handle());
    auto&& stream = handle->stream();
    auto&& cublasLt_handle = handle->cublasLt_handle();
    auto batched_hgemm = [&]() {
        auto zero_half = handle->zero_device_h();
        auto one_half = handle->one_device_h();
        megdnn_assert(
                ws_bundle.nr_workspace() == 1,
                "workspace bundle size should be 1(ws_algo)");
        cublas_check(cublasLtMatmul(
                cublasLt_handle, desc.matmul_desc, one_half,
                static_cast<const __half*>(args.tensor_b.raw_ptr()), desc.layout_b,
                static_cast<const __half*>(args.tensor_a.raw_ptr()), desc.layout_a,
                zero_half, static_cast<const __half*>(args.tensor_c.raw_ptr()),
                desc.layout_c, static_cast<__half*>(args.tensor_c.raw_ptr()),
                desc.layout_c, &algo, ws_bundle.get(0), ws_bundle.get_size(0), stream));
    };
    auto batched_sgemm = [&]() {
        auto zero = handle->zero_device();
        auto one = handle->one_device();
        auto dev_b = (desc.dt_b == CUDA_R_16F)
                           ? static_cast<void*>(args.tensor_b.ptr<dt_float16>())
                           : static_cast<void*>(args.tensor_b.ptr<dt_float32>());
        auto dev_a = (desc.dt_a == CUDA_R_16F)
                           ? static_cast<void*>(args.tensor_a.ptr<dt_float16>())
                           : static_cast<void*>(args.tensor_a.ptr<dt_float32>());
        auto dev_c = static_cast<void*>(args.tensor_c.raw_ptr());
        megdnn_assert(
                ws_bundle.nr_workspace() == 1,
                "workspace bundle size should be 1(ws_algo)");
        cublas_check(cublasLtMatmul(
                cublasLt_handle, desc.matmul_desc, one, dev_b, desc.layout_b, dev_a,
                desc.layout_a, zero, dev_c, desc.layout_c, dev_c, desc.layout_c, &algo,
                ws_bundle.get(0), ws_bundle.get_size(0), stream));
    };

    auto batched_igemm = [&]() {
        auto zero = handle->zero_device();
        auto one = handle->one_device();
        megdnn_assert(
                ws_bundle.nr_workspace() == 4,
                "workspace bundle size should be 4(ws_algo, ws_a, ws_b, ws_c)");
        void* ws_b = ws_bundle.get(1);
        void* ws_a = ws_bundle.get(2);
        void* ws_c = ws_bundle.get(3);
        int32_t pm = CUBLAS_POINTER_MODE_DEVICE;
        cublasOperation_t trans_a = CUBLAS_OP_T, trans_c = CUBLAS_OP_N;
        cublasLtMatrixTransformDesc_t transform_desc = nullptr;
        cublas_check(cublasLtMatrixTransformDescCreate(&transform_desc, CUDA_R_32F));
        cublas_check(cublasLtMatrixTransformDescSetAttribute(
                transform_desc, CUBLASLT_MATRIX_TRANSFORM_DESC_POINTER_MODE, &pm,
                sizeof(pm)));
        cublas_check(cublasLtMatrixTransform(
                cublasLt_handle, transform_desc, one, args.tensor_b.raw_ptr(),
                desc.layout_b, zero, nullptr, nullptr, ws_b, desc.layout_trans_b,
                stream));
        cublas_check(cublasLtMatrixTransformDescSetAttribute(
                transform_desc, CUBLASLT_MATRIX_TRANSFORM_DESC_TRANSA, &trans_a,
                sizeof(trans_a)));
        cublas_check(cublasLtMatrixTransform(
                cublasLt_handle, transform_desc, one, args.tensor_a.raw_ptr(),
                desc.layout_a, zero, nullptr, nullptr, ws_a, desc.layout_trans_a,
                stream));
        cublas_check(cublasLtMatmul(
                cublasLt_handle, desc.matmul_desc, one, ws_b, desc.layout_trans_b, ws_a,
                desc.layout_trans_a, zero, ws_c, desc.layout_trans_c, ws_c,
                desc.layout_trans_c, &algo, ws_bundle.get(0), ws_bundle.get_size(0),
                stream));
        cublas_check(cublasLtMatrixTransformDescSetAttribute(
                transform_desc, CUBLASLT_MATRIX_TRANSFORM_DESC_TRANSA, &trans_c,
                sizeof(trans_c)));
        cublas_check(cublasLtMatrixTransform(
                cublasLt_handle, transform_desc, one, ws_c, desc.layout_trans_c, zero,
                nullptr, nullptr, args.tensor_c.raw_ptr(), desc.layout_c, stream));
        cublas_check(cublasLtMatrixTransformDescDestroy(transform_desc));
    };

    ws_bundle.set(args.workspace.raw_ptr);
#if CUDA_VERSION >= 11000
    if (desc.dt_compute == CUBLAS_COMPUTE_32I) {
        batched_igemm();
    } else if (desc.dt_compute == CUBLAS_COMPUTE_16F) {
        batched_hgemm();
    } else if (desc.dt_compute == CUBLAS_COMPUTE_32F) {
        batched_sgemm();
    } else {
        megdnn_throw("compute_type must be int32/float16/float32");
    }
#else
    if (desc.dt_compute == CUDA_R_32I) {
        batched_igemm();
    } else if (desc.dt_compute == CUDA_R_16F) {
        batched_hgemm();
    } else if (desc.dt_compute == CUDA_R_32F) {
        batched_sgemm();
    } else {
        megdnn_throw("compute_type must be int32/float16/float32");
    }
#endif
}
#endif