bmm_v8.cpp

#include <iostream>
#include <string.h>
#include <vector>

#include "kernel.h"

using namespace std;

#define min(x, y) ((x) < (y) ? (x) : (y))

#ifndef NTHREADS
    #define NTHREADS 6 // TO DO: dynamically
#endif

#define MC (16 * NTHREADS * 8)
#define NC (6 * NTHREADS * 70)
#define KC 792
#define MC_D (8 * NTHREADS * 8)
#define NC_D (6 * NTHREADS * 35)

static float blockA_packed[MC * KC] __attribute__((aligned(64)));
static float blockB_packed[NC * KC] __attribute__((aligned(64)));
static double blockA_packed_d[MC_D * KC] __attribute__((aligned(64)));
static double blockB_packed_d[NC_D * KC] __attribute__((aligned(64)));


void pack_panelB_d(double* B, double* blockB_packed, int nr, int kc, int jj, int pp, const int* offset_right_b, const int* offset_contract_b) {
    for (int p = 0; p < kc; p++) {
        for (int j = 0; j < nr; j++) {
            *blockB_packed++ = B[offset_right_b[jj+j] + offset_contract_b[pp+p]];
        }
        for (int j = nr; j < 6; j++) {
            *blockB_packed++ = 0;
        }
    }
}

void pack_blockB_d(double* B, double* blockB_packed, int nc, int kc, int p, int jj, const int* offset_right_b, const int* offset_contract_b) {
#pragma omp parallel for num_threads(NTHREADS)
    for (int j = 0; j < nc; j += 6) {
        int nr = min(6, nc - j);
        pack_panelB_d(B, &blockB_packed[j * kc], nr, kc, jj+j, p, offset_right_b, offset_contract_b );
    }
}

void pack_panelA_d(double* A, double* blockA_packed, int mr, int kc, int pp, int ii, const int* offset_left_a, const int* offset_contract_a) {
    for (int p = 0; p < kc; p++) {
        for (int i = 0; i < mr; i++) {
            *blockA_packed++ = A[offset_left_a[ii + i]+ offset_contract_a[pp+p]];
        }
        for (int i = mr; i < 8; i++) {
            *blockA_packed++ = 0;
        }
    }
}

void pack_blockA_d(double* A, double* blockA_packed, int mc, int kc, int p, int ii, const int* offset_left_a, const int* offset_contract_a) {
#pragma omp parallel for num_threads(NTHREADS)
    for (int i = 0; i < mc; i += 8) {
        int mr = min(8, mc - i);
        pack_panelA_d(A, &blockA_packed[i * kc], mr, kc, p, ii+i, offset_left_a, offset_contract_a);
    }
}

