From aa470e3f90945142e1bc5833d5cc1df1846fd0be Mon Sep 17 00:00:00 2001
From: Kaushik Kulkarni <kaushikcfd@gmail.com>
Date: Fri, 16 Dec 2022 21:17:33 +0530
Subject: [PATCH] Test for transforms.reduction

---
 test/test_transform.py | 42 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 42 insertions(+)

diff --git a/test/test_transform.py b/test/test_transform.py
index 61b16d064..84e303b0c 100644
--- a/test/test_transform.py
+++ b/test/test_transform.py
@@ -1519,6 +1519,48 @@ def test_prefetch_to_same_temp_var(ctx_factory):
     lp.auto_test_vs_ref(ref_tunit, ctx, t_unit)
 
 
+def test_sum_redn_algebraic_transforms(ctx_factory):
+    from pymbolic import variables
+    from loopy.symbolic import Reduction
+
+    t_unit = lp.make_kernel(
+        "{[e,i,j,x,r]: 0<=e<N_e and 0<=i,j<35 and 0<=x,r<3}",
+        """
+        y[i] = sum([r,j], J[x, r, e]*D[r,i,j]*u[e,j])
+        """,
+        [lp.GlobalArg("J,D,u", dtype=np.float64, shape=lp.auto),
+         ...],
+    )
+    knl = t_unit.default_entrypoint
+
+    knl = lp.split_reduction_inward(knl, "j")
+    knl = lp.hoist_invariant_multiplicative_terms_in_sum_reduction(
+        knl,
+        reduction_inames="j"
+    )
+    knl = lp.extract_multiplicative_terms_in_sum_reduction_as_subst(
+        knl,
+        within=None,
+        subst_name="grad_without_jacobi_subst",
+        arguments=variables("r i e"),
+        terms_filter=lambda x: isinstance(x, Reduction)
+    )
+
+    transformed_t_unit = t_unit.with_kernel(knl)
+    transformed_t_unit = lp.precompute(
+        transformed_t_unit,
+        "grad_without_jacobi_subst",
+        sweep_inames=["r", "i"],
+        precompute_outer_inames=frozenset({"e"}),
+        temporary_address_space=lp.AddressSpace.PRIVATE)
+
+    x1 = lp.get_op_map(t_unit, subgroup_size=1).eval_and_sum({"N_e": 1})
+    x2 = lp.get_op_map(transformed_t_unit, subgroup_size=1).eval_and_sum({"N_e": 1})
+
+    assert x1 == 33075
+    assert x2 == 7980  # i.e. demonstrates a 4.14x fewer flops
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])