test dgemm transformation with rectangular tile prefetch

test/test_loopy.py

@@ -3483,6 +3483,52 @@ def test_einsum_parsing(ctx_factory):
                         parameters={"Ni": 10, "Nj": 10, "Nk": 10})
 
 
+def test_dgemm_with_rectangular_tile_prefetch(ctx_factory):
+    # See <https://github.com/inducer/loopy/issues/724>
+    t_unit = lp.make_kernel(
+        "{[i,j,k]: 0<=i,j<72 and 0<=k<32}",
+        """
+        C[i,j] = sum(k, A[i,k] * B[k,j])
+        """,
+        [lp.GlobalArg("A,B", dtype=np.float64, shape=lp.auto),
+         ...],
+    )
+    ref_t_unit = t_unit
+
+    tx = 8
+    ty = 23
+    tk = 11
+
+    t_unit = lp.split_iname(t_unit, "i", tx, inner_tag="l.0", outer_tag="g.0")
+    t_unit = lp.split_iname(t_unit, "j", ty, inner_tag="l.1", outer_tag="g.1")
+    t_unit = lp.split_iname(t_unit, "k", tk)
+    t_unit = lp.add_prefetch(
+        t_unit, "A",
+        sweep_inames=["i_inner", "k_inner"],
+        temporary_address_space=lp.AddressSpace.LOCAL,
+        fetch_outer_inames=frozenset({"i_outer", "j_outer", "k_outer"}),
+        dim_arg_names=["iprftch_A", "kprftch_A"],
+        default_tag=None,
+    )
+
+    t_unit = lp.add_prefetch(
+        t_unit, "B",
+        sweep_inames=["k_inner", "j_inner"],
+        temporary_address_space=lp.AddressSpace.LOCAL,
+        fetch_outer_inames=frozenset({"i_outer", "j_outer", "k_outer"}),
+        dim_arg_names=["kprftch_B", "jprftch_B"],
+        default_tag=None,
+    )
+
+    t_unit = lp.split_iname(t_unit, "kprftch_A", tx, inner_tag="l.0")
+    t_unit = lp.split_iname(t_unit, "iprftch_A", ty, inner_tag="l.1")
+    t_unit = lp.split_iname(t_unit, "jprftch_B", tx, inner_tag="l.0")
+    t_unit = lp.split_iname(t_unit, "kprftch_B", ty, inner_tag="l.1")
+
+    ctx = cl.create_some_context()
+    lp.auto_test_vs_ref(ref_t_unit, ctx, t_unit)
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])