flexflow-train/lib/kernels/include/kernels/batch_matmul_kernels.h at 29144947fc5aff6c68863fc5b610c3e2ce29a54d · flexflow/flexflow-train · GitHub

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
#ifndef _FLEXFLOW_OPS_KERNELS_BATCH_MATMUL_KERNELS_H
#define _FLEXFLOW_OPS_KERNELS_BATCH_MATMUL_KERNELS_H

#include "device.h"
#include "kernels/allocation.h"
#include "kernels/ff_handle.h"

namespace FlexFlow::Kernels::BatchMatmul {

void forward_kernel(ffStream_t stream,
                    PerDeviceFFHandle const &handle,
                    float *output_ptr,
                    float const *a_input_ptr,
                    float const *b_input_ptr,
                    int m,
                    int n,
                    int k,
                    int batch,
                    int seq_length,
                    int a_seq_length_dim,
                    int b_seq_length_dim);

void backward_kernel(ffStream_t stream,
                     PerDeviceFFHandle const &handle,
                     float const *o_ptr,
                     float const *o_grad_ptr,
                     float const *a_ptr,
                     float *a_grad_ptr,
                     float const *b_ptr,
                     float *b_grad_ptr,
                     int m,
                     int n,
                     int k,
                     int batch);

} // namespace FlexFlow::Kernels::BatchMatmul

#endif