Repository: exo-lang/exo

Issue:
Add div support to DoNormalize
e.g., `(wii + 8 * wio) / 4 / 2` should be simplified

Discussion:
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/exo-lang/exo/pull/473?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) Report
> Merging [#473](https://app.codecov.io/gh/exo-lang/exo/pull/473?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) (44cac1b) into [master](https://app.codecov.io/gh/exo-lang/exo/commit/20805c55eb9a977f9d9ac62e95043c3b411d91b9?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) (20805c5) will **increase** coverage by `0.01%`.
> The diff coverage is `100.00%`.

:exclamation: Your organization is not using the GitHub App Integration. As a result you may experience degraded service beginning May 15th. Please [install the Github App Integration](https://github.com/apps/codecov) for your organization. [Read more](https://about.codecov.io/blog/codecov-is-updating-its-github-integration/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang).

```diff
@@            Coverage Diff             @@
##           master     #473      +/-   ##
==========================================
+ Coverage   86.38%   86.39%   +0.01%     
==========================================
  Files          78       78              
  Lines       17910    17926      +16     
==========================================
+ Hits        15471    15487      +16     
  Misses       2439     2439              
```


| [Impacted Files](https://app.codecov.io/gh/exo-lang/exo/pull/473?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang) | Coverage Δ | |
|---|---|---|
| [src/exo/LoopIR\_scheduling.py](https://app.codecov.io/gh/exo-lang/exo/pull/473?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-c3JjL2V4by9Mb29wSVJfc2NoZWR1bGluZy5weQ==) | `90.61% <100.00%> (+0.06%)` | :arrow_up: |
| [tests/test\_schedules.py](https://app.codecov.io/gh/exo-lang/exo/pull/473?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=exo-lang#diff-dGVzdHMvdGVzdF9zY2hlZHVsZXMucHk=) | `99.36% <100.00%> (+<0.01%)` | :arrow_up: |



Patch:
diff --git a/src/exo/LoopIR_scheduling.py b/src/exo/LoopIR_scheduling.py
index 1b5d585..cb3ae54 100644
--- a/src/exo/LoopIR_scheduling.py
+++ b/src/exo/LoopIR_scheduling.py
@@ -2657,6 +2657,36 @@ class _DoNormalize(Cursor_Rewrite):
             new_lhs = generate_loopIR(e.lhs, constant, normalization_list)
             return LoopIR.BinOp("/", new_lhs, e.rhs, e.type, e.srcinfo)
 
+        def division_denominator_simplification(e):
+            assert e.op == "/"
+
+            def has_nested_const_denominator(expr):
+                # (n / c1) / c2
+                if expr.op == "/" and isinstance(
+                    expr.rhs, LoopIR.Const
+                ):  # (something / c2)
+                    if (
+                        isinstance(expr.lhs, LoopIR.BinOp) and expr.lhs.op == "/"
+                    ):  # (n / c1)
+                        if isinstance(expr.lhs.rhs, LoopIR.Const):
+                            return True
+
+                return False
+
+            new_e = e
+            # call division_denominator_simplification recursively
+            while has_nested_const_denominator(new_e):
+                new_e = new_e.update(
+                    lhs=new_e.lhs.lhs,
+                    rhs=LoopIR.Const(
+                        new_e.lhs.rhs.val * new_e.rhs.val,
+                        new_e.lhs.type,
+                        new_e.lhs.srcinfo,
+                    ),
+                )
+
+            return new_e
+
         def division_simplification_and_try_spliting_denominator(e):
             def still_division(e):
                 return isinstance(e, LoopIR.BinOp) and e.op == "/"
@@ -2756,13 +2786,16 @@ class _DoNormalize(Cursor_Rewrite):
             new_rhs = self.index_start(e.rhs)
             e = e.update(lhs=new_lhs, rhs=new_rhs)
 
-        if isinstance(e, LoopIR.BinOp) and e.op in ("/", "%"):
+        if isinstance(e, LoopIR.BinOp) and e.op == "/":
             assert isinstance(e.rhs, LoopIR.Const)
             if self.has_div_mod_config(e.lhs):
-                return e
+                e = division_denominator_simplification(e)
+            return division_simplification_and_try_spliting_denominator(e)
 
-            if e.op == "/":
-                return division_simplification_and_try_spliting_denominator(e)
+        if isinstance(e, LoopIR.BinOp) and e.op == "%":
+            assert isinstance(e.rhs, LoopIR.Const)
+            if self.has_div_mod_config(e.lhs):
+                return e
 
             return modulo_simplification(e)
 


Test patch:
diff --git a/tests/golden/test_apps/test_neon_sgemm.txt b/tests/golden/test_apps/test_neon_sgemm.txt
index 23e51dc..384c16a 100644
--- a/tests/golden/test_apps/test_neon_sgemm.txt
+++ b/tests/golden/test_apps/test_neon_sgemm.txt
@@ -157,13 +157,13 @@ EXO_ASSUME(K >= 1);
 float *Atile = malloc(64 * 64 * sizeof(*Atile));
 float *Btile = malloc(64 * 64 * sizeof(*Btile));
 for (int_fast32_t ko = 0; ko < ((K) / (64)); ko++) {
-  for (int_fast32_t io = 0; io < ((((M) / (4))) / (16)); io++) {
+  for (int_fast32_t io = 0; io < ((M) / (64)); io++) {
     for (int_fast32_t i0 = 0; i0 < 64; i0++) {
       for (int_fast32_t i1 = 0; i1 < 64; i1++) {
         Atile[i0 * 64 + i1] = A[(i0 + 64 * io) * K + i1 + 64 * ko];
       }
     }
-    for (int_fast32_t jo = 0; jo < ((((N) / (16))) / (4)); jo++) {
+    for (int_fast32_t jo = 0; jo < ((N) / (64)); jo++) {
       for (int_fast32_t i0 = 0; i0 < 64; i0++) {
         for (int_fast32_t i1 = 0; i1 < 64; i1++) {
           Btile[i0 * 64 + i1] = B[(i0 + 64 * ko) * N + i1 + 64 * jo];
@@ -180,16 +180,16 @@ for (int_fast32_t ko = 0; ko < ((K) / (64)); ko++) {
 free(Btile);
 free(Atile);
 for (int_fast32_t ko = 0; ko < ((K) / (64)); ko++) {
-  for (int_fast32_t io = 0; io < ((((M) / (4))) / (16)); io++) {
+  for (int_fast32_t io = 0; io < ((M) / (64)); io++) {
     for (int_fast32_t jm = 0; jm < ((N) / (16)) % 4; jm++) {
       for (int_fast32_t im = 0; im < 16; im++) {
-        neon_microkernel(ctxt,64,(struct exo_win_2f32c){ &A[(4 * im + 64 * io) * K + 64 * ko], { K, 1 } },(struct exo_win_2f32c){ &B[(64 * ko) * N + 16 * (jm + ((N / 16) / 4) * 4)], { N, 1 } },(struct exo_win_2f32){ &C[(4 * im + 64 * io) * N + 16 * (jm + ((N / 16) / 4) * 4)], { N, 1 } });
+        neon_microkernel(ctxt,64,(struct exo_win_2f32c){ &A[(4 * im + 64 * io) * K + 64 * ko], { K, 1 } },(struct exo_win_2f32c){ &B[(64 * ko) * N + 16 * (jm + (N / 64) * 4)], { N, 1 } },(struct exo_win_2f32){ &C[(4 * im + 64 * io) * N + 16 * (jm + (N / 64) * 4)], { N, 1 } });
       }
     }
   }
   for (int_fast32_t jo = 0; jo < ((N) / (16)); jo++) {
     for (int_fast32_t im = 0; im < ((M) / (4)) % 16; im++) {
-      neon_microkernel(ctxt,64,(struct exo_win_2f32c){ &A[(4 * (im + ((M / 4) / 16) * 16)) * K + 64 * ko], { K, 1 } },(struct exo_win_2f32c){ &B[(64 * ko) * N + 16 * jo], { N, 1 } },(struct exo_win_2f32){ &C[(4 * (im + ((M / 4) / 16) * 16)) * N + 16 * jo], { N, 1 } });
+      neon_microkernel(ctxt,64,(struct exo_win_2f32c){ &A[(4 * (im + (M / 64) * 16)) * K + 64 * ko], { K, 1 } },(struct exo_win_2f32c){ &B[(64 * ko) * N + 16 * jo], { N, 1 } },(struct exo_win_2f32){ &C[(4 * (im + (M / 64) * 16)) * N + 16 * jo], { N, 1 } });
     }
   }
 }
diff --git a/tests/golden/test_schedules/test_simplify_nested_div.txt b/tests/golden/test_schedules/test_simplify_nested_div.txt
new file mode 100644
index 0000000..de1452f
--- /dev/null
+++ b/tests/golden/test_schedules/test_simplify_nested_div.txt
@@ -0,0 +1,4 @@
+def foo(n: size):
+    x: f32 @ DRAM
+    for i in seq(0, n / 24):
+        x = 0.0
\ No newline at end of file
diff --git a/tests/golden/test_schedules/test_simplify_nested_div_2.txt b/tests/golden/test_schedules/test_simplify_nested_div_2.txt
new file mode 100644
index 0000000..d1ae304
--- /dev/null
+++ b/tests/golden/test_schedules/test_simplify_nested_div_2.txt
@@ -0,0 +1,4 @@
+def foo(n: size):
+    x: f32 @ DRAM
+    for i in seq(0, n / 4):
+        x = 0.0
\ No newline at end of file
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
index ba91baf..fd1411e 100644
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -319,6 +319,26 @@ def test_simplify_loop_bounds(golden):
     assert str(simplify(foo)) == golden
 
 
+def test_simplify_nested_div(golden):
+    @proc
+    def foo(n: size):
+        x: f32
+        for i in seq(0, (n / 4) / 3 / 2):
+            x = 0.0
+
+    assert str(simplify(foo)) == golden
+
+
+def test_simplify_nested_div_2(golden):
+    @proc
+    def foo(n: size):
+        x: f32
+        for i in seq(0, ((3 * n) / 4) / 3):
+            x = 0.0
+
+    assert str(simplify(foo)) == golden
+
+
 def test_pattern_match():
     @proc
     def foo(N1: size, M1: size, K1: size, N2: size, M2: size, K2: size):