Tackle many warnings

.github/workflows/build.yml

@@ -83,6 +83,7 @@ jobs:
             --install defcon \
             --install gadopt \
             --install asQ \
+            --package-branch loopy connorjward/merge-upstream \
             || (cat firedrake-install.log && /bin/false)
       - name: Install test dependencies
         run: |

.github/workflows/pyop2.yml

@@ -107,6 +107,8 @@ jobs:
         run: |
           source ../venv/bin/activate
           python -m pip install -v ".[test]"
+          python -m pip uninstall -y pytools loopy
+          python -m pip install git+https://github.com/firedrakeproject/loopy.git@connorjward/merge-upstream
 
       - name: Run tests
         shell: bash

pyop2/codegen/loopycompat.py

@@ -52,7 +52,7 @@ def map_subscript(self, expr):
             flattened_index -= (int(stride) * ind)
             new_indices.append(simplify_via_aff(ind))
 
-        return expr.aggregate.index(tuple(new_indices))
+        return expr.aggregate[tuple(new_indices)]
 
 
 def _match_caller_callee_argument_dimension_for_single_kernel(

pyop2/codegen/rep2loopy.py

@@ -1,5 +1,6 @@
 import ctypes
 import numpy
+from dataclasses import dataclass
 
 import loopy
 from loopy.symbolic import SubArrayRef
@@ -113,10 +114,12 @@ def __init__(self, name=None, arg_id_to_dtype=None,
                  arg_id_to_descr=None, name_in_target=None):
         if name is not None:
             assert name == self.name
+
+        name_in_target = name_in_target if name_in_target else self.name
         super(LACallable, self).__init__(self.name,
                                          arg_id_to_dtype=arg_id_to_dtype,
-                                         arg_id_to_descr=arg_id_to_descr)
-        self.name_in_target = name_in_target if name_in_target else self.name
+                                         arg_id_to_descr=arg_id_to_descr,
+                                         name_in_target=name_in_target)
 
     @abc.abstractproperty
     def name(self):
@@ -205,16 +208,18 @@ def __call__(self, preamble_info):
         yield ("0", self.preamble)
 
 
+@dataclass(frozen=True, init=False)
 class PyOP2KernelCallable(loopy.ScalarCallable):
     """Handles PyOP2 Kernel passed in as a string
     """
 
-    fields = set(["name", "parameters", "arg_id_to_dtype", "arg_id_to_descr", "name_in_target"])
     init_arg_names = ("name", "parameters", "arg_id_to_dtype", "arg_id_to_descr", "name_in_target")
 
+    parameters: tuple
+
     def __init__(self, name, parameters, arg_id_to_dtype=None, arg_id_to_descr=None, name_in_target=None):
-        super(PyOP2KernelCallable, self).__init__(name, arg_id_to_dtype, arg_id_to_descr, name_in_target)
-        self.parameters = parameters
+        super().__init__(name, arg_id_to_dtype, arg_id_to_descr, name_in_target)
+        object.__setattr__(self, "parameters", tuple(parameters))
 
     def with_types(self, arg_id_to_dtype, callables_table):
         new_arg_id_to_dtype = arg_id_to_dtype.copy()
@@ -288,8 +293,8 @@ def runtime_indices(expressions):
     for node in traversal(expressions):
         if isinstance(node, RuntimeIndex):
             indices.append(node.name)
-
-    return frozenset(indices)
+    # use a dict as an ordered set
+    return {i: None for i in indices}
 
 
 def imperatives(exprs):
@@ -325,7 +330,7 @@ def loop_nesting(instructions, deps, outer_inames, kernel_name):
 
     # boost inames, if one instruction is inside inner inames (free indices),
     # it should be inside the outer inames as dictated by other instructions.
-    index_nesting = defaultdict(frozenset)  # free index -> {runtime indices}
+    index_nesting = defaultdict(dict)  # free index -> {runtime indices}
     for insn in instructions:
         if isinstance(insn, When):
             key = insn.children[1]
@@ -338,7 +343,7 @@ def loop_nesting(instructions, deps, outer_inames, kernel_name):
     for insn in imperatives(instructions):
         outer = reduce(operator.or_,
                        iter(index_nesting[fi] for fi in traversal([insn]) if isinstance(fi, Index)),
-                       frozenset())
+                       {})
         nesting[insn] = nesting[insn] | outer
 
     return nesting
@@ -407,11 +412,10 @@ def generate(builder, wrapper_name=None):
     Materialise._count = itertools.count()
     RuntimeIndex._count = itertools.count()
 
+    # use a dict as an ordered set
+    outer_inames = {builder._loop_index.name: None}
     if builder.layer_index is not None:
-        outer_inames = frozenset([builder._loop_index.name,
-                                  builder.layer_index.name])
-    else:
-        outer_inames = frozenset([builder._loop_index.name])
+        outer_inames.update({builder.layer_index.name: None})
 
     instructions = list(builder.emit_instructions())
 
@@ -476,10 +480,17 @@ def renamer(expr):
     deps = instruction_dependencies(instructions, mapper.initialisers)
     within_inames = loop_nesting(instructions, deps, outer_inames, parameters.kernel_name)
 
