matmul bwd updated, test updated, fix issues

Author: tianrengao

examples/matmul.py

@@ -112,6 +112,74 @@ def matmul_bwd(
     return grad_mat1, grad_mat2
 
 
+@helion.kernel
+def addmm_bwd(
+    grad_out: Tensor,  # [m, n] gradient w.r.t output
+    input: Tensor,  # [m, n] or broadcastable bias tensor
+    mat1: Tensor,  # [m, k] first matrix
+    mat2: Tensor,  # [k, n] second matrix
+    alpha: float = 1.0,  # scalar multiplier for matmul
+    beta: float = 1.0,  # scalar multiplier for bias
+) -> tuple[Tensor, Tensor, Tensor]:
+    """
+    Backward pass for addmm operation following Triton reference pattern.
+
+    Forward: output = beta * input + alpha * (mat1 @ mat2)
+
+    Based on the Triton kernel analysis:
+    - grad_input = beta * grad_out (with proper reduction for broadcasting)
+    - grad_mat1 = alpha * (grad_out @ mat2.T)
+    - grad_mat2 = alpha * (mat1.T @ grad_out)
+
+    Args:
+        grad_out: Gradient w.r.t output [m, n]
+        input: Bias tensor [m, n] (or broadcastable)
+        mat1: First matrix [m, k]
+        mat2: Second matrix [k, n]
+        alpha: Scalar multiplier for matmul
+        beta: Scalar multiplier for bias
+
+    Returns:
+        tuple[Tensor, Tensor, Tensor]: (grad_input, grad_mat1, grad_mat2)
+    """
+    # Get all dimensions first
+    m, n = grad_out.size()
+    m2, k = mat1.size()
+    k2, n2 = mat2.size()
+
+    # All assertions at the top
+    assert m == m2 and n == n2 and k == k2, "Size mismatch in addmm backward"
+
+    # Declare ALL output tensors at the top before any loops
+    grad_input = torch.empty_like(input)
+    grad_mat1 = torch.empty_like(mat1)
+    grad_mat2 = torch.empty_like(mat2)
+
+    # Handle grad_input = beta * grad_out (assuming same shape for now)
+    for tile_m3, tile_n3 in hl.tile([m, n]):
+        grad_input[tile_m3, tile_n3] = beta * grad_out[tile_m3, tile_n3]
+
+    # First loop block: compute grad_mat1 = alpha * (grad_out @ mat2.T)
+    for tile_m1, tile_k1 in hl.tile([m, k]):
+        acc1 = hl.zeros([tile_m1, tile_k1], dtype=torch.float32)
+        for tile_n1 in hl.tile(n):
+            acc1 = torch.addmm(
+                acc1, grad_out[tile_m1, tile_n1], mat2[tile_k1, tile_n1].T
+            )
+        grad_mat1[tile_m1, tile_k1] = (alpha * acc1).to(mat1.dtype)
+
+    # Second loop block: compute grad_mat2 = alpha * (mat1.T @ grad_out)
+    for tile_k2, tile_n2 in hl.tile([k, n]):
+        acc2 = hl.zeros([tile_k2, tile_n2], dtype=torch.float32)
+        for tile_m2 in hl.tile(m):
+            acc2 = torch.addmm(
+                acc2, mat1[tile_m2, tile_k2].T, grad_out[tile_m2, tile_n2]
+            )
+        grad_mat2[tile_k2, tile_n2] = (alpha * acc2).to(mat2.dtype)
+
+    return grad_input, grad_mat1, grad_mat2
+
+
 # %%
 class MatMulFunction(torch.autograd.Function):
     @staticmethod
@@ -141,6 +209,45 @@ def matmul_autograd(mat1: Tensor, mat2: Tensor) -> Tensor:
     return MatMulFunction.apply(mat1, mat2)  # type: ignore[no-any-return]
 
 
+class AddMMFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx: Any,
+        input: Tensor,
+        mat1: Tensor,
+        mat2: Tensor,
+        alpha: float = 1.0,
+        beta: float = 1.0,
+    ) -> Tensor:
+        """Forward pass for addmm operation."""
+        result = torch.addmm(input, mat1, mat2, alpha=alpha, beta=beta)
+        ctx.save_for_backward(input, mat1, mat2)
+        ctx.alpha = alpha
+        ctx.beta = beta
+        return result
+
+    @staticmethod
+    def backward(
+        ctx: Any,
+        grad_out: Tensor,
+    ) -> tuple[Tensor | None, Tensor | None, Tensor | None, None, None]:
+        """Backward pass for addmm operation."""
+        input, mat1, mat2 = ctx.saved_tensors
+        alpha = ctx.alpha
+        beta = ctx.beta
+        grad_input, grad_mat1, grad_mat2 = addmm_bwd(
+            grad_out, input, mat1, mat2, alpha, beta
+        )
+        return grad_input, grad_mat1, grad_mat2, None, None
+
+
+def addmm_autograd(
+    input: Tensor, mat1: Tensor, mat2: Tensor, alpha: float = 1.0, beta: float = 1.0
+) -> Tensor:
+    """AddMM operation with forward + backward support."""
+    return AddMMFunction.apply(input, mat1, mat2, alpha, beta)  # type: ignore[no-any-return]
+
+
 # %%
 def autotune(m: int, k: int, n: int) -> None:
     """
@@ -230,6 +337,47 @@ def baseline_wrapper(x: Tensor, y: Tensor) -> Tensor:
         bwd=True,
     )
 
+    # Test addmm forward + backward pass
+    print("\n\n=== AddMM Forward + Backward Pass Test ===")
+    input_grad = torch.randn(
+        [m, n], device="cuda", dtype=torch.float16, requires_grad=True
+    )
+    mat1_grad = torch.randn(
+        [m, k], device="cuda", dtype=torch.float16, requires_grad=True
+    )
+    mat2_grad = torch.randn(
+        [k, n], device="cuda", dtype=torch.float16, requires_grad=True
+    )
+
+    # Use lambda to handle the keyword argument format for torch.addmm
+    run_example(
+        addmm_autograd,
+        lambda input, mat1, mat2, alpha, beta: torch.addmm(
+            input, mat1, mat2, alpha=alpha, beta=beta
+        ),
+        (input_grad, mat1_grad, mat2_grad, 1.0, 1.0),
+        kernel_name="helion_addmm_autograd",
+        baseline_name="torch",
+        rtol=1e-2,
+        atol=1e-2,
+        bwd=True,
+    )
+
+    # Test addmm forward + backward with different alpha/beta values
+    print("\n\n=== AddMM Forward + Backward Test (Alpha=2.0, Beta=0.5) ===")
+    run_example(
+        addmm_autograd,
+        lambda input, mat1, mat2, alpha, beta: torch.addmm(
+            input, mat1, mat2, alpha=alpha, beta=beta
+        ),
+        (input_grad, mat1_grad, mat2_grad, 2.0, 0.5),
+        kernel_name="helion_addmm_autograd_scaled",
+        baseline_name="torch",
+        rtol=1e-2,
+        atol=1e-2,
+        bwd=True,
+    )
+
 
 # %%
 def matmul_tritonbench(

test/test_examples.expected

@@ -2818,8 +2818,6 @@ def matmul_bwd(grad_out: Tensor, mat1: Tensor, mat2: Tensor, *, _launcher=_defau
     - grad_A = grad_C @ B.T
     - grad_B = A.T @ grad_C
 
-    Key: Declare ALL variables at the top before ANY loops to avoid TopLevelStatementBetweenLoops error
-
     Args:
         grad_out: Gradient w.r.t output [m, n]
         mat1: First matrix [m, k]

test/test_examples.py

@@ -50,19 +50,28 @@ def test_matmul(self):
 
     def test_matmul_bwd(self):
         """Test backward pass for matmul computation."""
+        # Create tensors with requires_grad=True like rms_norm_bwd test
+        mat1 = torch.randn(
+            [128, 128], device=DEVICE, dtype=torch.float32, requires_grad=True
+        )
+        mat2 = torch.randn(
+            [128, 128], device=DEVICE, dtype=torch.float32, requires_grad=True
+        )
         grad_out = torch.randn([128, 128], device=DEVICE, dtype=torch.float32)
-        mat1 = torch.randn([128, 128], device=DEVICE, dtype=torch.float32)
-        mat2 = torch.randn([128, 128], device=DEVICE, dtype=torch.float32)
 
-        args = (grad_out, mat1, mat2)
+        # Compute expected gradients with PyTorch
+        mat1_torch = mat1.detach().clone().requires_grad_(True)
+        mat2_torch = mat2.detach().clone().requires_grad_(True)
+        result_torch = torch.matmul(mat1_torch, mat2_torch)
+        result_torch.backward(grad_out)
 
-        expected = (grad_out @ mat2.T, mat1.T @ grad_out)
+        args = (grad_out, mat1, mat2)
 
         self.assertExpectedJournal(
             check_example(
                 "matmul",
                 args,
-                expected,
+                (mat1_torch.grad, mat2_torch.grad),  # Expected: (grad_mat1, grad_mat2)
                 fn_name="matmul_bwd",
                 block_sizes=[
                     16,