@triton.jit
def triton_add_mul_kernel(x_ptr, a_ptr, b_ptr, c, stride_x0, stride_x1,
    stride_a0, stride_a1, stride_b0, stride_b1, BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr, M: tl.constexpr, N: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    offsets_m = tl.arange(0, BLOCK_SIZE_M) + pid_m * BLOCK_SIZE_M
    offsets_n = tl.arange(0, BLOCK_SIZE_N) + pid_n * BLOCK_SIZE_N
    x = tl.load(x_ptr + offsets_m[:, None] * stride_x0 + offsets_n[None, :] *
        stride_x1, mask=(offsets_m[:, None] < M) & (offsets_n[None, :] < N))
    a = tl.load(a_ptr + offsets_m[:, None] * stride_a0 + offsets_n[None, :] *
        stride_a1, mask=(offsets_m[:, None] < M) & (offsets_n[None, :] < N))
    b = tl.load(b_ptr + offsets_m[:, None] * stride_b0 + offsets_n[None, :] *
        stride_b1, mask=(offsets_m[:, None] < M) & (offsets_n[None, :] < N))
    x_dtype = x.dtype
    x = (x.to(tl.float32) + a.to(tl.float32) * b.to(tl.float32) * c).to(x_dtype
        )
    tl.store(x_ptr + offsets_m[:, None] * stride_x0 + offsets_n[None, :] *
        stride_x1, x, mask=(offsets_m[:, None] < M) & (offsets_n[None, :] < N))
