Merge pull request #119 from coneoproject/fix-issue-118

Deterministically collect loop nests
coneoproject · Apr 11, 2017 · bdbc5fc · bdbc5fc
2 parents 50a0f59 + ef0f7f4
commit bdbc5fc
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/coffee/plan.py b/coffee/plan.py
@@ -43,7 +43,7 @@
 from .logger import log, warn, PERF_OK, PERF_WARN
 from coffee.visitors import Find, EstimateFlops
 
-from collections import defaultdict, OrderedDict
+from collections import OrderedDict
 import time
 
 
@@ -88,7 +88,7 @@ def plan_cpu(self, opts):
 
             info = visit(kernel, info_items=['decls', 'exprs'])
             # Collect expressions and related metadata
-            nests = defaultdict(OrderedDict)
+            nests = OrderedDict()
             for stmt, expr_info in info['exprs'].items():
                 parent, nest = expr_info
                 if not nest:
@@ -98,7 +98,7 @@ def plan_cpu(self, opts):
                 else:
                     typ = check_type(stmt, info['decls'])
                 metaexpr = MetaExpr(typ, parent, nest)
-                nests[nest[0]].update({stmt: metaexpr})
+                nests.setdefault(nest[0], OrderedDict()).update({stmt: metaexpr})
             loop_opts = [CPULoopOptimizer(loop, header, exprs)
                          for (loop, header), exprs in nests.items()]
 
@@ -185,7 +185,7 @@ def plan_gpu(self):
         for kernel in kernels:
             info = visit(kernel, info_items=['decls', 'exprs'])
             # Collect expressions and related metadata
-            nests = defaultdict(OrderedDict)
+            nests = OrderedDict()
             for stmt, expr_info in info['exprs'].items():
                 parent, nest = expr_info
                 if not nest:
@@ -195,7 +195,7 @@ def plan_gpu(self):
                 else:
                     typ = check_type(stmt, info['decls'])
                 metaexpr = MetaExpr(typ, parent, nest)
-                nests[nest[0]].update({stmt: metaexpr})
+                nests.setdefault(nest[0], OrderedDict()).update({stmt: metaexpr})
             loop_opts = [GPULoopOptimizer(loop, header, exprs)
                          for (loop, header), exprs in nests.items()]