int current_inner_pipe_idx = 0;
#pragma unroll
for (int n_tile = 0; n_tile < BLOCK_N_TILES_WMMA; ++n_tile) {
int next_inner_pipe_idx = 1 - current_inner_pipe_idx;
if (n_tile < BLOCK_N_TILES_WMMA - 1) {
int b_col_start_in_tile_next = (n_tile + 1) * WMMA_N;
wmma::load_matrix_sync(b_frag_inner_pipe[next_inner_pipe_idx], &Bsub_pipe[compute_pipe_idx][b_col_start_in_tile_next][0], WMMA_K + SKEW_HALF);
}
wmma::mma_sync(acc_frag[n_tile], a_frag, b_frag_inner_pipe[current_inner_pipe_idx], acc_frag[n_tile]);
current_inner_pipe_idx = next_inner_pipe_idx;
}
}
__syncthreads();
// Store results from accumulator fragments to global memory
#pragma unroll
for (int n_tile = 0; n_tile < BLOCK_N_TILES_WMMA; ++n_tile) {
wmma::store_matrix_sync(&C_shmem_output_buffers[warp_id][0][0], acc_frag[n_tile], WMMA_N, wmma::mem_row_major);
for (int elem_idx_in_frag = lane_id; elem_idx_in_frag < WMMA_M * WMMA_N; elem_idx_in_frag += warpSize) {
int r_frag = elem_idx_in_frag / WMMA_N;
int c_frag = elem_idx_in_frag % WMMA_N;
int oc_idx = block_row_gemm_start + (warp_id * WMMA_M) + r_frag;
int offset_in_block_N_processing = (n_tile * WMMA_N) + c_frag;
if (oc_idx < C_out && offset_in_block_N_processing < TILE_N_PER_BLOCK &&
n_params_sh[offset_in_block_N_processing].isValidPixel) {
const NDecomposed& current_n_params = n_params_sh[offset_in_block_N_processing];
int ow_eff = current_n_params.ow_eff;
int oh_eff = current_n_params.oh_eff;
int n_batch_idx = current_n_params.n_batch_idx;
float val = C_shmem_output_buffers[warp_id][r_frag][c_frag];
if (bias_ptr != nullptr) {
val += bias_ptr[oc_idx];
}
output_ptr[n_batch_idx * C_out * H_out * W_out +
oc_idx * H_out * W_out +
oh_eff * W_out +
ow_eff] = val;
}
}
}
}
torch::Tensor conv2d_implicit_gemm_cuda(
torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
int N_batch, int C_in, int H_in, int W_in,
int C_out, int K_h, int K_w,
int stride_h, int stride_w, int pad_h, int pad_w,
int H_out, int W_out) {
TORCH_CHECK(input.device().is_cuda(), "Input must be a CUDA tensor");
TORCH_CHECK(weight.device().is_cuda(), "Weight must be a CUDA tensor");
TORCH_CHECK(input.dtype() == torch::kFloat32, "Input must be float32");
TORCH_CHECK(weight.dtype() == torch::kFloat32, "Weight must be float32");
if (bias.defined()) {
TORCH_CHECK(bias.device().is_cuda(), "Bias must be a CUDA tensor");
TORCH_CHECK(bias.dtype() == torch::kFloat32, "Bias must be float32");
TORCH_CHECK(bias.dim() == 1 && bias.size(0) == C_out, "Bias has wrong shape");
}
TORCH_CHECK(input.dim() == 4, "Input must be 4D");
TORCH_CHECK(weight.dim() == 4, "Weight must be 4D");
TORCH_CHECK(input.size(0) == N_batch, "Input N_batch mismatch");
TORCH_CHECK(input.size(1) == C_in, "Input C_in mismatch");
TORCH_CHECK(input.size(2) == H_in, "Input H_in mismatch");
TORCH_CHECK(input.size(3) == W_in, "Input W_in mismatch");
TORCH_CHECK(weight.size(0) == C_out, "Weight C_out mismatch");
TORCH_CHECK(weight.size(1) == C_in, "Weight C_in mismatch");
TORCH_CHECK(weight.size(2) == K_h, "Weight K_h mismatch");
TORCH_CHECK(weight.size(3) == K_w, "Weight K_w mismatch");
auto output = torch::zeros({N_batch, C_out, H_out, W_out}, input.options());
const int M_gemm = C_out;
const int N_gemm = N_batch * H_out * W_out;
const int K_gemm = C_in * K_h * K_w;
if (M_gemm == 0 || N_gemm == 0) {
return output;
}
if (K_gemm == 0) {
if (bias.defined()) {
output = output + bias.reshape({1, C_out, 1, 1});
}
return output;
}
dim3 block_dim(THREADS_PER_BLOCK);
dim3 grid_dim(
(N_gemm + TILE_N_PER_BLOCK - 1) / TILE_N_PER_BLOCK,
(M_gemm + TILE_M_PER_BLOCK - 1) / TILE_M_PER_BLOCK
);
const float* bias_ptr_data = bias.defined() ? bias.data_ptr () : nullptr;
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
conv2d_implicit_gemm_wmma_kernel<< >>(
input.data_ptr (),
weight.data_ptr (),
bias_ptr_data,
output.data_ptr (),
N_batch, C_in, H_in, W_in,
C_out, K_h, K_w,
stride_h, stride_w, pad_h, pad_w,
H_out, W_out,
M_gemm, N_gemm, K_gemm
);
AT_CUDA_CHECK(cudaGetLastError());
return output;
}
"""
conv2d_implicit_gemm_cuda_declaration = r"""
torch::Tensor conv2d_implicit_gemm_cuda(
torch::Tensor input, torch::Tensor weight, torch::Tensor bias,
int N_batch, int C_in, int H_in, int W_in,
int C_out, int K_h, int K_w,
int stride_h, int stride_w, int pad_h, int pad_w,
int H_out, int W_out);
"""
# JIT compile the CUDA kernel
custom_conv2d_wmma_ops = load_inline(
name="custom_conv2d_wmma_ops_optimized_k_pipe_vec_smem", # Changed name to avoid collision
cpp_sources=conv2d_implicit_gemm_cuda_declaration,
cuda_sources=conv2d_implicit_gemm_cuda_source,
functions=["conv2d_implicit_gemm_cuda"],
verbose=True,
extra_cuda_cflags=["-arch=sm_70", "--use_fast_math", "-std=c++17"]
)
class ModelNew(nn.Module):
def __init__(self, num_classes=1000): # num_classes is part of original signature, kept for consistency
super(ModelNew, self).__init__()
# Define Conv1 parameters (matching the original model)
self.in_channels = 3
self.out_channels = 96
self.kernel_size_val = 11 # Assuming square kernel
self.stride_val = 4 # Assuming square stride
self.padding_val = 2 # Assuming square padding
# Create a temporary Conv2d layer to initialize weights and bias
temp_conv = nn.Conv2d(
in_channels=self.in_channels,
out_channels=self.out_channels,
kernel_size=self.kernel_size_val,
stride=self.stride_val,
padding=self.padding_val,
bias=True # nn.Conv2d has bias=True by default
)
self.conv1_weight = nn.Parameter(temp_conv.weight.detach().clone())
if temp_conv.bias is not None:
self.conv1_bias = nn.Parameter(temp_conv.bias.detach().clone())
else:
# Correctly register 'conv1_bias' as None if not present
self.register_parameter('conv1_bias', None)
self.custom_conv_op = custom_conv2d_wmma_ops.conv2d_implicit_gemm_cuda
def forward(self, x):
N_batch = x.size(0)
# C_in_runtime = x.size(1) # Should match self.in_channels
H_in = x.size(2)
W_in = x.size(3)
# Calculate output dimensions
H_out = (H_in + 2 * self.padding_val - self.kernel_size_val) // self.stride_val + 1
W_out = (W_in + 2 * self.padding_val - self.kernel_size_val) // self.stride_val + 1
# Bias tensor handling: pass an undefined tensor if bias is None.
# The C++ TORCH_CHECK(bias.defined()) handles this by providing nullptr to kernel.
bias_tensor = self.conv1_bias if self.conv1_bias is not None else torch.Tensor()
x = self.custom_conv_op(
x, self.conv1_weight, bias_tensor,
N_batch, self.in_channels, H_in, W_in,
self.out_channels, self.kernel_size_val, self.kernel_size_val, # K_h, K_w
self.stride_val, self.stride_val, # stride_h, stride_w
self.padding_val, self.padding_val, # pad_h, pad_w
H_out, W_out
)
return x</code></pre>
参考资料：
