callrun_correctness
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n_cpp_src = r'''\n#include <torch/extension.h>\ntorch::Tensor matmul_forward(torch::Tensor A, torch::Tensor B);\n'''\n_cuda_src = r'''\n#include <torch/extension.h>\n#include <ATen/cuda/CUDAContext.h>\n#include <cuda_runtime.h>\n#include <mma.h>\nusing namespace nvcuda;\n#define BM 64\n#define BN 64\n#define BK 64\n#define WM 16\n__device__ __forceinline__ float to_tf32(float x){ float y; asm(\"cvt.rna.tf32.f32 %0, %1;\" : \"=f\"(y) : \"f\"(x)); return y; }\n__global__ void wmma_tf32x3_kernel(const float* __restrict__ A,const float* __restrict__ B,float* __restrict__ C,int K,int split_chunk){\n int block_m=blockIdx.y*BM, block_n=blockIdx.x*BN;\n int tid=threadIdx.x, warp_id=tid>>5;\n extern __shared__ float smem[];\n float* As=smem; float* Bs=As+BM*BK; float* Al=Bs+BK*BN; float* Bl=Al+BM*BK;\n wmma::fragment<wmma::accumulator,16,16,8,float> acc0,acc1,acc2,acc3;\n wmma::fill_fragment(acc0,0.0f); wmma::fill_fragment(acc1,0.0f); wmma::fill_fragment(acc2,0.0f); wmma::fill_fragment(acc3,0.0f);\n int k_start=blockIdx.z*split_chunk; int k_end=min(k_start+split_chunk,K);\n for(int kb=k_start; kb<k_end; kb+=BK){\n for(int idx=tid; idx<BM*BK; idx+=128){ int r=idx/BK,c=idx-r*BK; float x=A[(block_m+r)*K+kb+c]; float hx=to_tf32(x); As[idx]=x; Al[idx]=x-hx; }\n for(int idx=tid; idx<BK*BN; idx+=128){ int r=idx/BN,c=idx-r*BN; float x=B[(kb+r)*256+block_n+c]; float hx=to_tf32(x); Bs[idx]=x; Bl[idx]=x-hx; }\n __syncthreads();\n #pragma unroll\n for(int kk=0; kk<BK; kk+=8){\n wmma::fragment<wmma::matrix_a,16,16,8,wmma::precision::tf32,wmma::row_major> ah,al;\n wmma::fragment<wmma::matrix_b,16,16,8,wmma::precision::tf32,wmma::row_major> bh0,bh1,bh2,bh3, bl0,bl1,bl2,bl3;\n wmma::load_matrix_sync(ah, As+(warp_id*WM)*BK+kk, BK); wmma::load_matrix_sync(al, Al+(warp_id*WM)*BK+kk, BK);\n wmma::load_matrix_sync(bh0, Bs+kk*BN+0, BN); wmma::load_matrix_sync(bl0, Bl+kk*BN+0, BN);\n wmma::load_matrix_sync(bh1, Bs+kk*BN+16, BN); wmma::load_matrix_sync(bl1, Bl+kk*BN+16, BN);\n wmma::load_matrix_sync(bh2, Bs+kk*BN+32, BN); wmma::load_matrix_sync(bl2, Bl+kk*BN+32, BN);\n wmma::load_matrix_sync(bh3, Bs+kk*BN+48, BN); wmma::load_matrix_sync(bl3, Bl+kk*BN+48, BN);\n wmma::mma_sync(acc0, ah, bh0, acc0); wmma::mma_sync(acc0, ah, bl0, acc0); wmma::mma_sync(acc0, al, bh0, acc0);\n wmma::mma_sync(acc1, ah, bh1, acc1); wmma::mma_sync(acc1, ah, bl1, acc1); wmma::mma_sync(acc1, al, bh1, acc1);\n wmma::mma_sync(acc2, ah, bh2, acc2); wmma::mma_sync(acc2, ah, bl2, acc2); wmma::mma_sync(acc2, al, bh2, acc2);\n wmma::mma_sync(acc3, ah, bh3, acc3); wmma::mma_sync(acc3, ah, bl3, acc3); wmma::mma_sync(acc3, al, bh3, acc3);\n }\n __syncthreads();\n }\n float* Cs=smem;\n wmma::store_matrix_sync(Cs+(warp_id*WM)*BN+0,acc0,BN,wmma::mem_row_major);\n wmma::store_matrix_sync(Cs+(warp_id*WM)*BN+16,acc1,BN,wmma::mem_row_major);\n wmma::store_matrix_sync(Cs+(warp_id*WM)*BN+32,acc2,BN,wmma::mem_row_major);\n wmma::store_matrix_sync(Cs+(warp_id*WM)*BN+48,acc3,BN,wmma::mem_row_major);\n __syncthreads();\n for(int idx=tid; idx<BM*BN; idx+=128){ int r=idx/BN,c=idx-r*BN; atomicAdd(C+(block_m+r)*256+block_n+c,Cs[idx]); }\n}\ntorch::Tensor matmul_forward(torch::Tensor A, torch::Tensor B){ int M=A.size(0),K=A.size(1),N=B.size(1); auto C=torch::empty({M,N},A.options()); cudaMemsetAsync(C.data_ptr<float>(),0,(size_t)M*N*sizeof(float),at::cuda::getCurrentCUDAStream()); int split_chunk=8192; int splits=(K+split_chunk-1)/split_chunk; dim3 grid((N+BN-1)/BN,(M+BM-1)/BM,splits); size_t smem=4*BM*BK*sizeof(float); wmma_tf32x3_kernel<<<grid,128,smem,at::cuda::getCurrentCUDAStream()>>>(A.data_ptr<float>(),B.data_ptr<float>(),C.data_ptr<float>(),K,split_chunk); return C; }\n'''\n_ext=load_inline(name='wmma_tf32x3_v3', cpp_sources=_cpp_src, cuda_sources=_cuda_src, functions=['matmul_forward'], extra_cuda_cflags=['-O3'], verbose=False)\nclass ModelNew(nn.Module):\n def __init__(self): super(ModelNew,self).__init__()\n def forward(self,A:torch.Tensor,B:torch.Tensor)->torch.Tensor: return _ext.matmul_forward(A,B)\n"
}