+    # used to avoid disadvantageous loop interchanges
+    loop_priorities = set()
+    for iname_nest in within_inames.values():
+        if len(iname_nest) > 1:
+            loop_priorities.add(tuple(iname_nest.keys()))
+    loop_priorities = frozenset(loop_priorities)
+
     # generate loopy
     context = Bag()
     context.parameters = parameters
-    context.within_inames = within_inames
+    context.within_inames = {k: frozenset(v.keys()) for k, v in within_inames.items()}
     context.conditions = []
     context.index_ordering = []
     context.instruction_dependencies = deps
@@ -544,11 +555,8 @@ def renamer(expr):
                                 options=options,
                                 assumptions=assumptions,
                                 lang_version=(2018, 2),
-                                name=wrapper_name)
-
-    # prioritize loops
-    for indices in context.index_ordering:
-        wrapper = loopy.prioritize_loops(wrapper, indices)
+                                name=wrapper_name,
+                                loop_priority=loop_priorities)
 
     # register kernel
     kernel = builder.kernel

pyop2/compilation.py

@@ -547,7 +547,7 @@ def check_source_hashes(compiler, jitmodule, extension, comm):
             output = Path(configuration["cache_dir"]).joinpath("mismatching-kernels")
             srcfile = output.joinpath(f"src-rank{icomm.rank}.{extension}")
             if icomm.rank == 0:
-                output.mkdir(exist_ok=True)
+                output.mkdir(parents=True, exist_ok=True)
             icomm.barrier()
             with open(srcfile, "w") as fh:
                 fh.write(jitmodule.code_to_compile)

pyop2/local_kernel.py

@@ -120,13 +120,11 @@ def cache_key(self):
     def _immutable_cache_key(self):
         # We need this function because self.accesses is mutable due to legacy support
         if isinstance(self.code, lp.TranslationUnit):
-            key_hash = hashlib.sha256()
-            self.code.update_persistent_hash(key_hash, LoopyKeyBuilder())
-            code = key_hash.hexdigest()
+            code_key = LoopyKeyBuilder()(self.code)
         else:
-            code = self.code
+            code_key = self.code
 
-        key = (code, self.name, self.cpp, self.flop_count,
+        key = (code_key, self.name, self.cpp, self.flop_count,
                self.headers, self.include_dirs, self.ldargs, sorted(self.opts.items()),
                self.requires_zeroed_output_arguments, self.user_code, version.__version__)
         return hashlib.md5(str(key).encode()).hexdigest()

pyop2/types/dat.py

@@ -358,14 +358,14 @@ def _op_kernel(self, op, globalp, dtype):
         _self = p.Variable("self")
         _ret = p.Variable("ret")
         i = p.Variable("i")
-        lhs = _ret.index(i)
+        lhs = _ret[i]
         if globalp:
-            rhs = _other.index(0)
+            rhs = _other[0]
             rshape = (1, )
         else:
-            rhs = _other.index(i)
+            rhs = _other[i]
             rshape = (self.cdim, )
-        insn = lp.Assignment(lhs, op(_self.index(i), rhs), within_inames=frozenset(["i"]))
+        insn = lp.Assignment(lhs, op(_self[i], rhs), within_inames=frozenset(["i"]))
         data = [lp.GlobalArg("self", dtype=self.dtype, shape=(self.cdim,)),
                 lp.GlobalArg("other", dtype=dtype, shape=rshape),
                 lp.GlobalArg("ret", dtype=self.dtype, shape=(self.cdim,))]
@@ -405,15 +405,15 @@ def _iop_kernel(self, op, globalp, other_is_self, dtype):
         _other = p.Variable("other")
         _self = p.Variable("self")
         i = p.Variable("i")
-        lhs = _self.index(i)
+        lhs = _self[i]
         rshape = (self.cdim, )
         if globalp:
-            rhs = _other.index(0)
+            rhs = _other[0]
             rshape = (1, )
         elif other_is_self:
-            rhs = _self.index(i)
+            rhs = _self[i]
         else:
-            rhs = _other.index(i)
+            rhs = _other[i]
         insn = lp.Assignment(lhs, op(lhs, rhs), within_inames=frozenset(["i"]))
         data = [lp.GlobalArg("self", dtype=self.dtype, shape=(self.cdim,))]
         if not other_is_self:
@@ -518,7 +518,7 @@ def _neg_kernel(self):
         lvalue = p.Variable("other")
         rvalue = p.Variable("self")
         i = p.Variable("i")
-        insn = lp.Assignment(lvalue.index(i), -rvalue.index(i), within_inames=frozenset(["i"]))
+        insn = lp.Assignment(lvalue[i], -rvalue[i], within_inames=frozenset(["i"]))
         data = [lp.GlobalArg("other", dtype=self.dtype, shape=(self.cdim,)),
                 lp.GlobalArg("self", dtype=self.dtype, shape=(self.cdim,))]
         knl = lp.make_function([domain], [insn], data, name=name, target=conf.target, lang_version=(2018, 2))

pyproject.toml

@@ -24,7 +24,7 @@ dependencies = [
   "pkgconfig",
   "progress",
   "pycparser",
-  "pytools",
+  "pytools[siphash]",
   "requests",
   "rtree>=1.2",
   "scipy",