gcn.cpp

#include "hip/hip_runtime.h"
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <chrono>
#include <fstream>
#include <vector>
using namespace std;
typedef std::chrono::time_point<std::chrono::steady_clock> TimePoint;

int v_num = 0;
int e_num = 0;
int F0 = 0, F1 = 0;
// coo graph
vector<vector<int>> edge_index;
vector<vector<double>> edge_val;
vector<int> degree;
vector<int> raw_graph;

// csr graph;
int *nodes_index;
int *edges;
double *edges_value;

// layer
double *X0, *W1, *X1, *X1_inter;
// layer on gpu
double *d_X0, *d_W1, *d_X1, *d_X1_inter;

// csr graph on gpu
int *d_index, *d_edges;
double *d_edges_val;

void readGraph(char *fname)
{
    ifstream infile(fname);
    int source;
    int end;
    infile >> v_num >> e_num;
    while (!infile.eof())
    {
        infile >> source >> end;
        if (infile.peek() == EOF)
            break;
        raw_graph.push_back(source);
        raw_graph.push_back(end);
    }
}

void to_csr()
{

    nodes_index = (int *)malloc(v_num * sizeof(int) + 1);

    int sum = 0;
    for (int i = 0; i < v_num; i++)
    {
        nodes_index[i] = sum;
        sum += degree[i];
    }
    nodes_index[v_num] = sum;

    edges = (int *)malloc(e_num * sizeof(int));
    for (int i = 0; i < v_num; i++)
    {
        memcpy(edges + nodes_index[i], edge_index[i].data(), sizeof(int) * edge_index[i].size());
    }

    edges_value = (double *)malloc(e_num * sizeof(double));
    for (int i = 0; i < v_num; i++)
    {
        memcpy(edges_value + nodes_index[i], edge_val[i].data(), sizeof(double) * edge_val[i].size());
    }
}

void raw_graph_to_AdjacencyList()
{
    int src;
    int dst;
    edge_index.resize(v_num);
    edge_val.resize(v_num);
    degree.resize(v_num, 0);

    for (int i = 0; i < raw_graph.size() / 2; i++)
    {
        src = raw_graph[2 * i];
        dst = raw_graph[2 * i + 1];
        edge_index[dst].push_back(src);
        degree[src]++;
    }
}

void edgeNormalization()
{
    for (int i = 0; i < v_num; i++)
    {
        for (int j = 0; j < edge_index[i].size(); j++)
        {
            double val = 1 / sqrt(degree[i]) / sqrt(degree[edge_index[i][j]]);
            edge_val[i].push_back(val);
        }
    }
}

void readdouble(char *fname, double *&dst, int num)
{
    dst = (double *)malloc(num * sizeof(double));
    FILE *fp = fopen(fname, "rb");
    fread(dst, num * sizeof(double), 1, fp);
    fclose(fp);
}

void initdouble(double *&dst, int num)
{
    dst = (double *)malloc(num * sizeof(double));
    memset(dst, 0, num * sizeof(double));
}

__global__ void XW_(int in_dim, int out_dim, double *in_X, double *out_X, double *W, int v_num)
{

    int tid = threadIdx.x + blockIdx.x * blockDim.x; // 控制v_vum

    if (tid >= v_num)
        return;

    double *tmp_in_X = in_X;
    double *tmp_out_X = out_X;
    double *tmp_W = W;

    for (int j = 0; j < out_dim; j++)
    {
        for (int k = 0; k < in_dim; k++)
        {
            tmp_out_X[tid * out_dim + j] += tmp_in_X[tid * in_dim + k] * tmp_W[k * out_dim + j];
        }
    }
}

__global__ void AX_(int dim, double *in_X, double *out_X, int *index, int *edges, double *edges_val, int v_num)
{

    int tid = threadIdx.x + blockIdx.x * blockDim.x;
    if (tid >= v_num)
        return;

    int *nbrs = &edges[index[tid]];
    double *nbrs_val = &edges_val[index[tid]];

    int degree = index[tid + 1] - index[tid];

    for (int j = 0; j < degree; j++)
    {
        int nbr = nbrs[j];
        for (int k = 0; k < dim; k++)
        {
            out_X[dim * tid + k] += in_X[nbr * dim + k] * nbrs_val[j];
        }
    }
}

