矩阵乘法

给定 M × K 的 float 矩阵 A 和 K × N 的 float 矩阵 B，把它们的乘积写进 M × N 的矩阵 C：

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

三个矩阵都按行存放：A[i][k] 在 A[i * K + k]，B[k][j] 在 B[k * N + j]，C[i][j] 在 C[i * N + j]。

要求

• 在所选框架的初始代码里实现 solve，签名保持不变。
  • 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 的形状分别是 (M, K)、(K, N)、(M, N)）
• A、B、C 已经在 GPU 上。不要修改 A 和 B。
• 结果原地写进 C；进入 solve 时 C 里的内容是未定义的。
• 用 float32 计算。A100 及更新的 GPU 上的 TF32 Tensor Core 会把输入舍入到 10 位尾数，误差远超容差。 Triton 的 tl.dot 对 float32 输入默认用 TF32，要传 input_precision="ieee"。

示例

输入：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
输出：C = [[ 58.0,  64.0],
           [139.0, 154.0]]

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

输入：A = [[1.0],     B = [[4.0, 5.0]]     M = 3, N = 2, K = 1
           [2.0],
           [3.0]]
输出：C = [[ 4.0,  5.0],
           [ 8.0, 10.0],
           [12.0, 15.0]]

约束

• 1 ≤ M, N, K ≤ 8192
• 测试里 A、B 的每个元素都在 [-1, 1] 内
• 输出与期望值之差不超过 1e-4 + 1e-4 × |期望值| 即判对
• 性能测试规模：M = 8192，N = 6144，K = 4096