Implement fused matrix multiplication and ReLU in Triton. Input matrices A and B have shapes (M,K) and (K,N), and the output Y has shape (M,N):
Ym,n=max(0,k=0∑K−1float32(Am,k)float32(Bk,n)).Here, M and N are output dimensions, K is the shared inner dimension, and m, n, and k index those dimensions. Float32 denotes the arithmetic precision, and the maximum applies ReLU. Complete only the empty kernel body; the supplied solve launcher is implemented. The pointer arguments identify the three buffers, the size arguments give M, N, and K, the stride arguments give each buffer's element strides, and BLOCK_M, BLOCK_N, and BLOCK_K are compile-time tile sizes. Return no value; write Y into the supplied output tensor using its preallocated dtype.
Input: a = [[1, -2], [3, 1]], b = [[2, -1, 0], [1, 2, -3]], dtype = float16, output_dtype = float16
Output: [[0, 0, 6], [7, 0, 0]]
Explanation: ReLU replaces the negative entries of the matrix product with zero.
Input: a = [[1, 2, 3], [0, -1, 2]], b = [[1, 0], [0, 1], [1, -1]], dtype = bfloat16, output_dtype = bfloat16
Output: [[4, 0], [2, 0]]
Input: a = [[-1, 2], [2, -1]], b = [[1, -3], [4, 2]], dtype = float16, output_dtype = float32
Output: [[7, 7], [0, 0]]
tl.program_id can address each dimension of a program grid.
tl.dot supports float32 accumulation of tile products.
tl.load and tl.store support masks for partial tiles.
Accepts: array
Accepts: array
Accepts: string
Accepts: string
Submit a correct solution to see your kernel performance.
Implement fused matrix multiplication and ReLU in Triton. Input matrices A and B have shapes (M,K) and (K,N), and the output Y has shape (M,N):
Ym,n=max(0,k=0∑K−1float32(Am,k)float32(Bk,n)).Here, M and N are output dimensions, K is the shared inner dimension, and m, n, and k index those dimensions. Float32 denotes the arithmetic precision, and the maximum applies ReLU. Complete only the empty kernel body; the supplied solve launcher is implemented. The pointer arguments identify the three buffers, the size arguments give M, N, and K, the stride arguments give each buffer's element strides, and BLOCK_M, BLOCK_N, and BLOCK_K are compile-time tile sizes. Return no value; write Y into the supplied output tensor using its preallocated dtype.
Input: a = [[1, -2], [3, 1]], b = [[2, -1, 0], [1, 2, -3]], dtype = float16, output_dtype = float16
Output: [[0, 0, 6], [7, 0, 0]]
Explanation: ReLU replaces the negative entries of the matrix product with zero.
Input: a = [[1, 2, 3], [0, -1, 2]], b = [[1, 0], [0, 1], [1, -1]], dtype = bfloat16, output_dtype = bfloat16
Output: [[4, 0], [2, 0]]
Input: a = [[-1, 2], [2, -1]], b = [[1, -3], [4, 2]], dtype = float16, output_dtype = float32
Output: [[7, 7], [0, 0]]
tl.program_id can address each dimension of a program grid.
tl.dot supports float32 accumulation of tile products.
tl.load and tl.store support masks for partial tiles.
Accepts: array
Accepts: array
Accepts: string
Accepts: string
Submit a correct solution to see your kernel performance.