diff --git a/pyop3/expr/base.py b/pyop3/expr/base.py index c115bb7745..1b4c88608c 100644 --- a/pyop3/expr/base.py +++ b/pyop3/expr/base.py @@ -15,7 +15,6 @@ from pyop3.axis_tree import UNIT_AXIS_TREE, AxisTree from pyop3.node import Node, Terminal - class Expression(Node, abc.ABC): # {{{ abstract methods @@ -333,7 +332,7 @@ def local_min(self) -> numbers.Number: from pyop3.expr.visitors import get_local_min return get_local_min(self.a) + get_local_min(self.b) - + # }}} diff --git a/pyop3/insn/exec.py b/pyop3/insn/exec.py index 1bb2b3283b..583580598b 100644 --- a/pyop3/insn/exec.py +++ b/pyop3/insn/exec.py @@ -56,6 +56,8 @@ class CompilerParameters: # TODO: handle these - need to build CompilerOptions + backend: str = "loopy" + """ Option to select 'loopy' or 'mlir' as code generation backends. """ # extra_cflags: tuple[str, ...] = () # extra_ldflags: tuple[str, ...] = () @@ -228,7 +230,7 @@ def compile(self) -> Callable[[int, ...], None]: ) def _compile(self) -> CompiledCodeExecutor: from pyop3.insn.visitors import collect_compiler_options - from pyop3.lower.loopy import _compile_static + from pyop3.lower.codegen import _compile_static # Preprocess the instruction. This is an expensive operation so we # want to avoid doing it if at all possible. @@ -683,7 +685,7 @@ def _(self, arr: np.ndarray, /) -> int: try: import cupy as cp - @_handle_to_pointer.register + @_handle_to_pointer.register(cp.ndarray) def _(self, arr: cp.ndarray, /) -> int: # NOTE: This gives a pointer to a GPU memory address. # Loopy cannot work with GPU so this will lead to a segfault. diff --git a/pyop3/lower/codegen.py b/pyop3/lower/codegen.py new file mode 100644 index 0000000000..b0c3d93b75 --- /dev/null +++ b/pyop3/lower/codegen.py @@ -0,0 +1,463 @@ +from __future__ import annotations + +import abc +import contextlib +import functools +import numbers +import os +from typing import Any + +import loopy as lp +import numpy as np +import pymbolic as pym +from immutabledict import immutabledict as idict +from petsc4py import PETSc + +import pyop3.axis_tree +import pyop3.buffer +import pyop3.cache +import pyop3.config +import pyop3.constants +import pyop3.dtypes +import pyop3.expr +from pyop3 import mpi, utils +from pyop3.axis_tree.tree import ( + UNIT_AXIS_TREE, + IndexedAxisTree, +) +from pyop3.buffer import ( + AbstractBuffer, + NullBuffer, + PetscMatBuffer, +) +from pyop3.constants import INC, MAX_RW, MAX_WRITE, MIN_RW, MIN_WRITE, READ, RW, WRITE +from pyop3.dtypes import IntType +from pyop3.insn.base import ( + AssignmentType, + Exscan, + InstructionList, + Loop, + NonEmptyArrayAssignment, + NullInstruction, + StandaloneCalledFunction, + assignment_type_as_intent, +) + +from pyop3.lower.loopy import LoopyCodegenContext + +# TODO: import other way around? +from pyop3.lower.transform import ( + _collect_temporary_shapes +) + +def _compile_static_hashkey(op: PreprocessedOperation, compiler_parameters: ParsedCompilerParameters) -> Hashable: + return (op.disk_cache_key, compiler_parameters, pyop3.config) + +@pyop3.cache.memory_and_disk_cache( + hashkey=_compile_static_hashkey, + get_comm=lambda op, *a, **kw: op.comm, +) +def _compile_static(op: InstructionExecutionContext, compiler_parameters: ParsedCompilerParameters) -> tuple: + """Compile the operation without regard for specific data values. + + This function is therefore suitable for disk caching. + Function passes compilation process to compiler-specified backend + + Returns + ------- + TU + datamap + + """ + insn = op.preprocess() + function_name = "pyop3_loop" # TODO: Provide as kwarg + + if isinstance(insn, InstructionList): + cs_expr = insn.instructions + else: + cs_expr = (insn,) + + if compiler_parameters.backend == "loopy": + make_context = LoopyCodegenContext + elif compiler_parameters.backend == "mlir": + raise NotImplementedError("Class is still being implemented.") + + context = make_context( + propagate_negatives=compiler_parameters.propagate_negatives, + mask_array_accesses=compiler_parameters.mask_array_accesses, + ) + # NOTE: so I think LoopCollection is a better abstraction here - don't want to be + # explicitly dealing with contexts at this point. Can always sniff them out again. + # for context, ex in cs_expr: + for ex in cs_expr: + # ex = expand_implicit_pack_unpack(ex) + + # add external loop indices as kernel arguments + # FIXME: removed because cs_expr needs to sniff the context now + loop_indices = {} + + for e in pyop3.collections.as_tuple(ex): # TODO: get rid of this loop + # context manager? + context.set_temporary_shapes(_collect_temporary_shapes(e)) + _compile(e, loop_indices, context) + + if not context.buffer_intents: + raise pyop3.exceptions.EffectlessComputationException( + "The generated kernel does not modify any global data, this may indicate that something has gone wrong" + ) + + translation_unit = context.finalize_kernel(function_name, compiler_parameters) + + # Extra information needed by the code executor + kernel_name_to_buffer_info = utils.invert_mapping(context.kernel_names) + buffer_intents = context.buffer_intents + + # Replace buffers with their indices, dropping any temporaries. Also + # match the calling order for the kernel. + kernel_name_to_global_buffer_info = {} + global_buffer_intents = {} + for kernel_arg in translation_unit.default_entrypoint.args: + buf_view = kernel_name_to_buffer_info[kernel_arg.name] + buf_index = op.preprocessed_buffers.index(buf_view.buffer) + kernel_name_to_global_buffer_info[kernel_arg.name] = (buf_index, buf_view.nest_indices) + global_buffer_intents[buf_index] = buffer_intents[buf_view.buffer] + + return translation_unit, kernel_name_to_global_buffer_info, global_buffer_intents + +@functools.singledispatch +def _compile(expr: Any, loop_indices: Dict, codegen_context: CodegenContext) -> None: + raise TypeError(f"No handler defined for {type(expr).__name__}") + +@_compile.register(NullInstruction) +def _(null, *args, **kwargs): + pass + +@_compile.register(InstructionList) +def _( + insn_list, + loop_indices, + codegen_context +) -> None: + for insn in insn_list: + _compile(insn, loop_indices, codegen_context) + +@_compile.register(Loop) +def _( + loop, + loop_indices, + codegen_context +) -> None: + _compile_loop( + loop, + loop.index.iterset, + loop_indices, + codegen_context, + ) + +@_compile.register(StandaloneCalledFunction) +def _(call, loop_indices, codegen_context): + args = [(arg, spec) for arg, spec in zip(call.arguments, call.argspec, strict=True)] + codegen_context.add_function_call(call.function.code, args) + subkernel = call.function.code.with_entrypoints(frozenset()) + codegen_context.add_subkernel(subkernel) + +@_compile.register(NonEmptyArrayAssignment) +def parse_assignment(assignment: NonEmptyArrayAssignment, loop_indices, codegen_context: CodegenContext): + if any(isinstance(arg, pyop3.expr.MatPetscMatBufferExpression) for arg in assignment.arguments): + _compile_petsc_mat(assignment, loop_indices, codegen_context) + else: + _compile_array_assignment( + assignment, + loop_indices, + assignment.axis_trees, + codegen_context, + ) + +def _compile_petsc_mat( + assignment, + loop_indices, + codegen_context, +): + if not isinstance(codegen_context, LoopyCodegenContext): + raise NotImplementedError("Only supported for Loopy") + # We need to know whether the matrix is the assignee or not because we need + # to know whether to put MatGetValues or MatSetValues + if isinstance(assignment.assignee.buffer_view.buffer, PetscMatBuffer): + mat = assignment.assignee + expr = assignment.expression + setting_mat_values = True + else: + mat = assignment.expression + expr = assignment.assignee + setting_mat_values = False + + row_axis_tree, column_axis_tree = assignment.axis_trees + + assert isinstance(expr, pyop3.expr.BufferExpression) + + # now emit the right line of code, this should properly be a lp.ScalarCallable + # https://petsc.org/release/manualpages/Mat/MatGetValuesLocal/ + mat_name = codegen_context.add_buffer(mat.buffer_view, assignment_type_as_intent(assignment.assignment_type)) + array_name = codegen_context.add_buffer(expr.buffer_view, READ) + + rsize = row_axis_tree.size + csize = column_axis_tree.size + + # these sizes can be expressions that need evaluating + rsize_var = codegen_context.register_extent( + rsize, + {}, + loop_indices, + ) + + csize_var = codegen_context.register_extent( + csize, + {}, + loop_indices, + ) + + # convert the generic expressions to + # for example: + # + # map0[3*i0 + i1] + # map0[3*i0 + i2 + 3] + # + # to the shared top-level layout: + # + # map0[3*i0] + # + # which is what Mat{Get,Set}Values() needs. + layout_exprs = [] + for layout in [mat.row_layout, mat.column_layout]: + subst_sublayout = layout.layouts[idict()] + subst_layout = pyop3.expr.LinearDatBufferExpression(layout.buffer, subst_sublayout) + layout_expr = codegen_context.lower_expr(subst_layout, ((),), loop_indices) + layout_exprs.append(layout_expr) + irow, icol = layout_exprs + + # FIXME: + blocked = False + + # hacky + myargs = [ + assignment, mat_name, array_name, rsize_var, csize_var, irow, icol, blocked + ] + if setting_mat_values: + match assignment.assignment_type: + case AssignmentType.WRITE: + call_str = _petsc_mat_store(*myargs) + case AssignmentType.INC: + call_str = _petsc_mat_add(*myargs) + case _: + raise AssertionError + else: + call_str = _petsc_mat_load(*myargs) + + codegen_context.add_cinstruction(call_str) + + +def _compile_array_assignment( + assignment, + loop_indices, + axis_trees, + codegen_context, + *, + iname_replace_maps=None, + # TODO document these under "Other Parameters" + axis_tree=None, + paths=None +): + if paths is None: + paths = [] + if iname_replace_maps is None: + iname_replace_maps = [] + + if axis_tree is None: + axis_tree, *axis_trees = axis_trees + + paths += [idict()] + iname_replace_maps += [idict()] + + if axis_tree.is_empty or axis_tree is UNIT_AXIS_TREE or isinstance(axis_tree, IndexedAxisTree): + if axis_trees: + raise NotImplementedError("Refactor needed") + + codegen_context.add_leaf_assignment( + assignment.assignee, + assignment.expression, + assignment.assignment_type, + paths, + iname_replace_maps, + loop_indices + ) + return + + axis = axis_tree.node_map[paths[-1]] + for component in axis.components: + new_paths = paths.copy() + new_paths[-1] = paths[-1] | {axis.label: component.label} + + if axis_tree.linearize(new_paths[-1], partial=True).size == 0: + continue + + if component.local_size != 1: + iname = codegen_context.unique_name("i") + ext = codegen_context.register_extent( + component.size, + iname_replace_maps[-1], + loop_indices + ) + codegen_context.add_domain(iname, ext) + new_maps = iname_replace_maps.copy() + new_maps[-1] = iname_replace_maps[-1] | {axis.label: codegen_context.var(iname)} + within_inames = {iname} + else: + new_maps = iname_replace_maps.copy() + new_maps[-1] = iname_replace_maps[-1] | {axis.label: 0} + within_inames = set() + + with codegen_context.within_inames(within_inames): + if axis_tree.node_map[new_paths[-1]]: + _compile_array_assignment( + assignment, + loop_indices, + axis_trees, + codegen_context, + iname_replace_maps=new_maps, + axis_tree=axis_tree, + paths=new_paths + ) + elif axis_trees: + _compile_array_assignment( + assignment, + loop_indices, + axis_trees, + codegen_context, + iname_replace_maps=new_maps, + axis_tree=None, + paths=new_paths + ) + else: + codegen_context.add_leaf_assignment( + assignment.assignee, + assignment.expression, + assignment.assignment_type, + new_paths, + new_maps, + loop_indices + ) + +def _compile_loop( + loop, + axis_tree, + loop_indices, + codegen_context, + *, + axis=None, + path=None, + iname_map=None, +) -> None: + if axis_tree is UNIT_AXIS_TREE: + # NOTE: might need an expression here sometimes + for statement in loop.statements: + _compile( + statement, + # loop_indices | dict(loop_exprs), + loop_indices, + codegen_context, + ) + return + + if utils.strictly_all(x is None for x in {axis, path, iname_map}): + axis = axis_tree.root + path = idict() + iname_map = idict() + + for component in axis.components: + path_ = path | {axis.label: component.label} + + if axis_tree.linearize(path_, partial=True).size == 0: + continue + elif component.size != 1: + iname = codegen_context.unique_name("i") + domain_var = codegen_context.register_extent( + component.size, + iname_map, + loop_indices + ) + codegen_context.add_domain(iname, domain_var) + iname_replace_map_ = iname_map | {axis.label: codegen_context.var(iname)} + within_inames = frozenset({iname}) + else: + iname_replace_map_ = iname_map | {axis.label: 0} + within_inames = set() + + with codegen_context.within_inames(within_inames): + if subaxis := axis_tree.node_map[path_]: + _compile_loop( + loop, + axis_tree, + loop_indices, + codegen_context, + axis=subaxis, + path=path_, + iname_map=iname_replace_map_ + ) + else: + loop_indices |= idict({ + (loop.index.id, axis_label): iname + for axis_label, iname in iname_replace_map_.items() + }) + for statement in loop.statements: + _compile( + statement, + loop_indices, + codegen_context, + ) + +@_compile.register(Exscan) +def _(exscan, loop_indices, codegen_context): + if not isinstance(codegen_context, LoopyCodegenContext): + raise NotImplementedError("Only supported for Loopy") + + if exscan.scan_type != "+": + raise NotImplementedError + + domain_var = codegen_context.register_extent( + exscan.extent, + {}, + loop_indices, + ) + + iname = codegen_context.unique_name("i") + codegen_context.add_domain(iname, domain_var) + + iname_var = codegen_context.var(iname) + iname_map = {exscan.scan_axis.label: codegen_context.var(iname)} + + lexpr = codegen_context.lower_expr(exscan.assignee, [iname_map], loop_indices, intent=RW) + rexpr = lexpr + codegen_context.lower_expr(exscan.expression, [iname_map], loop_indices) + + lexpr = pym.substitute(lexpr, {iname: iname_var+1}) + codegen_context.add_assignment(lexpr, rexpr) + + +def _petsc_mat_load(assignment, mat_name, array_name, nrow, ncol, irow, icol, blocked): + if blocked: + return f"MatGetValuesBlockedLocal({mat_name}, {nrow}, &({irow}), {ncol}, &({icol}), &({array_name}[0]));" + else: + return f"MatGetValuesLocal({mat_name}, {nrow}, &({irow}), {ncol}, &({icol}), &({array_name}[0]));" + + +def _petsc_mat_store(assignment, mat_name, array_name, nrow, ncol, irow, icol, blocked): + if blocked: + return f"MatSetValuesBlockedLocal({mat_name}, {nrow}, &({irow}), {ncol}, &({icol}), &({array_name}[0]), INSERT_VALUES);" + else: + return f"MatSetValuesLocal({mat_name}, {nrow}, &({irow}), {ncol}, &({icol}), &({array_name}[0]), INSERT_VALUES);" + + +def _petsc_mat_add(assignment, mat_name, array_name, nrow, ncol, irow, icol, blocked): + if blocked: + return f"MatSetValuesBlockedLocal({mat_name}, {nrow}, &({irow}), {ncol}, &({icol}), &({array_name}[0]), ADD_VALUES);" + else: + return f"MatSetValuesLocal({mat_name}, {nrow}, &({irow}), {ncol}, &({icol}), &({array_name}[0]), ADD_VALUES);" diff --git a/pyop3/lower/context.py b/pyop3/lower/context.py new file mode 100644 index 0000000000..45b4b07504 --- /dev/null +++ b/pyop3/lower/context.py @@ -0,0 +1,185 @@ +from abc import ABC, abstractmethod +from typing import Any, List, Dict, Tuple +import numbers +import functools + +import pymbolic as pym +from immutabledict import immutabledict as idict + +import pyop3.axis_tree +import pyop3.buffer +import pyop3.cache +import pyop3.config +import pyop3.constants +import pyop3.dtypes +import pyop3.expr + +from pyop3.axis_tree.tree import ( + UNIT_AXIS_TREE, + IndexedAxisTree, +) +from pyop3 import mpi, utils +from pyop3.buffer import IndexedBuffer +from pyop3.insn.base import Intent, assignment_type_as_intent +from pyop3.constants import INC, MAX_RW, MAX_WRITE, MIN_RW, MIN_WRITE, READ, RW, WRITE + +from pyop3.insn.base import ( + AbstractAssignment, + AssignmentType, + Exscan, + InstructionList, + Loop, + NonEmptyArrayAssignment, + NullInstruction, + StandaloneCalledFunction, + assignment_type_as_intent, +) + +class CodegenContext(ABC): + """ + Base class for code generation backends + + Class designed solely for use in codegen.py, as an interface to specific backends. + """ + + def __init__(self, *, propagate_negatives: bool, mask_array_accesses: bool) -> None: + self.propagate_negatives = propagate_negatives + self.mask_array_accesses = mask_array_accesses + + self._domains = [] + self._instructions = [] + self._arguments = [] + self._subkernels = [] + self._last_insn_id = None # determine dependence + + self._name_generator = utils.UniqueNameGenerator() + + # (buffer, nest_indices) -> name in kernel + self.kernel_names = {} + + # buffer name -> buffer + self.buffer_intents = {} + + # assignee name -> indirection expression + self._assignees = {} + + @property + def domains(self) -> Tuple: + return tuple(self._domains) + + @property + def instructions(self) -> Tuple: + return tuple(self._instructions) + + @property + def arguments(self) -> Tuple: + return tuple(sorted(self._arguments, key=lambda arg: getattr(arg, 'name', ''))) + + @property + def subkernels(self) -> Tuple: + return tuple(self._subkernels) + + @property + def _depends_on(self): + return frozenset({self._last_insn_id}) - {None} + + def _add_instruction(self, insn): + self._instructions.append(insn) + self._last_insn_id = insn.id + + # {{{ abstract methods + + + @abstractmethod + def var(self, iname: str, *args) -> str | pym.primitives.Variable: + """ + Implementation to represent symbolic variable for respective IR + """ + pass + + + @abstractmethod + def add_domain(self, iname: str, *args) -> None: + pass + + @abstractmethod + def add_assignment(self, assignee, expression, prefix: str = "insn") -> None: + pass + + @abstractmethod + def add_function_call(self, assignees, expression, prefix: str = "insn") -> None: + pass + + @abstractmethod + def add_buffer(self, buffer_view: IndexedBuffer, intent: Intent | None = None) -> str: + pass + + @abstractmethod + def add_subkernel(self, subkernel) -> None: + pass + + @abstractmethod + def set_temporary_shapes(self, shapes) -> None: + pass + + @abstractmethod + def lower_expr(self, expr, iname_maps, loop_indices, + intent: Intent | None = None, paths = None): + """ + Lower a PyOP3 expression to the target's IR representation. + + Returns: + - pymbolic for Loopy + - xDSL for MLIR + """ + pass + + @abstractmethod + def lower_buffer_access( + self, + buffer: IndexedBuffer, + layouts, + iname_maps, + loop_indices, + *, + intent + ): + """ + Determine indexing and lower buffer expression to respective IR + """ + pass + + @abstractmethod + def add_leaf_assignment( + self, + assignee, + expression, + assignment_type, + paths, + iname_maps, + loop_indices + ): + pass + + @abstractmethod + def register_extent(self, obj: Any, inames, loop_indices): + pass + + # }}} + + def add_subkernel(self, subkernel): + self._subkernels.append(subkernel) + + def unique_name(self, prefix: str) -> str: + return self._name_generator(prefix) + + def __str__(self) -> str: + ''' + Display key properties of CodegenContext + ''' + ctx = f"Domain: {str(self.domains)}\n\n" + ctx += f"Instructions: {str(self.instructions)}\n\n" + ctx += f"Arguments: {str(self.arguments)}\n\n" + ctx += f"Subkernels: {str(self.subkernels)}\n\n" + return ctx + diff --git a/pyop3/lower/loopy.py b/pyop3/lower/loopy.py index 21616af426..bd74bde6bd 100644 --- a/pyop3/lower/loopy.py +++ b/pyop3/lower/loopy.py @@ -44,50 +44,30 @@ assignment_type_as_intent, ) +from pyop3.lower.context import CodegenContext + # TODO: import other way around? from pyop3.lower.transform import ( with_attach_debugger, with_likwid_markers, with_petsc_event, ) - # FIXME this needs to be synchronised with TSFC, tricky # shared base package? or both set by Firedrake - better solution LOOPY_TARGET = lp.CWithGNULibcTarget() LOOPY_LANG_VERSION = (2018, 2) - -class CodegenContext(abc.ABC): - pass - - class LoopyCodegenContext(CodegenContext): def __init__(self, *, propagate_negatives: bool, mask_array_accesses: bool) -> None: - self.propagate_negatives = propagate_negatives - self.mask_array_accesses = mask_array_accesses - - self._domains = [] - self._instructions = [] - self._arguments = [] - self._subkernels = [] - + super().__init__( + propagate_negatives=propagate_negatives, + mask_array_accesses=mask_array_accesses + ) self._within_inames = frozenset() - self._last_insn_id = None - - self._name_generator = utils.UniqueNameGenerator() - - # (buffer, nest_indices) -> name in kernel - self.kernel_names = {} - - # buffer -> intent - self.buffer_intents = {} # initializer hash -> temporary name self._reusable_temporaries: dict[int, str] = {} - # assignee name -> indirection expression - self._assignees = {} - @property def domains(self) -> tuple: return tuple(self._domains) @@ -104,6 +84,12 @@ def arguments(self) -> tuple: def subkernels(self) -> tuple: return tuple(self._subkernels) + def add_subkernel(self, subkernel): + self._subkernels.append(subkernel) + + def var(self, iname: str, *args) -> pym.primitives.Variable: + return pym.var(iname) + def add_domain(self, iname, *args): nargs = len(args) if nargs == 1: @@ -137,7 +123,40 @@ def add_cinstruction(self, insn_str, read_variables=frozenset()): ) self._add_instruction(cinsn) - def add_function_call(self, assignees, expression, prefix="insn"): + def add_function_call(self, code, args, prefix="insn"): + subarrayrefs = {} + loopy_args = code.default_entrypoint.args + for loopy_arg, (arg, spec) in zip(loopy_args, args, strict=True): + name_in_kernel = self.add_buffer(arg.buffer_view, spec.intent) + if isinstance(loopy_arg, lp.ArrayArg): + # array arguments to an inner kernel require all strides to be defined + indices = [] + for s in loopy_arg.shape: + iname = self.unique_name("i") + self.add_domain(iname, s) + indices.append(pym.var(iname)) + indices = tuple(indices) + subarrayrefs[arg] = lp.symbolic.SubArrayRef( + indices, pym.var(name_in_kernel)[indices] + ) + else: + assert isinstance(loopy_arg, lp.ValueArg) + subarrayrefs[arg] = pym.var(name_in_kernel) + + assignees = tuple( + subarrayrefs[arg] + for arg, spec in args + if spec.intent in {WRITE, RW, INC, MIN_RW, MIN_WRITE, MAX_RW, MAX_WRITE} + ) + expression = pym.primitives.Call( + pym.var(code.default_entrypoint.name), + tuple( + subarrayrefs[arg] + for arg, spec in args + if spec.intent in {READ, RW, INC, MIN_RW, MAX_RW} + ), + ) + insn = lp.CallInstruction( assignees, expression, @@ -272,12 +291,6 @@ def add_opaque(self, opaque: OpaqueTerminal, intent) -> str: self.kernel_names[opaque] = name_in_kernel return name_in_kernel - def add_subkernel(self, subkernel): - self._subkernels.append(subkernel) - - def unique_name(self, prefix): - return self._name_generator(prefix) - @contextlib.contextmanager def within_inames(self, inames) -> None: orig_within_inames = self._within_inames @@ -288,15 +301,332 @@ def within_inames(self, inames) -> None: # FIXME, bad API but it is context-dependent def set_temporary_shapes(self, shapes): self._temporary_shapes = shapes + + def lower_buffer_access( + self, + buffer_view: pyop3.buffer.IndexedBuffer, + layouts, + iname_maps, + loop_indices, + *, + intent, + ) -> pym.Expression: + name_in_kernel = self.add_buffer(buffer_view, intent) - @property - def _depends_on(self): - return frozenset({self._last_insn_id}) - {None} + buffer = buffer_view.buffer + if isinstance(buffer, PetscMatBuffer): + buffer = buffer_view.denested.getPythonself().buffer + + # At this point we know how to address each axis of the underlying buffer. + # This is sufficient to address a flat buffer, but for a buffer with more + # dimensions (i.e. a matrix) we have to do more work. As an example + # consider accessing a 2D buffer with shape (5, 5) using layout functions + # '2*i+1' and 'j+2' for the rows and columns respectively, where + # '0<=i<2' and '0<=j<3'. The offset expression that we want from this is: + # + # 5*(2*i+1) + (j+2) + # + # Which we can only determine from knowing the underlying buffer shape. + offset_expr = sum( + stride * self.lower_expr(layout, [iname_map], loop_indices) + for stride, layout, iname_map in zip( + utils.strides(buffer.shape), + layouts, + iname_maps, + strict=True + ) + ) + + # Add some leading zeros to make loopy happy + indices = self.maybe_multiindex(buffer_view.buffer, offset_expr) - def _add_instruction(self, insn): - self._instructions.append(insn) - self._last_insn_id = insn.id + subscript = pym.subscript(pym.var(name_in_kernel), indices) + if self.propagate_negatives and intent == READ: + idx = indices[-1] # only the final index has meaning + is_negative = pym.primitives.Comparison(idx, "<", 0) + return pym.primitives.If(is_negative, -1, subscript) + else: + return subscript + + def maybe_multiindex(self, buffer, offset_expr): + # hack to handle the facbuffer.t that temporaries can have shape but we want to + # linearly index it here + if buffer in self._temporary_shapes: + shape = self._temporary_shapes[buffer] + rank = len(shape) + extra_indices = (0,) * (rank - 1) + + # also has to be a scalar, not an expression + temp_offset_name = self.add_temporary("j") + temp_offset_var = pym.var(temp_offset_name) + self.add_assignment(temp_offset_var, offset_expr) + indices = extra_indices + (temp_offset_var,) + else: + indices = (offset_expr,) + return indices + + # NOTE: This could probably be refactored + def add_leaf_assignment( + self, + assignee, + expression, + assignment_type, + paths, + iname_replace_maps, + loop_indices, + ): + intent = assignment_type_as_intent(assignment_type) + lexpr = self.lower_expr( + assignee, + iname_replace_maps, + loop_indices, + intent=intent, + paths=paths, + ) + rexpr = self.lower_expr( + expression, + iname_replace_maps, + loop_indices, + paths=paths, + ) + + match assignment_type: + case AssignmentType.WRITE: + pass + case AssignmentType.INC: + rexpr = lexpr + rexpr + case AssignmentType.MAX: + rexpr = pym.Variable("max")(lexpr, rexpr) + case AssignmentType.MIN: + rexpr = pym.Variable("min")(lexpr, rexpr) + case _: + raise NotImplementedError + + if self.mask_array_accesses: + # a[off_a] = off_b < 0 ? a[off_a] : b[off_b] + offset_expr = _min_subscript_offset(rexpr) + # if there are no subcripts then the mask is pointless + if offset_expr is not None: + cond = pym.primitives.Comparison(offset_expr, "<", 0) + rexpr = pym.primitives.If(cond, lexpr, rexpr) + + self.add_assignment(lexpr, rexpr) + + def lower_expr(self, expr, iname_maps, loop_indices, *, intent=READ, paths=None) -> pym.Expression: + return self._lower_expr(expr, iname_maps, loop_indices, intent=intent, paths=paths) + + # TODO: use overloadedexpressionevaluator + @functools.singledispatchmethod + def _lower_expr(self, obj: Any, /, *args, **kwargs) -> pym.Expression: + raise TypeError(f"No handler defined for {type(obj).__name__}") + + + @_lower_expr.register(numbers.Number) + def _(self, num: numbers.Number, /, *args, **kwargs) -> numbers.Number: + return num + + + @_lower_expr.register(pyop3.expr.Add) + def _(self, add: pyop3.expr.Add, /, *args, **kwargs) -> pym.Expression: + return self._lower_expr(add.a, *args, **kwargs) + self._lower_expr(add.b, *args, **kwargs) + + + @_lower_expr.register(pyop3.expr.Sub) + def _(self, sub: pyop3.expr.Sub, /, *args, **kwargs) -> pym.Expression: + return self._lower_expr(sub.a, *args, **kwargs) - self._lower_expr(sub.b, *args, **kwargs) + + + @_lower_expr.register(pyop3.expr.Mul) + def _(self, mul: pyop3.expr.Mul, /, *args, **kwargs) -> pym.Expression: + return self._lower_expr(mul.a, *args, **kwargs) * self._lower_expr(mul.b, *args, **kwargs) + + + @_lower_expr.register(pyop3.expr.Modulo) + def _(self, mod: pyop3.expr.Mod, /, *args, **kwargs) -> pym.Expression: + return self._lower_expr(mod.a, *args, **kwargs) % self._lower_expr(mod.b, *args, **kwargs) + + + @_lower_expr.register(pyop3.expr.Or) + def _(self, or_: pyop3.expr.Or, /, *args, **kwargs) -> pym.Expression: + return pym.primitives.LogicalOr((self._lower_expr(or_.a, *args, **kwargs), self._lower_expr(or_.b, *args, **kwargs))) + + + @_lower_expr.register(pyop3.expr.Neg) + def _(self, neg: pyop3.expr.Neg, /, *args, **kwargs) -> pym.Expression: + return -self._lower_expr(neg.a, *args, **kwargs) + + + @_lower_expr.register(pyop3.expr.FloorDiv) + def _(self, neg: pyop3.expr.Neg, /, *args, **kwargs) -> pym.Expression: + return self._lower_expr(neg.a, *args, **kwargs) // self._lower_expr(neg.b, *args, **kwargs) + + + @_lower_expr.register(pyop3.expr.Comparison) + def _(self, cond, /, *args, **kwargs) -> pym.Expression: + return pym.primitives.Comparison( + self._lower_expr(cond.a, *args, **kwargs), + cond._symbol, + self._lower_expr(cond.b, *args, **kwargs), + ) + + @_lower_expr.register(pyop3.expr.AxisVar) + def _(self, axis_var: pyop3.expr.AxisVar, /, iname_maps, *args, **kwargs) -> pym.Expression: + return utils.just_one(iname_maps)[axis_var.axis.label] + + + @_lower_expr.register(pyop3.expr.LoopIndexVar) + def _(self, loop_var: pyop3.expr.LoopIndexVar, /, iname_maps, loop_indices, *args, **kwargs) -> pym.Expression: + return loop_indices[(loop_var.loop_index.id, loop_var.axis.label)] + + + @_lower_expr.register(pyop3.expr.ScalarBufferExpression) + def _( + self, + expr: pyop3.expr.ScalarBufferExpression, + /, + iname_maps, + loop_indices, + *, + intent, + **kwargs, + ) -> pym.ExpressionNode: + return self.lower_buffer_access(expr.buffer_view, [0], iname_maps, loop_indices, intent=intent) + + + @_lower_expr.register(pyop3.expr.LinearDatBufferExpression) + def _(self, expr: pyop3.expr.LinearDatBufferExpression, /, iname_maps, loop_indices, *, intent, **kwargs) -> pym.Expression: + return self.lower_buffer_access(expr.buffer_view, [expr.layout], iname_maps, loop_indices, intent=intent) + + + @_lower_expr.register(pyop3.expr.NonlinearDatBufferExpression) + def _(self, expr: pyop3.expr.NonlinearDatBufferExpression, /, iname_maps, loop_indices, *, intent, paths, **kwargs) -> pym.Expression: + path = utils.just_one(paths) + return self.lower_buffer_access( + expr.buffer_view, + [expr.layouts[path]], + iname_maps, + loop_indices, + intent=intent, + ) + + + @_lower_expr.register(pyop3.expr.MatPetscMatBufferExpression) + def _(self, mat_expr: pyop3.expr.MatPetscMatBufferExpression, /, iname_maps, loop_indices, *, intent, paths) -> pym.Expression: + row_path, column_path = paths + layouts = ( + mat_expr.row_layout.linearize(row_path), + mat_expr.column_layout.linearize(column_path), + ) + return self.lower_buffer_access( + mat_expr.buffer_view, + layouts, + iname_maps, + loop_indices, + intent=intent, + ) + + + @_lower_expr.register(pyop3.expr.MatArrayBufferExpression) + def _(self, expr: pyop3.expr.MatArrayBufferExpression, /, iname_maps, loop_indices, *, intent, paths) -> pym.Expression: + row_path, column_path = paths + layouts = (expr.row_layouts[row_path], expr.column_layouts[column_path]) + return self.lower_buffer_access( + expr.buffer_view, + layouts, + iname_maps, + loop_indices, + intent=intent, + ) + + def finalize_kernel(self, function_name, compiler_parameters): + preambles = [ + ("20_debug", "#include "), # dont always inject + ("30_petsc", "#include "), # perhaps only if petsc callable used? + ] + + # add a no-op instruction touching all of the kernel arguments so they are + # not silently dropped + noop = lp.CInstruction( + (), + "", + read_variables=frozenset({a.name for a in self.arguments}), + within_inames=frozenset(), + within_inames_is_final=True, + depends_on=self._depends_on, + ) + self._instructions.append(noop) + + translation_unit = lp.make_kernel( + self.domains, + self.instructions, + self.arguments, + name=function_name, + target=LOOPY_TARGET, + lang_version=LOOPY_LANG_VERSION, + preambles=preambles, + ) + translation_unit = lp.merge((translation_unit, *self.subkernels)) + + entrypoint = translation_unit.default_entrypoint + if compiler_parameters.add_likwid_markers: + entrypoint = with_likwid_markers(entrypoint) + if compiler_parameters.add_petsc_event: + entrypoint = with_petsc_event(entrypoint) + if compiler_parameters.attach_debugger: + entrypoint = with_attach_debugger(entrypoint) + translation_unit = translation_unit.with_kernel(entrypoint) + + return translation_unit.with_kernel(entrypoint) + + @functools.singledispatchmethod + def register_extent(self, obj: Any, *args, **kwargs): + raise TypeError(f"No handler defined for {type(obj).__name__}") + + + @register_extent.register(numbers.Integral) + def _(self, num: numbers.Integral, *args, **kwargs): + return num + + @register_extent.register(pyop3.expr.Expression) + def _(self, expr: pyop3.expr.Expression, inames, loop_indices): + pym_expr = self.lower_expr(expr, [inames], loop_indices) + extent_name = self.add_temporary("p") + self.add_assignment(pym.var(extent_name), pym_expr) + return extent_name + +class _MinSubscriptOffsetMapper(pym.mapper.IdentityMapper): + + def __init__(self): + self.subscript_found = False + super().__init__() + + def map_sum(self, expr): + assert len(expr.children) == 2 + a, b = map(self.rec, expr.children) + return pym.primitives.If(pym.primitives.Comparison(a, "<", b), a, b) + + def map_subscript(self, expr): + self.subscript_found = True + # do not recurse + return utils.just_one(expr.index_tuple) + + +def _min_subscript_offset(expr: pym.ExpressionNode) -> pym.ExpressionNode | None: + """Return an expression for the minimum subscript offset in an expression. + + This is important because we sometimes need to be able to check if we are + indexing with negative values (and hence might want to mask the access). + + If no subscripts are found then `None` is returned. + + """ + mapper = _MinSubscriptOffsetMapper() + mapped_expr = mapper(expr) + if mapper.subscript_found: + return mapped_expr + else: + return None class LACallable(lp.ScalarCallable, metaclass=abc.ABCMeta): """ @@ -310,10 +640,12 @@ def __init__(self, name=None, arg_id_to_dtype=None, assert name == self.name name_in_target = name_in_target if name_in_target else self.name - super().__init__(self.name, - arg_id_to_dtype=arg_id_to_dtype, - arg_id_to_descr=arg_id_to_descr, - name_in_target=name_in_target) + super().__init__( + self.name, + arg_id_to_dtype=arg_id_to_dtype, + arg_id_to_descr=arg_id_to_descr, + name_in_target=name_in_target + ) @abc.abstractproperty def name(self): @@ -404,830 +736,3 @@ class SolveCallable(LACallable): def generate_preambles(self, target): assert isinstance(target, type(target)) yield ("solve", solve_preamble) - - -def _compile_static_hashkey(op: PreprocessedOperation, compiler_parameters: ParsedCompilerParameters) -> Hashable: - return (op.disk_cache_key, compiler_parameters, pyop3.config) - - -# NOTE: Some of this code is not specific to loopy, could be refactored -# This is generally a bit nasty and abstraction breaking because it relies on attrs -# of the InstructionExecutionContext -@pyop3.cache.memory_and_disk_cache( - hashkey=_compile_static_hashkey, - get_comm=lambda op, *a, **kw: op.comm, -) -def _compile_static(op: InstructionExecutionContext, compiler_parameters: ParsedCompilerParameters) -> tuple: - """Compile the operation without regard for specific data values. - - This function is therefore suitable for disk caching. - - Returns - ------- - TU - datamap - - """ - insn = op.preprocess() - function_name = "pyop3_loop" # TODO: Provide as kwarg - - if isinstance(insn, InstructionList): - cs_expr = insn.instructions - else: - cs_expr = (insn,) - - context = LoopyCodegenContext( - propagate_negatives=compiler_parameters.propagate_negatives, - mask_array_accesses=compiler_parameters.mask_array_accesses, - ) - # NOTE: so I think LoopCollection is a better abstraction here - don't want to be - # explicitly dealing with contexts at this point. Can always sniff them out again. - # for context, ex in cs_expr: - for ex in cs_expr: - # ex = expand_implicit_pack_unpack(ex) - - # add external loop indices as kernel arguments - # FIXME: removed because cs_expr needs to sniff the context now - loop_indices = {} - - for e in pyop3.collections.as_tuple(ex): # TODO: get rid of this loop - # context manager? - context.set_temporary_shapes(_collect_temporary_shapes(e)) - _compile(e, loop_indices, context) - - if not context.buffer_intents: - raise pyop3.exceptions.EffectlessComputationException( - "The generated kernel does not modify any global data, this may indicate that something has gone wrong" - ) - - # add a no-op instruction touching all of the kernel arguments so they are - # not silently dropped - noop = lp.CInstruction( - (), - "", - read_variables=frozenset({a.name for a in context.arguments}), - within_inames=frozenset(), - within_inames_is_final=True, - depends_on=context._depends_on, - ) - context._instructions.append(noop) - - preambles = [ - ("20_debug", "#include "), # dont always inject - ("30_petsc", "#include "), # perhaps only if petsc callable used? - ] - - translation_unit = lp.make_kernel( - context.domains, - context.instructions, - context.arguments, - name=function_name, - target=LOOPY_TARGET, - lang_version=LOOPY_LANG_VERSION, - preambles=preambles, - ) - translation_unit = lp.merge((translation_unit, *context.subkernels)) - - entrypoint = translation_unit.default_entrypoint - if compiler_parameters.add_likwid_markers: - entrypoint = with_likwid_markers(entrypoint) - if compiler_parameters.add_petsc_event: - entrypoint = with_petsc_event(entrypoint) - if compiler_parameters.attach_debugger: - entrypoint = with_attach_debugger(entrypoint) - translation_unit = translation_unit.with_kernel(entrypoint) - - # Extra information needed by the code executor - kernel_name_to_buffer_info = utils.invert_mapping(context.kernel_names) - buffer_intents = context.buffer_intents - - # Replace buffers with their indices, dropping any temporaries. Also - # match the calling order for the kernel. - kernel_name_to_global_buffer_info = {} - global_buffer_intents = {} - for kernel_arg in translation_unit.default_entrypoint.args: - buf_view = kernel_name_to_buffer_info[kernel_arg.name] - buf_index = op.preprocessed_buffers.index(buf_view.buffer) - kernel_name_to_global_buffer_info[kernel_arg.name] = (buf_index, buf_view.nest_indices) - global_buffer_intents[buf_index] = buffer_intents[buf_view.buffer] - - return translation_unit, kernel_name_to_global_buffer_info, global_buffer_intents - - -# put into a class in transform.py? -@functools.singledispatch -def _collect_temporary_shapes(expr): - raise TypeError(f"No handler defined for {type(expr).__name__}") - - -@_collect_temporary_shapes.register(InstructionList) -def _(insn_list: InstructionList, /) -> idict: - return utils.merge_dicts(_collect_temporary_shapes(insn) for insn in insn_list) - - -@_collect_temporary_shapes.register(Loop) -def _(loop: Loop, /): - shapes = {} - for stmt in loop.statements: - for temp, shape in _collect_temporary_shapes(stmt).items(): - if shape is None: - continue - if temp in shapes: - assert shapes[temp] == shape - else: - shapes[temp] = shape - return shapes - - -@_collect_temporary_shapes.register(AbstractAssignment) -@_collect_temporary_shapes.register(NullInstruction) -@_collect_temporary_shapes.register(Exscan) # assume we are fine -def _(assignment: AbstractAssignment, /) -> idict: - return idict() - - -# NOTE: This is a bit of a misnomer, we care about the shapes of things that -# we give to loopy, not just temporaries per se. We should be able to detect -# everything from the local kernels -@_collect_temporary_shapes.register -def _(call: StandaloneCalledFunction): - return idict( - { - arg.buffer: lp_arg.shape - for lp_arg, arg in zip( - call.function.code.default_entrypoint.args, call.arguments, strict=True - ) - if isinstance(lp_arg, lp.ArrayArg) - } - ) - - -@functools.singledispatch -def _compile(expr: Any, loop_indices, ctx: LoopyCodegenContext) -> None: - raise TypeError(f"No handler defined for {type(expr).__name__}") - - -@_compile.register(NullInstruction) -def _(null: NullInstruction, *args, **kwargs): - pass - - -@_compile.register(InstructionList) -def _(insn_list: InstructionList, /, loop_indices, ctx) -> None: - for insn in insn_list: - _compile(insn, loop_indices, ctx) - - -@_compile.register(Loop) -def _( - loop, - loop_indices, - codegen_context: LoopyCodegenContext, -) -> None: - parse_loop_properly_this_time( - loop, - loop.index.iterset, - loop_indices, - codegen_context, - ) - - -def parse_loop_properly_this_time( - loop, - axis_tree, - loop_indices, - codegen_context, - *, - axis=None, - path=None, - iname_map=None, -) -> None: - if axis_tree is UNIT_AXIS_TREE: - # NOTE: might need an expression here sometimes - for statement in loop.statements: - _compile( - statement, - # loop_indices | dict(loop_exprs), - loop_indices, - codegen_context, - ) - return - - if utils.strictly_all(x is None for x in {axis, path, iname_map}): - axis = axis_tree.root - path = idict() - iname_map = idict() - - for component in axis.components: - path_ = path | {axis.label: component.label} - - if axis_tree.linearize(path_, partial=True).size == 0: - continue - elif component.size != 1: - iname = codegen_context.unique_name("i") - domain_var = register_extent( - component.size, - iname_map, - loop_indices, - codegen_context, - ) - codegen_context.add_domain(iname, domain_var) - iname_replace_map_ = iname_map | {axis.label: pym.var(iname)} - within_inames = frozenset({iname}) - else: - iname_replace_map_ = iname_map | {axis.label: 0} - within_inames = set() - - with codegen_context.within_inames(within_inames): - if subaxis := axis_tree.node_map[path_]: - parse_loop_properly_this_time( - loop, - axis_tree, - loop_indices, - codegen_context, - axis=subaxis, - path=path_, - iname_map=iname_replace_map_, - ) - else: - loop_indices |= idict({ - (loop.index.id, axis_label): iname - for axis_label, iname in iname_replace_map_.items() - }) - for statement in loop.statements: - _compile( - statement, - loop_indices, - codegen_context, - ) - - -@_compile.register -def _(call: StandaloneCalledFunction, loop_indices, context: LoopyCodegenContext) -> None: - subarrayrefs = {} - loopy_args = call.function.code.default_entrypoint.args - for loopy_arg, arg, spec in zip(loopy_args, call.arguments, call.argspec, strict=True): - name_in_kernel = context.add_buffer(arg.buffer_view, spec.intent) - if isinstance(loopy_arg, lp.ArrayArg): - # array arguments to an inner kernel require all strides to be defined - indices = [] - for s in loopy_arg.shape: - iname = context.unique_name("i") - context.add_domain(iname, s) - indices.append(pym.var(iname)) - indices = tuple(indices) - subarrayrefs[arg] = lp.symbolic.SubArrayRef( - indices, pym.var(name_in_kernel)[indices] - ) - else: - assert isinstance(loopy_arg, lp.ValueArg) - subarrayrefs[arg] = pym.var(name_in_kernel) - - assignees = tuple( - subarrayrefs[arg] - for arg, spec in zip(call.arguments, call.argspec, strict=True) - if spec.intent in {WRITE, RW, INC, MIN_RW, MIN_WRITE, MAX_RW, MAX_WRITE} - ) - expression = pym.primitives.Call( - pym.var(call.function.code.default_entrypoint.name), - tuple( - subarrayrefs[arg] - for arg, spec in zip(call.arguments, call.argspec, strict=True) - if spec.intent in {READ, RW, INC, MIN_RW, MAX_RW} - ), - ) - - context.add_function_call(assignees, expression) - subkernel = call.function.code.with_entrypoints(frozenset()) - context.add_subkernel(subkernel) - - -@_compile.register(NonEmptyArrayAssignment) -def parse_assignment(assignment: NonEmptyArrayAssignment, loop_indices, context: CodegenContext): - if any(isinstance(arg, pyop3.expr.MatPetscMatBufferExpression) for arg in assignment.arguments): - _compile_petsc_mat(assignment, loop_indices, context) - else: - compile_array_assignment( - assignment, - loop_indices, - context, - assignment.axis_trees, - ) - - -def _compile_petsc_mat(assignment: NonEmptyArrayAssignment, loop_indices, context) -> None: - # We need to know whether the matrix is the assignee or not because we need - # to know whether to put MatGetValues or MatSetValues - if isinstance(assignment.assignee.buffer_view.buffer, PetscMatBuffer): - mat = assignment.assignee - expr = assignment.expression - setting_mat_values = True - else: - mat = assignment.expression - expr = assignment.assignee - setting_mat_values = False - - - row_axis_tree, column_axis_tree = assignment.axis_trees - - assert isinstance(expr, pyop3.expr.BufferExpression) - - # now emit the right line of code, this should properly be a lp.ScalarCallable - # https://petsc.org/release/manualpages/Mat/MatGetValuesLocal/ - mat_name = context.add_buffer(mat.buffer_view, assignment_type_as_intent(assignment.assignment_type)) - array_name = context.add_buffer(expr.buffer_view, READ) - - rsize = row_axis_tree.size - csize = column_axis_tree.size - - # these sizes can be expressions that need evaluating - rsize_var = register_extent( - rsize, - {}, - loop_indices, - context, - ) - - csize_var = register_extent( - csize, - {}, - loop_indices, - context, - ) - - # convert the generic expressions to - # for example: - # - # map0[3*i0 + i1] - # map0[3*i0 + i2 + 3] - # - # to the shared top-level layout: - # - # map0[3*i0] - # - # which is what Mat{Get,Set}Values() needs. - layout_exprs = [] - for layout in [mat.row_layout, mat.column_layout]: - subst_sublayout = layout.layouts[idict()] - subst_layout = pyop3.expr.LinearDatBufferExpression(layout.buffer, subst_sublayout) - layout_expr = lower_expr(subst_layout, ((),), loop_indices, context) - layout_exprs.append(layout_expr) - irow, icol = layout_exprs - - # FIXME: - blocked = False - - # hacky - myargs = [ - assignment, mat_name, array_name, rsize_var, csize_var, irow, icol, blocked - ] - if setting_mat_values: - match assignment.assignment_type: - case AssignmentType.WRITE: - call_str = _petsc_mat_store(*myargs) - case AssignmentType.INC: - call_str = _petsc_mat_add(*myargs) - case _: - raise AssertionError - else: - call_str = _petsc_mat_load(*myargs) - - context.add_cinstruction(call_str) - - -def _petsc_mat_load(assignment, mat_name, array_name, nrow, ncol, irow, icol, blocked): - if blocked: - return f"MatGetValuesBlockedLocal({mat_name}, {nrow}, &({irow}), {ncol}, &({icol}), &({array_name}[0]));" - else: - return f"MatGetValuesLocal({mat_name}, {nrow}, &({irow}), {ncol}, &({icol}), &({array_name}[0]));" - - -def _petsc_mat_store(assignment, mat_name, array_name, nrow, ncol, irow, icol, blocked): - if blocked: - return f"MatSetValuesBlockedLocal({mat_name}, {nrow}, &({irow}), {ncol}, &({icol}), &({array_name}[0]), INSERT_VALUES);" - else: - return f"MatSetValuesLocal({mat_name}, {nrow}, &({irow}), {ncol}, &({icol}), &({array_name}[0]), INSERT_VALUES);" - - -def _petsc_mat_add(assignment, mat_name, array_name, nrow, ncol, irow, icol, blocked): - if blocked: - return f"MatSetValuesBlockedLocal({mat_name}, {nrow}, &({irow}), {ncol}, &({icol}), &({array_name}[0]), ADD_VALUES);" - else: - return f"MatSetValuesLocal({mat_name}, {nrow}, &({irow}), {ncol}, &({icol}), &({array_name}[0]), ADD_VALUES);" - -# TODO now I attach a lot of info to the context-free array, do I need to pass axes around? -def compile_array_assignment( - assignment, - loop_indices, - codegen_context, - axis_trees, - *, - iname_replace_maps=None, - # TODO document these under "Other Parameters" - axis_tree=None, - paths=None, -): - if paths is None: - paths = [] - if iname_replace_maps is None: - iname_replace_maps = [] - - if axis_tree is None: - axis_tree, *axis_trees = axis_trees - - paths += [idict()] - iname_replace_maps += [idict()] - - if axis_tree.is_empty or axis_tree is UNIT_AXIS_TREE or isinstance(axis_tree, IndexedAxisTree): - if axis_trees: - raise NotImplementedError("need to refactor code here") - - add_leaf_assignment( - assignment, - paths, - iname_replace_maps, - codegen_context, - loop_indices, - ) - return - - axis = axis_tree.node_map[paths[-1]] - - for component in axis.components: - new_paths = paths.copy() - new_paths[-1] = paths[-1] | {axis.label: component.label} - - # If the subtree below this is zero-sized then don't do anything - if axis_tree.linearize(new_paths[-1], partial=True).size == 0: - continue - elif component.size != 1: - iname = codegen_context.unique_name("i") - - extent_var = register_extent( - component.size, - iname_replace_maps[-1], - loop_indices, - codegen_context, - ) - codegen_context.add_domain(iname, extent_var) - new_iname_replace_maps = iname_replace_maps.copy() - new_iname_replace_maps[-1] = iname_replace_maps[-1] | {axis.label: pym.var(iname)} - within_inames = {iname} - else: - new_iname_replace_maps = iname_replace_maps.copy() - new_iname_replace_maps[-1] = iname_replace_maps[-1] | {axis.label: 0} - within_inames = set() - - with codegen_context.within_inames(within_inames): - if axis_tree.node_map[new_paths[-1]]: - compile_array_assignment( - assignment, - loop_indices, - codegen_context, - axis_trees, - iname_replace_maps=new_iname_replace_maps, - axis_tree=axis_tree, - paths=new_paths, - ) - elif axis_trees: - compile_array_assignment( - assignment, - loop_indices, - codegen_context, - axis_trees, - iname_replace_maps=new_iname_replace_maps, - axis_tree=None, - paths=new_paths, - ) - else: - add_leaf_assignment( - assignment, - new_paths, - new_iname_replace_maps, - codegen_context, - loop_indices, - ) - - -def add_leaf_assignment( - assignment, - paths, - iname_replace_maps, - codegen_context, - loop_indices, -): - intent = assignment_type_as_intent(assignment.assignment_type) - lexpr = lower_expr( - assignment.assignee, - iname_replace_maps, - loop_indices, - codegen_context, - intent=intent, - paths=paths, - ) - rexpr = lower_expr( - assignment.expression, - iname_replace_maps, - loop_indices, - codegen_context, - paths=paths, - ) - - match assignment.assignment_type: - case AssignmentType.WRITE: - pass - case AssignmentType.INC: - rexpr = lexpr + rexpr - case AssignmentType.MAX: - rexpr = pym.Variable("max")(lexpr, rexpr) - case AssignmentType.MIN: - rexpr = pym.Variable("min")(lexpr, rexpr) - case _: - raise NotImplementedError - - if codegen_context.mask_array_accesses: - # a[off_a] = off_b < 0 ? a[off_a] : b[off_b] - offset_expr = _min_subscript_offset(rexpr) - # if there are no subcripts then the mask is pointless - if offset_expr is not None: - cond = pym.primitives.Comparison(offset_expr, "<", 0) - rexpr = pym.primitives.If(cond, lexpr, rexpr) - - codegen_context.add_assignment(lexpr, rexpr) - - -@_compile.register(Exscan) -def _(exscan: Exscan, loop_indices, context) -> None: - if exscan.scan_type != "+": - raise NotImplementedError - domain_var = register_extent( - exscan.extent, - {}, - loop_indices, - context, - ) - iname = context.unique_name("i") - context.add_domain(iname, domain_var) - - iname_var = pym.var(iname) - iname_map = {exscan.scan_axis.label: pym.var(iname)} - - lexpr = lower_expr(exscan.assignee, [iname_map], loop_indices, context, intent=RW) - rexpr = lexpr + lower_expr(exscan.expression, [iname_map], loop_indices, context) - lexpr = pym.substitute(lexpr, {iname: iname_var+1}) - context.add_assignment(lexpr, rexpr) - - -def lower_expr(expr, iname_maps, loop_indices, ctx, *, intent=READ, paths=None) -> pym.Expression: - return _lower_expr(expr, iname_maps, loop_indices, ctx, intent=intent, paths=paths) - - -# TODO: use overloadedexpressionevaluator -@functools.singledispatch -def _lower_expr(obj: Any, /, *args, **kwargs) -> pym.Expression: - raise TypeError(f"No handler defined for {type(obj).__name__}") - - -@_lower_expr.register(numbers.Number) -def _(num: numbers.Number, /, *args, **kwargs) -> numbers.Number: - return num - - -@_lower_expr.register(pyop3.expr.Add) -def _(add: pyop3.expr.Add, /, *args, **kwargs) -> pym.Expression: - return _lower_expr(add.a, *args, **kwargs) + _lower_expr(add.b, *args, **kwargs) - - -@_lower_expr.register(pyop3.expr.Sub) -def _(sub: pyop3.expr.Sub, /, *args, **kwargs) -> pym.Expression: - return _lower_expr(sub.a, *args, **kwargs) - _lower_expr(sub.b, *args, **kwargs) - - -@_lower_expr.register(pyop3.expr.Mul) -def _(mul: pyop3.expr.Mul, /, *args, **kwargs) -> pym.Expression: - return _lower_expr(mul.a, *args, **kwargs) * _lower_expr(mul.b, *args, **kwargs) - - -@_lower_expr.register(pyop3.expr.Modulo) -def _(mod: pyop3.expr.Mul, /, *args, **kwargs) -> pym.Expression: - return _lower_expr(mod.a, *args, **kwargs) % _lower_expr(mod.b, *args, **kwargs) - - -@_lower_expr.register(pyop3.expr.Or) -def _(or_: pyop3.expr.Or, /, *args, **kwargs) -> pym.Expression: - return pym.primitives.LogicalOr((_lower_expr(or_.a, *args, **kwargs), _lower_expr(or_.b, *args, **kwargs))) - - -@_lower_expr.register(pyop3.expr.Neg) -def _(neg: pyop3.expr.Neg, /, *args, **kwargs) -> pym.Expression: - return -_lower_expr(neg.a, *args, **kwargs) - - -@_lower_expr.register(pyop3.expr.FloorDiv) -def _(neg: pyop3.expr.Neg, /, *args, **kwargs) -> pym.Expression: - return _lower_expr(neg.a, *args, **kwargs) // _lower_expr(neg.b, *args, **kwargs) - - -@_lower_expr.register(pyop3.expr.Comparison) -def _(cond, /, *args, **kwargs) -> pym.Expression: - return pym.primitives.Comparison( - _lower_expr(cond.a, *args, **kwargs), - cond._symbol, - _lower_expr(cond.b, *args, **kwargs), - ) - - -@_lower_expr.register(pyop3.expr.AxisVar) -def _(axis_var: pyop3.expr.AxisVar, /, iname_maps, *args, **kwargs) -> pym.Expression: - return utils.just_one(iname_maps)[axis_var.axis.label] - - -@_lower_expr.register(pyop3.expr.LoopIndexVar) -def _(loop_var: pyop3.expr.LoopIndexVar, /, iname_maps, loop_indices, *args, **kwargs) -> pym.Expression: - return loop_indices[(loop_var.loop_index.id, loop_var.axis.label)] - - -@_lower_expr.register(pyop3.expr.ScalarBufferExpression) -def _( - expr: pyop3.expr.ScalarBufferExpression, - /, - iname_maps, - loop_indices, - context, - *, - intent, - **kwargs, -) -> pym.ExpressionNode: - return lower_buffer_access(expr.buffer_view, [0], iname_maps, loop_indices, context, intent=intent) - - -@_lower_expr.register(pyop3.expr.LinearDatBufferExpression) -def _(expr: pyop3.expr.LinearDatBufferExpression, /, iname_maps, loop_indices, context, *, intent, **kwargs) -> pym.Expression: - return lower_buffer_access(expr.buffer_view, [expr.layout], iname_maps, loop_indices, context, intent=intent) - - -@_lower_expr.register(pyop3.expr.NonlinearDatBufferExpression) -def _(expr: pyop3.expr.NonlinearDatBufferExpression, /, iname_maps, loop_indices, context, *, intent, paths, **kwargs) -> pym.Expression: - path = utils.just_one(paths) - return lower_buffer_access( - expr.buffer_view, - [expr.layouts[path]], - iname_maps, - loop_indices, - context, - intent=intent, - ) - - -@_lower_expr.register(pyop3.expr.MatPetscMatBufferExpression) -def _(mat_expr: pyop3.expr.MatPetscMatBufferExpression, /, iname_maps, loop_indices, context, *, intent, paths) -> pym.Expression: - row_path, column_path = paths - layouts = ( - mat_expr.row_layout.linearize(row_path), - mat_expr.column_layout.linearize(column_path), - ) - return lower_buffer_access( - mat_expr.buffer_view, - layouts, - iname_maps, - loop_indices, - context, - intent=intent, - ) - - -@_lower_expr.register(pyop3.expr.MatArrayBufferExpression) -def _(expr: pyop3.expr.MatArrayBufferExpression, /, iname_maps, loop_indices, context, *, intent, paths) -> pym.Expression: - row_path, column_path = paths - layouts = (expr.row_layouts[row_path], expr.column_layouts[column_path]) - return lower_buffer_access( - expr.buffer_view, - layouts, - iname_maps, - loop_indices, - context, - intent=intent, - ) - - -def lower_buffer_access( - buffer_view: pyop3.buffer.IndexedBuffer, - layouts, - iname_maps, - loop_indices, - context, - *, - intent, -) -> pym.Expression: - name_in_kernel = context.add_buffer(buffer_view, intent) - - buffer = buffer_view.buffer - if isinstance(buffer, PetscMatBuffer): - buffer = buffer_view.denested.getPythonContext().buffer - - # At this point we know how to address each axis of the underlying buffer. - # This is sufficient to address a flat buffer, but for a buffer with more - # dimensions (i.e. a matrix) we have to do more work. As an example - # consider accessing a 2D buffer with shape (5, 5) using layout functions - # '2*i+1' and 'j+2' for the rows and columns respectively, where - # '0<=i<2' and '0<=j<3'. The offset expression that we want from this is: - # - # 5*(2*i+1) + (j+2) - # - # Which we can only determine from knowing the underlying buffer shape. - offset_expr = sum( - stride * lower_expr(layout, [iname_map], loop_indices, context) - for stride, layout, iname_map in zip( - utils.strides(buffer.shape), - layouts, - iname_maps, - strict=True - ) - ) - - # Add some leading zeros to make loopy happy - indices = maybe_multiindex(buffer_view.buffer, offset_expr, context) - - subscript = pym.subscript(pym.var(name_in_kernel), indices) - if context.propagate_negatives and intent == READ: - idx = indices[-1] # only the final index has meaning - is_negative = pym.primitives.Comparison(idx, "<", 0) - return pym.primitives.If(is_negative, -1, subscript) - else: - return subscript - - -def maybe_multiindex(buffer, offset_expr, context): - # hack to handle the facbuffer.t that temporaries can have shape but we want to - # linearly index it here - if buffer in context._temporary_shapes: - shape = context._temporary_shapes[buffer] - rank = len(shape) - extra_indices = (0,) * (rank - 1) - - # also has to be a scalar, not an expression - temp_offset_name = context.add_temporary("j") - temp_offset_var = pym.var(temp_offset_name) - context.add_assignment(temp_offset_var, offset_expr) - indices = extra_indices + (temp_offset_var,) - else: - indices = (offset_expr,) - - return indices - - -@_lower_expr.register(pyop3.expr.Conditional) -def _(cond: pyop3.expr.Conditional, /, *args, **kwargs) -> pym.Expression: - return pym.primitives.If(_lower_expr(cond.a, *args, **kwargs), _lower_expr(cond.b, *args, **kwargs), _lower_expr(cond.c, *args, **kwargs)) - - -class _MinSubscriptOffsetMapper(pym.mapper.IdentityMapper): - - def __init__(self): - self.subscript_found = False - super().__init__() - - def map_sum(self, expr): - assert len(expr.children) == 2 - a, b = map(self.rec, expr.children) - return pym.primitives.If(pym.primitives.Comparison(a, "<", b), a, b) - - def map_subscript(self, expr): - self.subscript_found = True - # do not recurse - return utils.just_one(expr.index_tuple) - - -def _min_subscript_offset(expr: pym.ExpressionNode) -> pym.ExpressionNode | None: - """Return an expression for the minimum subscript offset in an expression. - - This is important because we sometimes need to be able to check if we are - indexing with negative values (and hence might want to mask the access). - - If no subscripts are found then `None` is returned. - - """ - mapper = _MinSubscriptOffsetMapper() - mapped_expr = mapper(expr) - if mapper.subscript_found: - return mapped_expr - else: - return None - - -@functools.singledispatch -def register_extent(obj: Any, *args, **kwargs): - raise TypeError(f"No handler defined for {type(obj).__name__}") - - -@register_extent.register(numbers.Integral) -def _(num: numbers.Integral, *args, **kwargs): - return num - - -@register_extent.register(pyop3.expr.Expression) -def _(expr: pyop3.expr.Expression, inames, loop_indices, context): - pym_expr = lower_expr(expr, [inames], loop_indices, context) - extent_name = context.add_temporary("p") - context.add_assignment(pym.var(extent_name), pym_expr) - return extent_name diff --git a/pyop3/lower/transform.py b/pyop3/lower/transform.py index 65c1bc2046..ccfd832404 100644 --- a/pyop3/lower/transform.py +++ b/pyop3/lower/transform.py @@ -1,7 +1,18 @@ import textwrap import loopy as lp +import functools +from immutabledict import immutabledict as idict + +from pyop3.insn.base import ( + AbstractAssignment, + Exscan, + InstructionList, + Loop, + NullInstruction, + StandaloneCalledFunction +) def with_likwid_markers(knl): """ @@ -62,3 +73,44 @@ def with_attach_debugger(kernel): *(insn.copy(depends_on=insn.depends_on | {debug_insn.id}) for insn in kernel.instructions), ) return kernel.copy(instructions=insns) + +# NOTE: Make this overloaded function into class in transform.py +# Only issue may be loopy-specific standalone_function overloading. +@functools.singledispatch +def _collect_temporary_shapes(expr): + raise TypeError(f"No handler defined for {type(expr).__name__}") + +@_collect_temporary_shapes.register(InstructionList) +def _(insn_list): + return utils.merge_dicts(_collect_temporary_shapes(insn) for insn in insn_list) + +@_collect_temporary_shapes.register(Loop) +def _(loop): + shapes = {} + for stmt in loop.statements: + for temp, shape in _collect_temporary_shapes(stmt).items(): + if shape is None: + continue + if temp in shapes: + assert shapes[temp] == shape + else: + shapes[temp] = shape + return shapes + +@_collect_temporary_shapes.register(AbstractAssignment) +@_collect_temporary_shapes.register(NullInstruction) +@_collect_temporary_shapes.register(Exscan) +def _(assignment: AbstractAssignment, /) -> idict: + return idict() + +@_collect_temporary_shapes.register +def _(call: StandaloneCalledFunction): + return idict( + { + arg.buffer: lp_arg.shape + for lp_arg, arg in zip( + call.function.code.default_entrypoint.args, call.arguments, strict=True + ) + if isinstance(lp_arg, lp.ArrayArg) + } + ) diff --git a/requirements-build.txt b/requirements-build.txt index 811c170c0f..2d2cdf2044 100644 --- a/requirements-build.txt +++ b/requirements-build.txt @@ -13,4 +13,4 @@ setuptools>=77.0.3 # Transitive build dependencies hatchling meson-python -scikit_build_core +scikit_build_core \ No newline at end of file