CUDA任意维度的矩阵相乘

__global__ void matrixMul( float* A, float* B, float* C, int hA,int wA, int wB)
{
int bx = blockIdx.x;
int by = blockIdx.y;
int tx = threadIdx.x;
int ty = threadIdx.y;


float Csub = 0.0f;



for (int j=0;j<wA;j+=BLOCK_DIM)
{
__shared__ float AS[BLOCK_DIM][BLOCK_DIM];
__shared__ float BS[BLOCK_DIM][BLOCK_DIM];


if(((by*BLOCK_DIM+ty)<hA)&&((tx+j)<wA))
{
AS[ty][tx] = A[(by*BLOCK_DIM+ty)*wA+tx+j];
}
else
{
AS[ty][tx] = 0;
}
if(((ty+j)<wA)&&((bx*BLOCK_DIM+tx)<wB))
{
BS[ty][tx] = B[(ty+j)*wB+bx*BLOCK_DIM+tx];
}
else
{
BS[ty][tx] = 0;
}


__syncthreads();


//Kahan's Summation Formula
for (int k = 0; k < BLOCK_DIM; ++k)
{
Csub += AS[ty][k]*BS[k][tx];
}


__syncthreads();
}

if(((by*BLOCK_DIM+ty)<hA)&&((bx*BLOCK_DIM+tx)<wB))
{
int c = wB*BLOCK_DIM*by + BLOCK_DIM*bx; 
C[c + wB * ty + tx] = Csub; 


}


}




//



int wA =  src.feat[t].size[2];;                       //1的宽 2的高
int wB = src.feat[t].size[0]*src.feat[t].size[1];  //2 3矩阵的宽
int hA = pmatrix.size[1];                  //1 3矩阵的高


dim3 mygrid(((wB+BLOCK_DIM-1)/BLOCK_DIM),(hA+BLOCK_DIM-1)/BLOCK_DIM);
dim3 myblock(BLOCK_DIM,BLOCK_DIM);
matrixMul<<<mygrid,myblock>>>(dev_matrix,dev_pyra[t],dev_pyrapro[t],hA,wA,wB);