3D Tensor-Matrix Multiplication

Perform 3D tensor-matrix multiplication of two tensors:

$$C[i][j][l] = \sum_{k=0}^{K-1} A[i][j][k] \cdot B[k][l]$$

Input

• Tensor $A$ of size $N \times M \times K$
• Matrix $B$ of size $K \times L$

Output

• Tensor $C$ of size $N \times M \times L$

Notes:

• All tensors $\text{A}$, $\text{B}$, and $\text{C}$ are stored in row-major order
• This problem is adapted from KernelBench