From d50730415071cdb7c3f0c4c7cca141bd6f390ef7 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Sat, 23 May 2015 15:08:00 -0400
Subject: [PATCH] Add fallback to counting routines, minor cleanups

---
 loopy/statistics.py     | 41 ++++++++++++++++++++++++++-----
 test/test_statistics.py | 53 +++++++++++++++++++++++++++++++++--------
 2 files changed, 78 insertions(+), 16 deletions(-)

diff --git a/loopy/statistics.py b/loopy/statistics.py
index a1e8c7648..10e5c4c8d 100755
--- a/loopy/statistics.py
+++ b/loopy/statistics.py
@@ -1,6 +1,4 @@
-from __future__ import division
-from __future__ import absolute_import
-import six
+from __future__ import division, absolute_import
 
 __copyright__ = "Copyright (C) 2015 James Stevens"
 
@@ -24,6 +22,8 @@ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 THE SOFTWARE.
 """
 
+import six  # noqa
+
 import loopy as lp
 import warnings
 from islpy import dim_type
@@ -218,6 +218,36 @@ class SubscriptCounter(CombineMapper):
         return 0
 
 
+def count(kernel, bset):
+    try:
+        return bset.card()
+    except AttributeError:
+        pass
+
+    if not bset.is_box():
+        from loopy.diagnostic import warn
+        warn(kernel, "count_overestimate",
+                "Barvinok wrappers are not installed. "
+                "Counting routines may overestimate the "
+                "number of integer points in your loop "
+                "domain.")
+
+    result = None
+
+    for i in range(bset.dim(isl.dim_type.set)):
+        dmax = bset.dim_max(i)
+        dmin = bset.dim_min(i)
+
+        length = isl.PwQPolynomial.from_pw_aff(dmax - dmin + 1)
+
+        if result is None:
+            result = length
+        else:
+            result = result * length
+
+    return result
+
+
 # to evaluate poly: poly.eval_with_dict(dictionary)
 def get_op_poly(knl):
     from loopy.preprocess import preprocess_kernel, infer_unknown_types
@@ -233,7 +263,7 @@ def get_op_poly(knl):
         inames_domain = knl.get_inames_domain(insn_inames)
         domain = (inames_domain.project_out_except(insn_inames, [dim_type.set]))
         ops = op_counter(insn.expression)
-        op_poly = op_poly + ops*domain.card()
+        op_poly = op_poly + ops*count(knl, domain)
     return op_poly
 
 
@@ -245,6 +275,5 @@ def get_DRAM_access_poly(knl):  # for now just counting subscripts
         insn_inames = knl.insn_inames(insn)
         inames_domain = knl.get_inames_domain(insn_inames)
         domain = (inames_domain.project_out_except(insn_inames, [dim_type.set]))
-        poly += subscript_counter(insn.expression) * domain.card()
+        poly += subscript_counter(insn.expression) * count(knl, domain)
     return poly
-
diff --git a/test/test_statistics.py b/test/test_statistics.py
index e68700833..f3919db2c 100644
--- a/test/test_statistics.py
+++ b/test/test_statistics.py
@@ -1,4 +1,4 @@
-from __future__ import division
+from __future__ import division, print_function
 
 __copyright__ = "Copyright (C) 2015 James Stevens"
 
@@ -23,14 +23,14 @@ THE SOFTWARE.
 """
 
 import sys
-from pyopencl.tools import (
+from pyopencl.tools import (  # noqa
         pytest_generate_tests_for_pyopencl
         as pytest_generate_tests)
 from loopy.statistics import *  # noqa
 import numpy as np
 
 
-def test_op_counter_basic(ctx_factory):
+def test_op_counter_basic():
 
     knl = lp.make_kernel(
             "[n,m,l] -> {[i,k,j]: 0<=i<n and 0<=k<m and 0<=j<l}",
@@ -56,7 +56,7 @@ def test_op_counter_basic(ctx_factory):
     assert i32 == n*m
 
 
-def test_op_counter_reduction(ctx_factory):
+def test_op_counter_reduction():
 
     knl = lp.make_kernel(
             "{[i,k,j]: 0<=i<n and 0<=k<m and 0<=j<l}",
@@ -74,10 +74,10 @@ def test_op_counter_reduction(ctx_factory):
     assert f32 == 2*n*m*l
 
 
-def test_op_counter_logic(ctx_factory):
+def test_op_counter_logic():
 
     knl = lp.make_kernel(
-            "[n,m,l] -> {[i,k,j]: 0<=i<n and 0<=k<m and 0<=j<l}",
+            "{[i,k,j]: 0<=i<n and 0<=k<m and 0<=j<l}",
             [
                 """
                 e[i,k] = if(not(k<l-2) and k>6 or k/2==l, g[i,k]*2, g[i,k]+h[i,k]/2)
@@ -98,10 +98,10 @@ def test_op_counter_logic(ctx_factory):
     assert i32 == n*m
 
 
-def test_op_counter_specialops(ctx_factory):
+def test_op_counter_specialops():
 
     knl = lp.make_kernel(
-            "[n,m,l] -> {[i,k,j]: 0<=i<n and 0<=k<m and 0<=j<l}",
+            "{[i,k,j]: 0<=i<n and 0<=k<m and 0<=j<l}",
             [
                 """
                 c[i, j, k] = (2*a[i,j,k])%(2+b[i,j,k]/3.0)
@@ -124,10 +124,10 @@ def test_op_counter_specialops(ctx_factory):
     assert i32 == n*m
 
 
-def test_op_counter_bitwise(ctx_factory):
+def test_op_counter_bitwise():
 
     knl = lp.make_kernel(
-            "[n,m,l] -> {[i,k,j]: 0<=i<n and 0<=k<m and 0<=j<l}",
+            "{[i,k,j]: 0<=i<n and 0<=k<m and 0<=j<l}",
             [
                 """
                 c[i, j, k] = (a[i,j,k] | 1) + (b[i,j,k] & 1)
@@ -143,10 +143,43 @@ def test_op_counter_bitwise(ctx_factory):
     m = 256
     l = 128
     i32 = poly.dict[np.dtype(np.int32)].eval_with_dict({'n': n, 'm': m, 'l': l})
+    print(poly.dict[np.dtype(np.int32)])
     not_there = poly[np.dtype(np.float64)].eval_with_dict({'n': n, 'm': m, 'l': l})
     assert i32 == 3*n*m+n*m*l
     assert not_there == 0
 
+
+def test_op_counter_triangular_domain():
+
+    knl = lp.make_kernel(
+            "{[i,j]: 0<=i<n and 0<=j<m and i<j}",
+            """
+            a[i, j] = b[i,j] * 2
+            """,
+            name="bitwise", assumptions="n,m >= 1")
+
+    knl = lp.add_and_infer_dtypes(knl,
+            dict(b=np.float64))
+
+    expect_fallback = False
+    import islpy as isl
+    try:
+        isl.BasicSet.carod
+    except AttributeError:
+        expect_fallback = True
+    else:
+        expect_fallback = False
+
+    poly = get_op_poly(knl)[np.dtype(np.float64)]
+    value_dict = dict(m=13, n=200)
+    flops = poly.eval_with_dict(value_dict)
+
+    if expect_fallback:
+        assert flops == 144
+    else:
+        assert flops == 78
+
+
 if __name__ == "__main__":
     if len(sys.argv) > 1:
         exec(sys.argv[1])
-- 
GitLab