void LogSoftmax(int dim, double *X)
{

    for (int i = 0; i < v_num; i++)
    {
        double max = X[i * dim + 0];
        for (int j = 1; j < dim; j++)
        {
            if (X[i * dim + j] > max)
                max = X[i * dim + j];
        }

        double sum = 0;
        for (int j = 0; j < dim; j++)
        {
            sum += exp(X[i * dim + j] - max);
        }
        sum = log(sum);

        for (int j = 0; j < dim; j++)
        {
            X[i * dim + j] = X[i * dim + j] - max - sum;
        }
    }
}

double MaxRowSum(double *X, int dim)
{

    double max = -__FLT_MAX__;

    for (int i = 0; i < v_num; i++)
    {
        double sum = 0;
        for (int j = 0; j < dim; j++)
        {
            sum += X[i * dim + j];
        }
        if (sum > max)
            max = sum;
    }
    return max;
}

void freedoubles()
{
    free(X0);
    free(W1);
    free(X1);
    free(X1_inter);
    free(nodes_index);
    free(edges);
    free(edges_value);
    hipFree(d_X0);
    hipFree(d_X1_inter);
    hipFree(d_W1);
    hipFree(d_X1);
    hipFree(d_index);
    hipFree(d_edges);
    hipFree(d_edges_val);
}

void initGPUMemory()
{

    hipFree(0);

    hipMalloc(&d_X0, v_num * F0 * sizeof(double));
    hipMemcpy(d_X0, X0, v_num * F0 * sizeof(double), hipMemcpyHostToDevice);

    hipMalloc(&d_X1_inter, v_num * F1 * sizeof(double));
    hipMemcpy(d_X1_inter, X1_inter, v_num * F1 * sizeof(double), hipMemcpyHostToDevice);

    hipMalloc(&d_W1, F0 * F1 * sizeof(double));
    hipMemcpy(d_W1, W1, F0 * F1 * sizeof(double), hipMemcpyHostToDevice);

    hipMalloc(&d_X1, F1 * v_num * sizeof(double));
    hipMemcpy(d_X1, X1, F1 * v_num * sizeof(double), hipMemcpyHostToDevice);

    //    d_index, d_edge, d_edge_val

    hipMalloc(&d_index, (v_num + 1) * sizeof(int));
    hipMemcpy(d_index, nodes_index, (v_num + 1) * sizeof(int), hipMemcpyHostToDevice);

    hipMalloc(&d_edges, e_num * sizeof(int));
    hipMemcpy(d_edges, edges, e_num * sizeof(int), hipMemcpyHostToDevice);

    hipMalloc(&d_edges_val, e_num * sizeof(double));
    hipMemcpy(d_edges_val, edges_value, e_num * sizeof(double), hipMemcpyHostToDevice);
}