void bmm_double(double* A, double* B, double* C, const int b, const int m, const int n, const int k, \
const int* offset_batch_a, const int* offset_left_a, const int* offset_contract_a, \
const int* offset_batch_b, const int* offset_contract_b, const int* offset_right_b) {

    const int bsC = m * n;        // batch size for C

    if (n <= NC_D){
        for (int batch = 0; batch < b; batch++) {
            for (int j = 0; j < n; j += NC_D) {
                const int nc = min(NC_D, n - j);
                for (int p = 0; p < k; p += KC) {
                    const int kc = min(KC, k - p);
                    pack_blockB_d(&B[offset_batch_b[batch]], blockB_packed_d, nc, kc, p, j, offset_right_b, offset_contract_b);
                    for (int i = 0; i < m; i += MC_D) {
                        const int mc = min(MC_D, m - i);
                        pack_blockA_d(&A[offset_batch_a[batch]], blockA_packed_d, mc, kc, p, i, offset_left_a, offset_contract_a);
    #pragma omp parallel for collapse(2) num_threads(NTHREADS)
                        for (int jr = 0; jr < nc; jr += 6) {
                            for (int ir = 0; ir < mc; ir += 8) {
                                const int nr = min(6, nc - jr);
                                const int mr = min(8, mc - ir);
                                kernel_8x6_double(&blockA_packed_d[ir * kc],
                                            &blockB_packed_d[jr * kc],
                                            &C[batch * bsC + (j + jr) * m + (i + ir)],
                                            mr,
                                            nr,
                                            kc,
                                            m);
                            }
                        }
                    }
                }
            }
        }
    }
    else{
        double* A_ordered = static_cast<double*>(_aligned_malloc(m * k * sizeof(double), 64));
        if (!A_ordered) {
            throw std::bad_alloc();
        }
        double* curr_block = A_ordered;

        for (int batch = 0; batch < b; batch++) {

            // first block in B (also C): here we need pack A once
            const int nc = NC_D;
            int j = 0;
            for (int p = 0; p < k; p += KC) {
                const int kc = min(KC, k - p);
                pack_blockB_d(&B[offset_batch_b[batch]], blockB_packed_d, nc, kc, p, j, offset_right_b, offset_contract_b);
                for (int i = 0; i < m; i += MC_D) {
                    const int mc = min(MC_D, m - i);
                    pack_blockA_d(&A[offset_batch_a[batch]], blockA_packed_d, mc, kc, p, i, offset_left_a, offset_contract_a);
                    memcpy(curr_block, blockA_packed_d, mc * kc * sizeof(double));
                    curr_block += mc*kc;
                    #pragma omp parallel for collapse(2) num_threads(NTHREADS)
                    for (int jr = 0; jr < nc; jr += 6) {
                        for (int ir = 0; ir < mc; ir += 8) {
                            const int nr = min(6, nc - jr);
                            const int mr = min(8, mc - ir);
                            kernel_8x6_double(&blockA_packed_d[ir * kc],
                                        &blockB_packed_d[jr * kc],
                                        &C[batch * bsC + (j + jr) * m + (i + ir)],
                                        mr,
                                        nr,
                                        kc,
                                        m);
                        }
                    }
                }
            }

            // all remaining blocks of B (also C) are now calculated with A_ordered
            for (int j = NC_D; j < n; j += NC_D) {
                double* curr_block = A_ordered;
                const int nc = min(NC_D, n - j);
                for (int p = 0; p < k; p += KC) {
                    const int kc = min(KC, k - p);
                    pack_blockB_d(&B[offset_batch_b[batch]], blockB_packed_d, nc, kc, p, j, offset_right_b, offset_contract_b);
                    for (int i = 0; i < m; i += MC_D) {
                        const int mc = min(MC_D, m - i);
                        memcpy(blockA_packed_d, curr_block, mc * kc * sizeof(double));
                        curr_block += mc*kc;
                        #pragma omp parallel for collapse(2) num_threads(NTHREADS)
                        for (int jr = 0; jr < nc; jr += 6) {
                            for (int ir = 0; ir < mc; ir += 8) {
                                const int nr = min(6, nc - jr);
                                const int mr = min(8, mc - ir);
                                kernel_8x6_double(&blockA_packed_d[ir * kc],
                                            &blockB_packed_d[jr * kc],
                                            &C[batch * bsC + (j + jr) * m + (i + ir)],
                                            mr,
                                            nr,
                                            kc,
                                            m);
                            }
                        }
                    }
                }
            }
        }
        _aligned_free(A_ordered);
    }
}


void fill(int* offsets, const int* dims, const int* strides, const int* sizes_block,
          int start, int level, int value, const int n) {
    if (level == n - 1) {
        for (int i = 0; i < dims[level]; ++i) {
            offsets[start] = i * strides[level] + value;
            start += 1;
        }
    } else {
        for (int i = 0; i < dims[level]; ++i) {
            fill(offsets, dims, strides, sizes_block, \
                start + i * sizes_block[level], level + 1, value + i * strides[level], n);
        }
    }
}

extern "C" {
    void my_bmm(
    void* A,       // Pointer to A with dims (b, m, k)
    void* B,       // Pointer to B with dims (b, k, n)
    void* C,       // Pointer to C with dims (b, m, n)
    int b,
    int m,
    int n,
    int k,
    const int* offset_batch_a, const int* offset_left_a, const int* offset_contract_a,
    const int* offset_batch_b, const int* offset_contract_b, const int* offset_right_b,
    const char* d_type
) {
        if(*d_type == *"d"){
            // Trick: bmm_double expects colum major layout. compute B.T @ A.T
            bmm_double(static_cast<double*>(B), static_cast<double*>(A), static_cast<double*>(C), b, n, m, k, \
            offset_batch_b, offset_right_b, offset_contract_b, offset_batch_a, offset_contract_a, offset_left_a);
        }
        // if(*d_type == *"f") {
        //     // Trick: bmm_float expects colum major layout. compute B.T @ A.T
        //     bmm_float(static_cast<float*>(B), static_cast<float*>(A), static_cast<float*>(C), b, n, m, k);
        // }
    }
    void my_fill(void* offsets, void* dims, void* strides, void* sizes_block, const int n){
        fill(static_cast<int*>(offsets), static_cast<int*>(dims), static_cast<int*>(strides), \
            static_cast<int*>(sizes_block), 0, 0, 0, n );
    }
}