Batched Matrix Multiplication

题面

Implement a batched matrix multiplication in FP32. 给定 A∈R^{B×M×K} 与 B∈R^{B×K×N}，计算 C∈R^{B×M×N}，对每个 batch 索引 b：C_b = A_b × B_b。

Implementation Requirements

• External libraries are not permitted
• The solve function signature must remain unchanged
• The final result must be stored in the C array

Examples

B = 2, M = 2, K = 3, N = 2 → C 如页面示例

Constraints

• 1 ≤ B ≤ 128
• 1 ≤ M, N, K ≤ 1024
• Performance is measured with K = 256, M = 256, N = 256