// self defined kernel -------------------------------------------------
template <const int BM, const int BN, const int BK, const int TM, const int TN>
__global__ void XWSGEMM_2D(int in_dim, int out_dim, double *in_X, double *out_X, double *W, int v_num)
{
    const uint cRow = blockIdx.x;
    const uint cCol = blockIdx.y;

    const uint totalResultsBlocktile = BM * BN;
    const uint numThreadsBlocktile = totalResultsBlocktile / (TM * TN);

    // assert(numThreadsBlocktile == blockDim.x);

    const int threadCol = threadIdx.x % (BN / TN);
    const int threadRow = threadIdx.x / (BN / TN);

    int vid = cRow * BM + threadRow;
    if (vid >= v_num) return;

    __shared__ double in_Xs[BM * BK];
    __shared__ double Ws[BK * BN];

    double* tmp_in_X = in_X;
    double* tmp_W = W;
    double* tmp_out_X = out_X;

    tmp_in_X += cRow * BM * in_dim;
    tmp_W += cCol * BN;
    tmp_out_X += cRow * BM * out_dim + cCol * BN;

    const uint innerRowA = threadIdx.x / BK;
    const uint innerColA = threadIdx.x % BK;
    const uint strideA = numThreadsBlocktile / BK;

    const uint innerRowB = threadIdx.x / BN;
    const uint innerColB = threadIdx.x % BN;
    const uint strideB = numThreadsBlocktile / BN;

    // allocate cache
    double threadResults[TM * TN] = {0.0};
    // register caches
    double regM[TM] = {0.0};
    double regN[TN] = {0.0};

    for (uint bkIdx = 0; bkIdx < in_dim; bkIdx += BK) 
    {
        for (uint loadOffset = 0; loadOffset < BM; loadOffset += strideA) 
        {
            in_Xs[(innerRowA + loadOffset) * BK + innerColA] =
                tmp_in_X[(innerRowA + loadOffset) * in_dim + innerColA];
        }

        for (uint loadOffset = 0; loadOffset < BK; loadOffset += strideB)
        {
            Ws[(innerRowB + loadOffset) * BN + innerColB] =
                tmp_W[(innerRowB + loadOffset) * out_dim + innerColB];
        }
        __syncthreads();

        // advance blocktile
        tmp_in_X += BK;     // move BK columns to right
        tmp_W += BK * out_dim; // move BK rows down

        // calculate per-thread results
        for (uint dotIdx = 0; dotIdx < BK; ++dotIdx) 
        {
            // block into registers
            for (uint i = 0; i < TM; ++i) {
                regM[i] = in_Xs[(threadRow * TM + i) * BK + dotIdx];
            }

            for (uint i = 0; i < TN; ++i) {
                regN[i] = Ws[dotIdx * BN + threadCol * TN + i];
            }

            for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
                for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) {
                threadResults[resIdxM * TN + resIdxN] +=
                    regM[resIdxM] * regN[resIdxN];
                }
            }
        }
        __syncthreads();
    }

    // write out the results
    for (uint resIdxM = 0; resIdxM < TM; ++resIdxM) {
        for (uint resIdxN = 0; resIdxN < TN; ++resIdxN) 
        {
            tmp_out_X[(threadRow * TM + resIdxM) * out_dim + threadCol * TN + resIdxN] = threadResults[resIdxM * TN + resIdxN];
        }
    }
}


template<const int AX_BM, const int AX_BN, const int AX_TM>
__global__ void AXLogSlice_(int dim, double *in_X, double *out_X, int *index, int *edges, double *edges_val, int v_num)
{
    const int cRow = blockIdx.x;
    const int cCol = blockIdx.y;

    // assert(AX_BN == dim);
    const int threadRow = threadIdx.x / AX_BN; // total 4 (0~3)
    const int threadCol = threadIdx.x % AX_BN; // total 16 (0~15)

    __shared__ double Ns[AX_BM * AX_BN];
    __shared__ double NVs[AX_BM * AX_BN];
    __shared__ double Xs[AX_BM * AX_BN];

    double* tmp_in = in_X + cCol * AX_BN;

    double* tmp_out_X = out_X;
    tmp_out_X += cRow * dim * AX_BM;

    int* tmp_index0 = index + cRow * AX_BM;
    int* tmp_index1 = tmp_index0 + 1;

    #pragma unroll
    for(int resIdx = 0; resIdx < AX_TM; resIdx++)
    {
        int innerRow = threadRow * AX_TM + resIdx;
        int vid = cRow * AX_BM + innerRow;

        if (vid >= v_num) return;

        int index0 = tmp_index0[innerRow];
        int index1 = tmp_index1[innerRow];

        int *nbrs = &edges[index0];
        double *nbrs_val = &edges_val[index0];

        int degree = index1 - index0;

        // AX
        double temp = 0;
        int upper_bound = 0;
        for(int i = 0; i<degree; i+=AX_BN)
        {
            if(threadCol + i < degree)
            {
                Ns[innerRow * AX_BN + threadCol] = nbrs[threadCol + i];
                NVs[innerRow * AX_BN + threadCol] = nbrs_val[threadCol + i];
            }

            upper_bound = degree;
            if(degree > (i+AX_BN))
            {
                upper_bound = i + AX_BN;
            }
            
            for (int j = 0; j < (upper_bound - i); j++)
            {
                int nbr = Ns[innerRow * AX_BN + j];
                temp += tmp_in[nbr * dim + threadCol] * NVs[innerRow * AX_BN + j];
            }

            __syncthreads();
        }

        // LogSoftmax
        double max = 0.0;
        double sum = 0.0;

        Xs[innerRow * dim + threadCol] = temp;

        double temp_max = 0;
        for(int j = 0; j < dim; j++)
        {
            int scan_index = (threadCol + j) % AX_BN;
            temp_max = Xs[innerRow * dim + scan_index];
            if(max < temp_max) max = temp_max;  
        }

        // __syncthreads();

        #pragma unroll
        for(int j = 0; j < dim; j++)
        {
            int sum_index = (threadCol + j) % AX_BN;
            sum += __expf(Xs[innerRow * dim + sum_index] - max);
        }
        sum = log(sum);

        // final output
        tmp_out_X[innerRow * dim + threadCol] = Xs[innerRow * dim + threadCol] - max - sum;
    }
}

