forked from ccf-ai-infra/GPUCodeForces
26 lines
769 B
Plaintext
26 lines
769 B
Plaintext
Write a custom CUDA kernel for optimized matrix multiplication (GEMM).
|
||
|
||
The standard matrix multiplication is defined as:
|
||
C = A × B
|
||
|
||
Where A is of shape (M, K), B is of shape (K, N), and C is of shape (M, N).
|
||
|
||
You should optimize the matrix multiplication using:
|
||
1. Shared memory tiling for better memory access patterns
|
||
2. Coalesced memory access
|
||
3. Thread block tiling to maximize parallelism
|
||
4. Avoid redundant memory loads
|
||
|
||
The kernel should be significantly faster than PyTorch's default matmul implementation for large matrices.
|
||
|
||
You are given the following architecture:
|
||
|
||
import torch
|
||
import torch.nn as nn
|
||
|
||
class Model(nn.Module):
|
||
def __init__(self):
|
||
super(Model, self).__init__()
|
||
|
||
def forward(self, A, B):
|
||
return torch.matmul(A, B) |