blas.h 3.17 KB
Newer Older
PanZezhongQY's avatar
PanZezhongQY committed
1
2
3
4
5
6
#ifndef __BLAS_H__
#define __BLAS_H__

#include "infiniop/operator.h"
#include <algorithm>

7
8
namespace matmul {
struct BlasMatrix {
PanZezhongQY's avatar
PanZezhongQY committed
9
10
    size_t ndim;
    size_t batch;
11
    ptrdiff_t stride;
PanZezhongQY's avatar
PanZezhongQY committed
12
13
    size_t rows;
    size_t cols;
14
15
    ptrdiff_t row_stride;
    ptrdiff_t col_stride;
PanZezhongQY's avatar
PanZezhongQY committed
16

17
    BlasMatrix() = default;
PanZezhongQY's avatar
PanZezhongQY committed
18

PanZezhong's avatar
PanZezhong committed
19
    BlasMatrix(infiniopTensorDescriptor_t layout, infiniStatus_t *status) {
PanZezhongQY's avatar
PanZezhongQY committed
20
        if (layout->ndim == 2) {
21
22
23
24
25
26
27
            ndim = 2;
            batch = 1;
            stride = 0;
            rows = layout->shape[0];
            cols = layout->shape[1];
            row_stride = layout->strides[0];
            col_stride = layout->strides[1];
PanZezhongQY's avatar
PanZezhongQY committed
28
        } else if (layout->ndim == 3) {
29
30
31
32
33
34
35
            ndim = 3;
            batch = layout->shape[0];
            stride = batch == 1 ? 0 : layout->strides[0];
            rows = layout->shape[1];
            cols = layout->shape[2];
            row_stride = layout->strides[1];
            col_stride = layout->strides[2];
PanZezhongQY's avatar
PanZezhongQY committed
36
        } else {
PanZezhong's avatar
PanZezhong committed
37
            *status = INFINI_STATUS_BAD_TENSOR_SHAPE;
PanZezhongQY's avatar
PanZezhongQY committed
38
39
40
            return;
        }

41
        if (row_stride != 1 && col_stride != 1) {
PanZezhong's avatar
PanZezhong committed
42
            *status = INFINI_STATUS_BAD_TENSOR_STRIDES;
PanZezhongQY's avatar
PanZezhongQY committed
43
44
45
            return;
        }

PanZezhong's avatar
PanZezhong committed
46
        *status = INFINI_STATUS_SUCCESS;
PanZezhongQY's avatar
PanZezhongQY committed
47
48
    }

49
    bool match_batch(size_t _batch) const {
50
        return batch == _batch || batch == 1;
PanZezhongQY's avatar
PanZezhongQY committed
51
52
53
54
55
56
57
    }

    void transpose() {
        std::swap(rows, cols);
        std::swap(row_stride, col_stride);
    }

58
    ptrdiff_t ld() const {
59
        return row_stride == 1 ? col_stride : row_stride;
PanZezhongQY's avatar
PanZezhongQY committed
60
    }
61
62
};

63
enum class MatrixLayout : char {
64
65
66
    COL_MAJOR,
    ROW_MAJOR,
};
PanZezhongQY's avatar
PanZezhongQY committed
67
68
69
70
71
72
73
74
75
76

struct MatmulInfo {
    BlasMatrix a_matrix;
    BlasMatrix b_matrix;
    BlasMatrix c_matrix;

    size_t m, n, k, batch;

    bool is_transed = false;

77
78
79
    MatmulInfo(infiniopTensorDescriptor_t c_desc,
               infiniopTensorDescriptor_t a_desc,
               infiniopTensorDescriptor_t b_desc,
PanZezhong's avatar
PanZezhong committed
80
               infiniStatus_t *status,
81
               MatrixLayout layout) {
PanZezhongQY's avatar
PanZezhongQY committed
82
        a_matrix = BlasMatrix(a_desc, status);
PanZezhong's avatar
PanZezhong committed
83
        if (*status != INFINI_STATUS_SUCCESS) {
PanZezhongQY's avatar
PanZezhongQY committed
84
85
86
            return;
        }
        b_matrix = BlasMatrix(b_desc, status);
PanZezhong's avatar
PanZezhong committed
87
        if (*status != INFINI_STATUS_SUCCESS) {
PanZezhongQY's avatar
PanZezhongQY committed
88
89
90
            return;
        }
        c_matrix = BlasMatrix(c_desc, status);
PanZezhong's avatar
PanZezhong committed
91
        if (*status != INFINI_STATUS_SUCCESS) {
PanZezhongQY's avatar
PanZezhongQY committed
92
93
94
            return;
        }

95
        if (c_matrix.rows != a_matrix.rows || c_matrix.cols != b_matrix.cols || a_matrix.cols != b_matrix.rows) {
PanZezhong's avatar
PanZezhong committed
96
            *status = INFINI_STATUS_BAD_TENSOR_SHAPE;
PanZezhongQY's avatar
PanZezhongQY committed
97
98
99
100
101
            return;
        }

        batch = c_matrix.batch;
        if (!a_matrix.match_batch(batch) || !b_matrix.match_batch(batch)) {
PanZezhong's avatar
PanZezhong committed
102
            *status = INFINI_STATUS_BAD_TENSOR_SHAPE;
PanZezhongQY's avatar
PanZezhongQY committed
103
104
105
            return;
        }

106
107
        if ((layout == MatrixLayout::COL_MAJOR && c_matrix.col_stride == 1)
            || (layout == MatrixLayout::ROW_MAJOR && c_matrix.row_stride == 1)) {
PanZezhongQY's avatar
PanZezhongQY committed
108
109
110
111
112
113
114
115
116
117
118
119
            c_matrix.transpose();
            b_matrix.transpose();
            a_matrix.transpose();
            std::swap(a_matrix, b_matrix);
            is_transed = true;
        }

        m = c_matrix.rows;
        n = c_matrix.cols;
        k = a_matrix.cols;
    }
};
120
} // namespace matmul
PanZezhongQY's avatar
PanZezhongQY committed
121

122
#endif // __BLAS_H__