// --------------------------------------------------------------

double GCN()
{
    hipMemset(d_X1_inter, 0, v_num * F1 * sizeof(double));
    hipMemset(d_X1, 0, F1 * v_num * sizeof(double));

    TimePoint start = chrono::steady_clock::now();

    // use pinned memory
    hipHostRegister(X1, v_num * F1 * sizeof(double), hipHostRegisterMapped);

#define CEIL_DIV(M, N) (((M) + (N)-1) / (N))
    const uint BM = 16;
    const uint BN = 16;
    const uint BK = 8;
    const uint TM = 2;
    const uint TN = 2;
    dim3 gridDim(CEIL_DIV(v_num, BM), CEIL_DIV(F1, BN));
    dim3 blockDim((BM * BN) / (TM * TN));

    const uint AX_TM = 2;
    dim3 ax_grid_size(CEIL_DIV(v_num, BM), CEIL_DIV(F1, BN));
    dim3 ax_block_size(BM * BN / AX_TM);

    // XW_2D
    XWSGEMM_2D<BM, BN, BK, TM, TN>
        <<<gridDim, blockDim>>>(F0, F1, d_X0, d_X1_inter, d_W1, v_num);

    // AX_LogSoftmax
    AXLogSlice_<BM, BN, AX_TM>
        <<<ax_grid_size, ax_block_size>>>(F1, d_X1_inter, d_X1, d_index, d_edges, d_edges_val, v_num);
    
    hipDeviceSynchronize();

    //Memcpy
    hipMemcpyAsync(X1, d_X1, sizeof(double) * v_num * F1, hipMemcpyDeviceToHost);

    TimePoint end = chrono::steady_clock::now();
    chrono::duration<double> l_durationSec = end - start;
    double l_timeMs = l_durationSec.count() * 1e3;

    return l_timeMs;
}

void XW_verify(int in_dim, int out_dim, double *in_X, double *out_X, double *W)
{
    double *tmp_in_X = in_X;
    double *tmp_out_X = out_X;
    double *tmp_W = W;

    for (int i = 0; i < v_num; i++)
    {   
        for (int j = 0; j < out_dim; j++)
        {
            for (int k = 0; k < in_dim; k++)
            {
                tmp_out_X[i * out_dim + j] += tmp_in_X[i * in_dim + k] * tmp_W[k * out_dim + j];
            }
        }
    }
}
void AX_verify(int dim, double *in_X, double *out_X)
{
    for (int i = 0; i < v_num; i++)

    {
        int *nbrs = &edges[nodes_index[i]];
        double *nbrs_val = &edges_value[nodes_index[i]];
        int degree = nodes_index[i + 1] - nodes_index[i];
        
        for (int j = 0; j < degree; j++)
        {
            int nbr = nbrs[j];
            for (int k = 0; k < dim; k++)
            {
                out_X[dim * i + k] += in_X[nbr * dim + k] * nbrs_val[j];
            }
        }
    }
}

