test

fix test

Author: yf225

test/test_indexing.py

@@ -1604,6 +1604,130 @@ def load_store_kernel(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
         self.assertEqual(code3, code4)
         self.assertExpectedJournal(code4)
 
+    def test_indirect_indexing_2d(self):
+        @helion.kernel()
+        def test(
+            col: torch.Tensor,   # [M, K] int64
+            val: torch.Tensor,   # [M, K] fp32
+            B: torch.Tensor,     # [K, N] fp32
+        ) -> torch.Tensor:       # [M, N] fp32
+            M, K = col.shape
+            _, N = B.shape
+            out_dtype = torch.promote_types(val.dtype, B.dtype)
+            C = torch.empty((M, N), dtype=out_dtype, device=B.device)
+            B_flat = B.reshape(-1)  # [K*N]
+
+            for tile_m, tile_n in hl.tile([M, N]):
+                # [tile_m, tile_n]
+                acc = hl.zeros([tile_m, tile_n], dtype=torch.float32)
+
+                for tile_k in hl.tile(K):
+                    # [tile_m, tile_k]
+                    cols_2d = col[tile_m, tile_k]
+                    # [tile_m, tile_k, tile_n]
+                    B_slice = hl.load(
+                        B_flat,
+                        [(cols_2d * N)[:, :, None] + tile_n.index[None, None, :]]
+                    )
+                    # [tile_m, tile_k]
+                    vals_2d = val[tile_m, tile_k]
+                    # [tile_m, tile_k, tile_n]
+                    contrib = vals_2d[:, :, None] * B_slice
+                    # [tile_m, tile_n]
+                    contrib = contrib.sum(dim=1)
+                    # [tile_m, tile_n]
+                    acc = acc + contrib
+
+                C[tile_m, tile_n] = acc.to(out_dtype)
+
+            return C
+
+        M, K, N = 32, 16, 24
+        col = torch.randint(0, K, (M, K), device=DEVICE, dtype=torch.int64)
+        val = torch.rand((M, K), device=DEVICE, dtype=torch.float32)
+        B = torch.rand((K, N), device=DEVICE, dtype=torch.float32)
+
+        code, result = code_and_output(
+            test,
+            (col, val, B),
+            block_size=[8, 8, 4],
+        )
+
+        # For each output position (i,j), compute sum over k: val[i,k] * B[col[i,k], j]
+        expected = torch.zeros((M, N), device=DEVICE, dtype=torch.float32)
+        for i in range(M):
+            for j in range(N):
+                for k in range(K):
+                    expected[i, j] += val[i, k] * B[col[i, k], j]
+        
+        torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-5)
+        self.assertExpectedJournal(code)
+
+    def test_indirect_indexing_3d(self):
+        @helion.kernel()
+        def test(
+            col: torch.Tensor,   # [M, N, K] int64 - indices for first dimension of B
+            val: torch.Tensor,   # [M, N, K] fp32 - values to multiply
+            B: torch.Tensor,     # [K, P, Q] fp32 - tensor to index into
+        ) -> torch.Tensor:       # [M, N, P, Q] fp32
+            M, N, K = col.shape
+            _, P, Q = B.shape
+            out_dtype = torch.promote_types(val.dtype, B.dtype)
+            C = torch.empty((M, N, P, Q), dtype=out_dtype, device=B.device)
+            
+            for tile_m, tile_n, tile_p, tile_q in hl.tile([M, N, P, Q]):
+                # [tile_m, tile_n, tile_p, tile_q]
+                acc = hl.zeros([tile_m, tile_n, tile_p, tile_q], dtype=torch.float32)
+                
+                for tile_k in hl.tile(K):
+                    # [tile_m, tile_n, tile_k]
+                    cols_3d = col[tile_m, tile_n, tile_k]
+                    
+                    # [tile_m, tile_n, tile_k, tile_p, tile_q]
+                    # Direct indexing into B using gather
+                    B_slice = B[
+                        cols_3d[:, :, :, None, None],
+                        tile_p.index[None, None, :, None],
+                        tile_q.index[None, None, None, :],
+                    ]
+                    
+                    # [tile_m, tile_n, tile_k]
+                    vals_3d = val[tile_m, tile_n, tile_k]
+                    
+                    # [tile_m, tile_n, tile_k, tile_p, tile_q]
+                    contrib = vals_3d[:, :, :, None, None] * B_slice
+                    
+                    # [tile_m, tile_n, tile_p, tile_q] - sum over k dimension
+                    contrib = contrib.sum(dim=2)
+                    
+                    # [tile_m, tile_n, tile_p, tile_q]
+                    acc = acc + contrib
+                    
+                C[tile_m, tile_n, tile_p, tile_q] = acc.to(out_dtype)
+            return C
+        
+        M, N, K, P, Q = 16, 12, 8, 10, 14
+        col = torch.randint(0, K, (M, N, K), device=DEVICE, dtype=torch.int64)
+        val = torch.rand((M, N, K), device=DEVICE, dtype=torch.float32)
+        B = torch.rand((K, P, Q), device=DEVICE, dtype=torch.float32)
+        
+        code, result = code_and_output(
+            test,
+            (col, val, B),
+            block_size=[4, 4, 4, 4, 4],  # 5D tiling for M, N, P, Q, K
+        )
+        
+        # For each output position (i,j,p,q), compute sum over k: val[i,j,k] * B[col[i,j,k], p, q]
+        expected = torch.zeros((M, N, P, Q), device=DEVICE, dtype=torch.float32)
+        for i in range(M):
+            for j in range(N):
+                for p in range(P):
+                    for q in range(Q):
+                        for k in range(K):
+                            expected[i, j, p, q] += val[i, j, k] * B[col[i, j, k], p, q]
+        
+        torch.testing.assert_close(result, expected, rtol=1e-5, atol=1e-5)
+        self.assertExpectedJournal(code)
 
 if __name__ == "__main__":
     unittest.main()