added mma and st inst. Only different b mode needs to be added

1 files changed, 85 insertions, 16 deletions

diff --git a/cuda-kernels/v4p_kernel.cu b/cuda-kernels/v4p_kernel.cu
index ced0e44..83055de 100644
--- a/cuda-kernels/v4p_kernel.cu
+++ b/cuda-kernels/v4p_kernel.cu
@@ -30,30 +30,99 @@ const int WMMA_M = 16;
 const int WMMA_N = 16;
 const int WMMA_K = 16;
 
-__global__ void v4p_example(int *a_int32, int *b_int8, int *c,int *d_int32, int M, int N, int K) {
+__global__ void v4p_example(int *a_int32, int *b_int4, int *c,int *d_int32, int M, int N, int K) {
 	
 	int registers_a[8];
+	int registers_b[8];
+	int registers_c[8];
+	int registers_d[8];
 	int register_b;		//contains 8 4bit b elements
    	int idx = blockDim.x * blockIdx.x + threadIdx.x;
 	
 	asm("/*");
 	asm("CPTX_BEGIN");
-	asm("vp.load.b4.sync.row.m16n16k16.s32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8],%9;" : 
+	asm("vp.load.a.sync.row.m16n16k16.s32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8],%9;" : 
 	"=r"(registers_a[0]), "=r"(registers_a[1]),"=r"(registers_a[2]),"=r"(registers_a[3]),
 	"=r"(registers_a[4]),"=r"(registers_a[5]),"=r"(registers_a[6]),"=r"(registers_a[7]):
-	"l"(b_int8),"r"(M)
+	"l"(a_int32),"r"(M)
 	);
 	asm("CPTX_END");
 	asm("*/");
+	asm("/*");
+	asm("CPTX_BEGIN");
+	asm("vp.load.b4.sync.row.m16n16k16.s32 {%0},[%1],%2;" : 
+	"=r"(registers_b[0]):
+	"l"(b_int4),"r"(M)
+	);
+	asm("CPTX_END");
+	asm("*/");
+	asm("/*");
+	asm("CPTX_BEGIN");
+	asm("vp.load.c.sync.row.m16n16k16.s32 {%0,%1,%2,%3,%4,%5,%6,%7},[%8],%9;" : 
+	"=r"(registers_c[0]), "=r"(registers_c[1]),"=r"(registers_c[2]),"=r"(registers_c[3]),
+	"=r"(registers_c[4]),"=r"(registers_c[5]),"=r"(registers_c[6]),"=r"(registers_c[7]):
+	"l"(c),"r"(M)
+	);
+	asm("CPTX_END");
+	asm("*/");
+	//B4
+	asm("/*");
+	asm("CPTX_BEGIN");
+	asm("vp.mma.sync.row.row.m16n16k16.s32 {%0, %1, %2, %3, %4, %5, %6, %7}, {%8, %9, %10, %11, %12, %13, %14, %15}, {%16}, {%17, %18, %19, %20, %21, %22, %23, %24};" : 
+	"=r"(registers_d[0]), "=r"(registers_d[1]),"=r"(registers_d[2]),"=r"(registers_d[3]),
+	"=r"(registers_d[4]),"=r"(registers_d[5]),"=r"(registers_d[6]),"=r"(registers_d[7]):
+	"r"(registers_a[0]),"r"(registers_a[1]),"r"(registers_a[2]),"r"(registers_a[3]),
+	"r"(registers_a[4]),"r"(registers_a[5]),"r"(registers_a[6]),"r"(registers_a[7]),
+	"r"(registers_b[0]),
+	"r"(registers_c[0]),"r"(registers_c[1]),"r"(registers_c[2]),"r"(registers_c[3]),
+	"r"(registers_c[4]),"r"(registers_c[5]),"r"(registers_c[6]),"r"(registers_c[7])
+	);
+	asm("CPTX_END");
+	asm("*/");
+
+	//B8
+	//asm("CPTX_END");
+	//asm("*/");
+	//asm("/*");
+	//asm("CPTX_BEGIN");
+	//asm("vp.mma.sync.row.row.m16n16k16.s32 {%0, %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};" : 
+	//"=r"(registers_d[0]), "=r"(registers_d[1]),"=r"(registers_d[2]),"=r"(registers_d[3]),
+	//"=r"(registers_d[4]),"=r"(registers_d[5]),"=r"(registers_d[6]),"=r"(registers_d[7]):
+	//"r"(registers_a[0]),"r"(registers_a[1]),"r"(registers_a[2]),"r"(registers_a[3]),
+	//"r"(registers_a[4]),"r"(registers_a[5]),"r"(registers_a[6]),"r"(registers_a[7]),
+	//"r"(registers_b[0]),"r"(registers_b[1]),
+	//"r"(registers_c[0]),"r"(registers_c[1]),"r"(registers_c[2]),"r"(registers_c[3]),
+	//"r"(registers_c[4]),"r"(registers_c[5]),"r"(registers_c[6]),"r"(registers_c[7])
+	//);
+	//asm("CPTX_END");
+	//asm("*/");
+
+	//B16
+	//asm("CPTX_END");
+	//asm("*/");
+	//asm("/*");
+	//asm("CPTX_BEGIN");
+	//asm("vp.mma.sync.row.row.m16n16k16.s32 {%0, %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};" : 
+	//"=r"(registers_d[0]), "=r"(registers_d[1]),"=r"(registers_d[2]),"=r"(registers_d[3]),
+	//"=r"(registers_d[4]),"=r"(registers_d[5]),"=r"(registers_d[6]),"=r"(registers_d[7]):
+	//"r"(registers_a[0]),"r"(registers_a[1]),"r"(registers_a[2]),"r"(registers_a[3]),
+	//"r"(registers_a[4]),"r"(registers_a[5]),"r"(registers_a[6]),"r"(registers_a[7]),
+	//"r"(registers_b[0]),"r"(registers_b[1]),"r"(registers_b[2]),"r"(registers_b[3]),
+	//"r"(registers_c[0]),"r"(registers_c[1]),"r"(registers_c[2]),"r"(registers_c[3]),
+	//"r"(registers_c[4]),"r"(registers_c[5]),"r"(registers_c[6]),"r"(registers_c[7])
+	//);
+	//asm("CPTX_END");
+	//asm("*/");
+
 
-	b_int8[0]=registers_a[7];
-	b_int8[1]=registers_a[7];
-	b_int8[2]=registers_a[7];
-	b_int8[3]=registers_a[7];
-	b_int8[4]=registers_a[7];
-	b_int8[5]=registers_a[7];
-	b_int8[6]=registers_a[7];
-	b_int8[7]=registers_a[7];
+	d_int32[0]=registers_d[0];
+	d_int32[1]=registers_d[1];
+	d_int32[2]=registers_d[2];
+	d_int32[3]=registers_d[3];
+	d_int32[4]=registers_d[4];
+	d_int32[5]=registers_d[5];
+	d_int32[6]=registers_d[6];
+	d_int32[7]=registers_d[7];
 }
 
 __global__ void convertFp32ToFp16 (half *out, float *in, int n) {
@@ -200,7 +269,6 @@ int main(int argc, char* argv[]) {
    cudaErrCheck(cudaMemcpy(b_int32,b_host_wmma,  MATRIX_K * MATRIX_N * sizeof(int), cudaMemcpyHostToDevice));
    cudaErrCheck(cudaMemcpy(c_int32,c_host_wmma,  MATRIX_M * MATRIX_N * sizeof(int), cudaMemcpyHostToDevice));
 
-   #define TEST4
    #ifdef TEST16
    	convertInt32ToInt16 <<< (MATRIX_M * MATRIX_K + 255) / 256, 256 >>> (a_int16, a_int32, MATRIX_M * MATRIX_K);
    	convertInt16ToInt32 <<< (MATRIX_M * MATRIX_K + 255) / 256, 256 >>> (d_int32, a_int16, MATRIX_M * MATRIX_K);
@@ -216,6 +284,7 @@ int main(int argc, char* argv[]) {
   	convertInt4ToInt32 <<< (MATRIX_M * MATRIX_K + 255) / 256, 256 >>> (d_int32, a_int4, MATRIX_M * MATRIX_K);
  	cudaErrCheck(cudaMemcpy(d_host_wmma, d_int32, MATRIX_M * MATRIX_N * sizeof(int), cudaMemcpyDeviceToHost));
    #endif
+   convertInt32ToInt4 <<< (MATRIX_M * MATRIX_K + 255) / 256, 256 >>> (b_int4, b_int32, MATRIX_M * MATRIX_K);
    //convertFp32ToFp16 <<< (MATRIX_K * MATRIX_N + 255) / 256, 256 >>> (b_fp16, b_fp32, MATRIX_K * MATRIX_N);
    //convertFp32ToFp16 <<< (MATRIX_M * MATRIX_N + 255) / 256, 256 >>> (c_fp16, c_fp32, MATRIX_K * MATRIX_N);
 
@@ -223,10 +292,10 @@ int main(int argc, char* argv[]) {
 //AAMIR   printf("\nM = %d, N = %d, K = %d. \n", MATRIX_M, MATRIX_N, MATRIX_K);
 //AAMIR   
 //AAMIR   printf("Running with wmma...\n");
-   //cudaErrCheck(cudaEventRecord(startWMMA));
-   //v4p_example <<< 1, 32>>> (a_int32, a_int8, a_int32, d_int32, MATRIX_M, MATRIX_N, MATRIX_K);
-   //cudaErrCheck(cudaEventRecord(stopWMMA));
-   //cudaErrCheck(cudaEventSynchronize(stopWMMA));
+   cudaErrCheck(cudaEventRecord(startWMMA));
+   v4p_example <<< 1, 32>>> (a_int32, b_int4, c_int32, d_int32, MATRIX_M, MATRIX_N, MATRIX_K);
+   cudaErrCheck(cudaEventRecord(stopWMMA));
+   cudaErrCheck(cudaEventSynchronize(stopWMMA));
 
 //AAMIR
 //AAMIR   // Error checking