Skip to content

Commit

Permalink
Tackle many warnings (#3898)
Browse files Browse the repository at this point in the history
* Various loopy fixes

* Remove deprecated pymbolic function

* Install siphash via pytools to avoid producing many warnings
  • Loading branch information
connorjward authored Dec 5, 2024
1 parent 452b342 commit 57217cf
Show file tree
Hide file tree
Showing 6 changed files with 42 additions and 36 deletions.
2 changes: 1 addition & 1 deletion pyop2/codegen/loopycompat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
46 changes: 27 additions & 19 deletions pyop2/codegen/rep2loopy.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import ctypes
import numpy
from dataclasses import dataclass

import loopy
from loopy.symbolic import SubArrayRef
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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]
Expand All @@ -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
Expand Down Expand Up @@ -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())

Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion pyop2/compilation.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
8 changes: 3 additions & 5 deletions pyop2/local_kernel.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
18 changes: 9 additions & 9 deletions pyop2/types/dat.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,))]
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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))
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ dependencies = [
"pkgconfig",
"progress",
"pycparser",
"pytools",
"pytools[siphash]",
"requests",
"rtree>=1.2",
"scipy",
Expand Down

0 comments on commit 57217cf

Please sign in to comment.