void LogSoftmax_verify(int dim, double *X)
{

    for (int i = 0; i < v_num; i++)
    {
        double max = X[i * dim + 0];
        for (int j = 1; j < dim; j++)
        {
            if (X[i * dim + j] > max)
                max = X[i * dim + j];
        }

        double sum = 0;
        for (int j = 0; j < dim; j++)
        {
            sum += exp(X[i * dim + j] - max);
        }
        sum = log(sum);

        for (int j = 0; j < dim; j++)
        {
            X[i * dim + j] = X[i * dim + j] - max - sum;
        }
    }
}

bool verify(double max_sum)
{

    memset(X1_inter, 0, v_num * F1 * sizeof(double));
    memset(X1, 0, F1 * v_num * sizeof(double));

    XW_verify(F0, F1, X0, X1_inter, W1);

    // printf("Layer1 AX\n");
    AX_verify(F1, X1_inter, X1);

    // printf("Layer1 ReLU\n");
    LogSoftmax_verify(F1, X1);
    double verify_max_sum = MaxRowSum(X1, F1);
    printf("CPU_max_sum,  %6f\n", verify_max_sum);
    printf("GPU_max_sum,  %6f\n", max_sum);
    return fabs(max_sum - verify_max_sum) < 0.000001;
}

int main(int argc, char **argv)
{
    // !!! Attention !!!
    // Datasets: web-stanford ak_2010 dblp
    // Downloaded from：

    // 编译：
	//      hipify-perl gcn.cu > gcn.cpp
	//      hipcc gcn.cpp -o gcn
    //
    // 执行：仅供测试参考，队伍提交直接执行slurm.sh 即可
    //      可执行程序需接收5个参数，分别为：
	//      输入顶点特征长度F0，第一层顶点特征长度F1，图结构文件名，输入顶点特征矩阵文件名，第一层权重矩阵文件名
    //      ./gcn 128 16 graph/web-stanford_nodes_281903_edges_1992636_core_71.txt embedding/web-stanford_F0_128.bin weight/web-stanford_F0_128_F1_16.bin
    //      ./gcn 128 16 graph/com-dblp_nodes_317080_edges_1049866_core_113.txt embedding/dblp_F0_128.bin weight/dblp_F0_128_F1_16.bin
    //      ./gcn 128 16 graph/ak_2010.txt embedding/ak_2010_F0_128.bin weight/ak_2010_F0_128_F1_16.bin
    
    // 要求： 
    //      只允许修改GCN()函数里包含的代码；其余代码不允许修改，一旦发现取消成绩。

    // 评分：
    //      计算耗时显示 程序运行后会循环计算五次，评分是主要查看平均耗时。

    // 提交：
    //      查看slurm.sh 文件
    F0 = atoi(argv[1]);
    F1 = atoi(argv[2]);
    readGraph(argv[3]);
    readdouble(argv[4], X0, v_num * F0);
    readdouble(argv[5], W1, F0 * F1);
    initdouble(X1, v_num * F1);
    initdouble(X1_inter, v_num * F1);

    raw_graph_to_AdjacencyList();
    edgeNormalization();
    to_csr();
    initGPUMemory();

    double max_sum = 0, ave_timeMs = 0;
    int ROUNDs = 20;

    // warm up
    GCN();

    for (int i = 0; i < ROUNDs; i++)
    {
        // ################
        //
        ave_timeMs += GCN();
        // ################
        // Time point at the end of the computation
        // Compute the max row sum for result verification
        max_sum = MaxRowSum(X1, F1);

        // The max row sum and the computing time should be print
    }

    printf("verify\n");

    if (verify(max_sum))
    {
        printf("True\n");
    }
    else
    {
        printf("False\n");
    }

    printf("%f\n", ave_timeMs / ROUNDs);

    freedoubles();
}