Matrix Multiplication

Given an M × K matrix A and a K × N matrix B of floats, write their product into the M × N matrix C:

C[i][j] = A[i][0] * B[0][j] + A[i][1] * B[1][j] + ... + A[i][K-1] * B[K-1][j]

All three matrices are stored row by row, so A[i][k] sits at A[i * K + k], B[k][j] at B[k * N + j] and C[i][j] at C[i * N + j].

Requirements

• Implement solve in the starter of the framework you choose; keep its signature.
  • CUDA: extern "C" void solve(const float* A, const float* B, float* C, int M, int N, int K)
  • Triton, PyTorch: def solve(A: torch.Tensor, B: torch.Tensor, C: torch.Tensor, M: int, N: int, K: int) (A, B, C have shapes (M, K), (K, N), (M, N))
• A, B and C are already on the GPU. Do not modify A or B.
• Write the result into C in place; its contents on entry are undefined.
• Compute in float32. TF32 Tensor Cores (A100 and newer) round the inputs to 10 mantissa bits, which is far outside the tolerance. Triton's tl.dot uses TF32 for float32 inputs by default; pass input_precision="ieee".

Examples

Input:  A = [[1.0, 2.0, 3.0],     B = [[ 7.0,  8.0],
             [4.0, 5.0, 6.0]]          [ 9.0, 10.0],
                                       [11.0, 12.0]]     M = 2, N = 2, K = 3
Output: C = [[ 58.0,  64.0],
             [139.0, 154.0]]

C[0][0] = 1 × 7 + 2 × 9 + 3 × 11 = 58.

Input:  A = [[1.0],     B = [[4.0, 5.0]]     M = 3, N = 2, K = 1
             [2.0],
             [3.0]]
Output: C = [[ 4.0,  5.0],
             [ 8.0, 10.0],
             [12.0, 15.0]]

Constraints

• 1 ≤ M, N, K ≤ 8192
• In the tests, every entry of A and B lies in [-1, 1]
• An output is accepted if it is within 1e-4 + 1e-4 × |expected| of the expected value
• Performance is measured with M = 8192, N = 6144, K = 4096