CUDA:行列の積

cudaで正方行列の積を計算してみた。ヘッダファイルのbook.hについては下記の本を参照してください。
sharedメモリを使うと高速に計算できるらしい。

CUDA by Example 汎用GPUプログラミング入門

#include "../common/book.h"

#define N 100

__global__ void mat(float *a,float *b,float *c){
    int xid=threadIdx.x+blockIdx.x*blockDim.x;
    int yid=threadIdx.y+blockIdx.y*blockDim.y;
    int idx=xid*N+yid;
    
    if(xid<N&&yid<N){
        float x=0.0;
        for(int i=0;i<N;i++){
            x+=a[xid*N+i]*b[i*N+yid];
        }
        c[idx]=x;
    }
        
}

int main(void){
    float *a,*b,*c;
    float *dev_a,*dev_b,*dev_c;
    dim3 grids(N,N);
    dim3 threads((N+grids.x-1)/grids.x,(N+grids.y-1)/grids.y);


    a=(float *)malloc(N*N*sizeof(float));
    b=(float *)malloc(N*N*sizeof(float));
    c=(float *)malloc(N*N*sizeof(float));

    cudaMalloc((void **)&dev_a,sizeof(float)*N*N);
    cudaMalloc((void **)&dev_b,sizeof(float)*N*N);
    cudaMalloc((void **)&dev_c,sizeof(float)*N*N);

    for(int i=0;i<N;i++){
        for(int j=0;j<N;j++){
            a[i*N+j]=2;
            if(i==j)
                b[i*N+j]=1;
            else
                b[i*N+j]=0;
        }
    }
    HANDLE_ERROR(cudaMemcpy(dev_a,a,N*N*sizeof(float),cudaMemcpyHostToDevice));
    HANDLE_ERROR(cudaMemcpy(dev_b,b,N*N*sizeof(float),cudaMemcpyHostToDevice));

    mat<<<grids,1>>>(dev_a,dev_b,dev_c);

    HANDLE_ERROR(cudaMemcpy(c,dev_c,N*N*sizeof(float),cudaMemcpyDeviceToHost));
    
    for(int i=0;i<N;i++){
        for(int j=0;j<N;j++){
            printf("%2.2f ",c[i*N+j]);
        }
        puts("");
    }

    free(a);
    free(b);
    free(c);
    cudaFree(dev_a);
    cudaFree(dev_b);
    cudaFree(dev_c);

    return 0;
}