From efe4ccb8db28a6e53913bfccd429a7f679db34fa Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Fri, 19 Jun 2026 12:21:05 +0100 Subject: [PATCH 01/30] sample mlir generation and runnable --- run_compiled_mlir.py | 52 +++++++++++++++++++++++ sample.py | 98 ++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 150 insertions(+) create mode 100644 run_compiled_mlir.py create mode 100644 sample.py diff --git a/run_compiled_mlir.py b/run_compiled_mlir.py new file mode 100644 index 0000000000..6905b37c82 --- /dev/null +++ b/run_compiled_mlir.py @@ -0,0 +1,52 @@ +import ctypes +import numpy as np +import cupy as cp + +from ctypes import c_void_p, c_longlong, Structure + +class MemRefDescriptor(Structure): + _fields_ = [ + ("allocated", c_void_p), + ("aligned", c_void_p), + ("offset", c_longlong), + ("shape", c_longlong * 1), + ("stride", c_longlong * 1), + ] + +def numpy_to_memref(arr): + if not arr.flags["C_CONTIGUOUS"]: + arr = np.ascontiguousarray(arr) + + desc = MemRefDescriptor() + desc.allocated = arr.ctypes.data_as(c_void_p) + desc.aligned = desc.allocated + desc.offset = 0 + desc.shape[0] = arr.shape[0] + desc.stride[0] = 1 + + return desc + + +if __name__ == "__main__": + lib = ctypes.CDLL("./liboutput.dylib") + + array_add = lib._mlir_ciface_add + array_add.argtypes = [ + ctypes.POINTER(MemRefDescriptor) + ] * 3 + + size = 8 + a = np.ones(size, dtype=np.float64) + b = np.ones(size, dtype=np.float64) * 2 + c = np.zeros(size, dtype=np.float64) + + a_desc = numpy_to_memref(a) + b_desc = numpy_to_memref(b) + c_desc = numpy_to_memref(c) + + array_add(ctypes.byref(a_desc), ctypes.byref(b_desc), ctypes.byref(c_desc)) + + expected = a + b + np.testing.assert_array_almost_equal(c, expected) + print("Array addition successful!") + print(f"First few elements: {c[:5]}") diff --git a/sample.py b/sample.py new file mode 100644 index 0000000000..338190d9ce --- /dev/null +++ b/sample.py @@ -0,0 +1,98 @@ +# Basic functionality +# Make MLIR in xDSL that adds two unranked tensor arrays +# Take this MLIR and lower to LLVM in xDSL +# Compile JIT with llvmlite? + +import sys +from xdsl.dialects import arith, func, memref, scf, tensor, linalg +from xdsl.dialects.arith import ConstantOp, AddfOp +from xdsl.dialects.tensor import DimOp, EmptyOp +from xdsl.dialects.builtin import ( + DYNAMIC_INDEX, + ModuleOp, + IndexType, + IntegerAttr, + f64, + TensorType, + ArrayAttr, + AffineMap, + AffineMapAttr, + AffineDimExpr, + UnitAttr +) +from xdsl.dialects.linalg import ( + IteratorTypeAttr, + YieldOp +) + +from xdsl.ir import Block, Region +from xdsl.context import Context +from xdsl.printer import Printer + +def build_array_add(n: int) -> ModuleOp: + tensor_type = TensorType(f64, [DYNAMIC_INDEX]) + + identity_1d = AffineMap(num_dims=1, num_symbols=0, results=(AffineDimExpr(0),)) + identity_attr = AffineMapAttr(identity_1d) + + parallel = IteratorTypeAttr.parallel() + + func_block = Block(arg_types=[tensor_type, tensor_type, tensor_type]) + a, b, out = func_block.args + + c0 = ConstantOp(IntegerAttr(0, IndexType())) + func_block.add_op(c0) + + body_block = Block(arg_types=[f64, f64, f64]) + x, y, _z = body_block.args + + add = AddfOp(x, y) + body_block.add_op(add) + + body_block.add_op(YieldOp(add.result)) + + generic = linalg.GenericOp( + inputs=[a, b], + outputs=[out], + body=Region([body_block]), + indexing_maps=[identity_attr, identity_attr, identity_attr], + iterator_types=[parallel], + result_types=[tensor_type], + ) + func_block.add_op(generic) + + func_block.add_op(func.ReturnOp()) + + func_region = Region([func_block]) + func_op = func.FuncOp( + "add", + ([tensor_type, tensor_type, tensor_type], []), + func_region, + ) + + func_op.attributes["llvm.emit_c_interface"] = UnitAttr() + + return ModuleOp([func_op]) + +def emit_mlir(module: ModuleOp) -> str: + """Return the MLIR text representation of a module.""" + import io + buf = io.StringIO() + Printer(stream=buf).print_op(module) + return buf.getvalue() + +if __name__ == "__main__": + n = int(sys.argv[1]) if len(sys.argv) > 1 else 8 + + # Register dialects so xDSL can verify the IR + ctx = Context() + ctx.load_dialect(func.Func) + ctx.load_dialect(arith.Arith) + ctx.load_dialect(linalg.Linalg) + ctx.load_dialect(tensor.Tensor) + + module = build_array_add(n) + mlir = emit_mlir(module) + print(mlir) + + From 39930f6f318acd25d5520e3a091f65b68ed91da6 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Fri, 19 Jun 2026 12:39:01 +0100 Subject: [PATCH 02/30] demo to satisfy --- gpu_offloading_demo.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) create mode 100644 gpu_offloading_demo.py diff --git a/gpu_offloading_demo.py b/gpu_offloading_demo.py new file mode 100644 index 0000000000..735f63dd4f --- /dev/null +++ b/gpu_offloading_demo.py @@ -0,0 +1,19 @@ +from firedrake import * +import pyop3 as op3 +import numpy as np, cupy as cp + +mesh = UnitSquareMesh(3,3) +V = FunctionSpace(mesh, "CG", 1) +f = Function(V).assign(10) +g = Function(V) + +gpu = op3.CUDAGPU() + +with op3.offloading(gpu): + g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile") + assert isinstance(g.dat.data_ro, cp.ndarray) # Device + assert (g.dat.data_ro == 20).all() + +assert isinstance(g.dat.data_ro, np.ndarray) # Host +assert (g.dat.data_ro == 20).all() + From 94e9c15db4fb14b8cb3b533cdbd6bfe3c2802992 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Wed, 8 Jul 2026 12:38:08 +0100 Subject: [PATCH 03/30] building plan and stepping through loopy.py to see inputs --- gpu_offloading_demo.py | 10 ++-- pyop3/insn/exec.py | 1 + pyop3/lower/codegen.py | 111 +++++++++++++++++++++++++++++++++++++++++ pyop3/lower/context.py | 4 ++ pyop3/lower/loopy.py | 10 +++- pyop3/lower/mlir.py | 6 +++ sample.py | 10 ++-- 7 files changed, 142 insertions(+), 10 deletions(-) create mode 100644 pyop3/lower/codegen.py create mode 100644 pyop3/lower/context.py create mode 100644 pyop3/lower/mlir.py diff --git a/gpu_offloading_demo.py b/gpu_offloading_demo.py index 735f63dd4f..ef52d35197 100644 --- a/gpu_offloading_demo.py +++ b/gpu_offloading_demo.py @@ -9,10 +9,12 @@ gpu = op3.CUDAGPU() -with op3.offloading(gpu): - g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile") - assert isinstance(g.dat.data_ro, cp.ndarray) # Device - assert (g.dat.data_ro == 20).all() +g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile") +print("DONE ASSIGN OPERATION\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n") +# with op3.offloading(gpu): +# g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile") +# assert isinstance(g.dat.data_ro, cp.ndarray) # Device +# assert (g.dat.data_ro == 20).all() assert isinstance(g.dat.data_ro, np.ndarray) # Host assert (g.dat.data_ro == 20).all() diff --git a/pyop3/insn/exec.py b/pyop3/insn/exec.py index ef53c05a30..3ddce126c9 100644 --- a/pyop3/insn/exec.py +++ b/pyop3/insn/exec.py @@ -235,6 +235,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. diff --git a/pyop3/lower/codegen.py b/pyop3/lower/codegen.py new file mode 100644 index 0000000000..17a6d89760 --- /dev/null +++ b/pyop3/lower/codegen.py @@ -0,0 +1,111 @@ +import abc +import collections +import contextlib +import ctypes +import dataclasses +import enum +import functools +import os +import numbers +import textwrap +import warnings +import weakref +from collections.abc import Mapping +from functools import cached_property +from typing import Any +from weakref import WeakValueDictionary + +# 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, *args, **kwargs: 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(check_negatives=compiler_parameters.check_negatives) + # 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 utils.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.global_buffers: + 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) + + kernel_to_buffer_names = utils.invert_mapping(context._kernel_names) + buffer_index_map = {} + for kernel_arg in entrypoint.args: + buffer_key = kernel_to_buffer_names[kernel_arg.name] + buffer_ref = context.global_buffers[buffer_key] + buffer_index = op.preprocessed_buffers.index(buffer_ref) + intent = context.global_buffer_intents[buffer_key] + buffer_index_map[kernel_arg.name] = (buffer_index, buffer_ref.nest_indices, intent) + + return translation_unit, buffer_index_map diff --git a/pyop3/lower/context.py b/pyop3/lower/context.py new file mode 100644 index 0000000000..84c1eb501e --- /dev/null +++ b/pyop3/lower/context.py @@ -0,0 +1,4 @@ +import abc + +class CodegenContext(abc.ABC): + pass diff --git a/pyop3/lower/loopy.py b/pyop3/lower/loopy.py index e522e7bbce..fe4b27978e 100644 --- a/pyop3/lower/loopy.py +++ b/pyop3/lower/loopy.py @@ -71,7 +71,6 @@ class CodegenContext(abc.ABC): pass - class LoopyCodegenContext(CodegenContext): def __init__(self, *, check_negatives): self.check_negatives = check_negatives @@ -115,6 +114,13 @@ def arguments(self) -> tuple: def subkernels(self) -> tuple: return tuple(self._subkernels) + def __str__(self) -> str: + 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 + def add_domain(self, iname, *args): nargs = len(args) if nargs == 1: @@ -444,6 +450,7 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed cs_expr = (insn,) context = LoopyCodegenContext(check_negatives=compiler_parameters.check_negatives) + breakpoint() # 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: @@ -458,6 +465,7 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed # context manager? context.set_temporary_shapes(_collect_temporary_shapes(e)) _compile(e, loop_indices, context) + breakpoint() if not context.global_buffers: raise pyop3.exceptions.EffectlessComputationException( diff --git a/pyop3/lower/mlir.py b/pyop3/lower/mlir.py new file mode 100644 index 0000000000..04a97b3cdf --- /dev/null +++ b/pyop3/lower/mlir.py @@ -0,0 +1,6 @@ + + +from pyop3.lower.context import CodegenContext + +class MLIRCodegenContext(CodegenContext): + pass diff --git a/sample.py b/sample.py index 338190d9ce..1e8c018eda 100644 --- a/sample.py +++ b/sample.py @@ -20,10 +20,10 @@ AffineDimExpr, UnitAttr ) -from xdsl.dialects.linalg import ( - IteratorTypeAttr, - YieldOp -) + + +from xdsl.dialects.linalg.attrs import IteratorTypeAttr +from xdsl.dialects.linalg.ops import YieldOp, GenericOp from xdsl.ir import Block, Region from xdsl.context import Context @@ -51,7 +51,7 @@ def build_array_add(n: int) -> ModuleOp: body_block.add_op(YieldOp(add.result)) - generic = linalg.GenericOp( + generic = GenericOp( inputs=[a, b], outputs=[out], body=Region([body_block]), From 285db41539928486dc6f4bad951afcbeb9367a3d Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Tue, 14 Jul 2026 16:42:12 +0100 Subject: [PATCH 04/30] designing make_kernel function --- assign_loop.txt | 42 +++++++++++++++ gpu_offloading_demo.py | 1 - pyop3/lower/context.py | 87 +++++++++++++++++++++++++++++-- pyop3/lower/loopy.py | 67 +++--------------------- pyop3/lower/mlir.py | 113 ++++++++++++++++++++++++++++++++++++++++- sample_2.py | 65 ++++++++++++++++++++++++ 6 files changed, 311 insertions(+), 64 deletions(-) create mode 100644 assign_loop.txt create mode 100644 sample_2.py diff --git a/assign_loop.txt b/assign_loop.txt new file mode 100644 index 0000000000..21b8e370be --- /dev/null +++ b/assign_loop.txt @@ -0,0 +1,42 @@ +******************************************************************************** +#include +#include +#include +#include + +void pyop3_loop(int64_t const *__restrict__ dat_0, int64_t const *__restrict__ dat_1, double *__restrict__ dat_2, double const *__restrict__ dat_3, int64_t const *__restrict__ dat_4, int64_t const *__restrict__ dat_5, int64_t const *__restrict__ dat_6, int64_t const *__restrict__ dat_7, int32_t const *__restrict__ idat_0, int32_t const *__restrict__ idat_1, int32_t const *__restrict__ idat_2, int32_t const *__restrict__ idat_3) +{ + int32_t p_0; + int32_t p_1; + int32_t p_2; + + p_0 = (int32_t) (dat_0[0] + dat_1[0]); + for (int32_t i_0 = 0; i_0 <= -1 + p_0; ++i_0) + { + } + p_1 = (int32_t) (dat_4[0] + dat_5[0]); + for (int32_t i_2 = 0; i_2 <= -1 + p_1; ++i_2) + dat_2[idat_0[idat_2[i_2]]] = 2.0 * dat_3[idat_0[idat_2[i_2]]]; + p_2 = (int32_t) (dat_6[0] + dat_7[0]); + for (int32_t i_3 = 0; i_3 <= -1 + p_2; ++i_3) + { + } + +} +******************************************************************************** +dat_0 (1) : [18] +dat_1 (1) : [0] +dat_2 (16) : [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] +dat_3 (16) : [10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10.] +dat_4 (1) : [16] +dat_5 (1) : [0] +dat_6 (1) : [33] +dat_7 (1) : [0] +idat_0 (67) : [ 0 0 0 0 0 1 2 3 3 3 3 4 4 4 4 5 5 5 5 6 6 6 6 7 + 7 7 7 8 8 8 8 9 9 9 9 9 9 10 10 10 10 11 11 11 11 12 12 12 + 12 13 13 13 13 13 13 13 13 14 14 14 14 15 15 15 15 15 15] +idat_1 (18) : [ 0 7 11 15 19 23 27 31 33 37 41 45 49 51 53 57 61 63] +idat_2 (16) : [ 4 5 6 10 14 18 22 26 30 36 40 44 48 56 60 66] +idat_3 (33) : [ 1 2 3 8 9 12 13 16 17 20 21 24 25 28 29 32 34 35 38 39 42 43 46 47 + 50 52 54 55 58 59 62 64 65] +******************************************************************************** diff --git a/gpu_offloading_demo.py b/gpu_offloading_demo.py index ef52d35197..f7133521cf 100644 --- a/gpu_offloading_demo.py +++ b/gpu_offloading_demo.py @@ -10,7 +10,6 @@ gpu = op3.CUDAGPU() g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile") -print("DONE ASSIGN OPERATION\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n") # with op3.offloading(gpu): # g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile") # assert isinstance(g.dat.data_ro, cp.ndarray) # Device diff --git a/pyop3/lower/context.py b/pyop3/lower/context.py index 84c1eb501e..664bd66eff 100644 --- a/pyop3/lower/context.py +++ b/pyop3/lower/context.py @@ -1,4 +1,85 @@ -import abc +from abc import ABC, abstractmethod -class CodegenContext(abc.ABC): - pass +from pyop3 import utils +from pyop3.insn.base import ( + Intent +) + +class CodegenContext(ABC): + def __init__(self, *, check_negatives): + self.check_negatives = check_negatives + + self._domains = [] + self._instructions = [] + self._arguments = [] + self._subkernels = [] + + self._name_generator = utils.UniqueNameGenerator() + + # buffer name -> name in kernel + self._kernel_names = {} + + # buffer name -> buffer + self.global_buffers = {} + self.global_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: arg.name)) + + @property + def subkernels(self) -> tuple: + return tuple(self._subkernels) + + def __str__(self) -> str: + 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 + + @abstractmethod + def add_domain(self, iname, *args): + pass + + @abstractmethod + def add_assignment(self, assigneee, expression, prefix="insn"): + pass + + @abstractmethod + def add_function_call(self, assignees, expression, prefix="insn"): + pass + + @abstractmethod + def add_buffer(self, buffer, intent: Intent | None = None) -> str: + pass + + @abstractmethod + def add_subkernel(self, subkernel): + pass + + @abstractmethod + def set_temporary_shapes(self, shapes): + pass + + def unique_name(self, prefix): + return self._name_generator(prefix) + + def _add_instruction(self, insn): + self._instructions.append(insn) + self._last_insn_id = insn.id + + @property + def _depends_on(self): + return frozenset({self._last_insn_id}) - {None} diff --git a/pyop3/lower/loopy.py b/pyop3/lower/loopy.py index fe4b27978e..178f4f00da 100644 --- a/pyop3/lower/loopy.py +++ b/pyop3/lower/loopy.py @@ -38,6 +38,8 @@ from pyop3.buffer import AbstractBuffer, ConcreteBuffer, PetscMatBuffer, ArrayBuffer, NullBuffer from pyop3.dtypes import IntType from pyop3.lower.transform import with_likwid_markers, with_petsc_event, with_attach_debugger +from pyop3.lower.context import CodegenContext +from pyop3.lower.mlir import MLIRCodegenContext # to remove from pyop3.insn.base import ( Intent, INC, @@ -67,60 +69,16 @@ LOOPY_TARGET = lp.CWithGNULibcTarget() LOOPY_LANG_VERSION = (2018, 2) - -class CodegenContext(abc.ABC): - pass - class LoopyCodegenContext(CodegenContext): def __init__(self, *, check_negatives): - self.check_negatives = check_negatives - - self._domains = [] - self._instructions = [] - self._arguments = [] - self._subkernels = [] + super().__init__(check_negatives=check_negatives) self._within_inames = frozenset() self._last_insn_id = None - self._name_generator = utils.UniqueNameGenerator() - - # buffer name -> name in kernel - self._kernel_names = {} - - # buffer name -> buffer - self.global_buffers = {} - self.global_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) - - @property - def instructions(self) -> tuple: - return tuple(self._instructions) - - @property - def arguments(self) -> tuple: - return tuple(sorted(self._arguments, key=lambda arg: arg.name)) - - @property - def subkernels(self) -> tuple: - return tuple(self._subkernels) - - def __str__(self) -> str: - 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 - def add_domain(self, iname, *args): nargs = len(args) if nargs == 1: @@ -238,7 +196,7 @@ def add_buffer(self, buffer, intent: Intent | None = None) -> str: name_in_kernel = self.unique_name("mat") loopy_arg = lp.ValueArg(name_in_kernel, dtype=pyop3.dtypes.OpaqueType("Mat")) - + self.global_buffers[buffer_key] = buffer self.global_buffer_intents[buffer_key] = intent self._arguments.append(loopy_arg) @@ -287,9 +245,6 @@ def add_opaque(self, opaque: OpaqueTerminal, intent) -> str: 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 @@ -301,13 +256,6 @@ def within_inames(self, inames) -> None: def set_temporary_shapes(self, shapes): self._temporary_shapes = shapes - @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 class LACallable(lp.ScalarCallable, metaclass=abc.ABCMeta): @@ -450,7 +398,8 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed cs_expr = (insn,) context = LoopyCodegenContext(check_negatives=compiler_parameters.check_negatives) - breakpoint() + mlir_context = MLIRCodegenContext(check_negatives=compiler_parameters.check_negatives) + # 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: @@ -465,7 +414,6 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed # context manager? context.set_temporary_shapes(_collect_temporary_shapes(e)) _compile(e, loop_indices, context) - breakpoint() if not context.global_buffers: raise pyop3.exceptions.EffectlessComputationException( @@ -489,6 +437,8 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed ("30_petsc", "#include "), # perhaps only if petsc callable used? ] + breakpoint() + translation_unit = lp.make_kernel( context.domains, context.instructions, @@ -870,7 +820,6 @@ def compile_array_assignment( continue elif component.size != 1: iname = codegen_context.unique_name("i") - extent_var = register_extent( component.size, iname_replace_maps[-1], diff --git a/pyop3/lower/mlir.py b/pyop3/lower/mlir.py index 04a97b3cdf..0d3a033f93 100644 --- a/pyop3/lower/mlir.py +++ b/pyop3/lower/mlir.py @@ -1,6 +1,117 @@ +from xdsl.dialects import arith, func, tensor, linalg +from xdsl.dialects.builtin import ( + ModuleOp, + IntegerAttr +) + +import pyop3 +from pyop3.buffer import AbstractBuffer, ConcreteBuffer, PetscMatBuffer, ArrayBuffer, NullBuffer +from pyop3.dtypes import IntType from pyop3.lower.context import CodegenContext +from pyop3.insn.base import ( + Intent +) + class MLIRCodegenContext(CodegenContext): - pass + + def __init__(self, *, check_negatives): + super().__init__(check_negatives=check_negatives) + + def add_domain(self, iname, *args): + nargs = len(args) + if nargs == 1: + start, stop = 0, args[0] + else: + assert nargs == 2 + start, stop = args[0], args[1] + self._domains.append((start, stop)) + + def add_assignment(self, assigneee, expression, prefix="insn"): + pass + + def add_function_call(self, assignees, expression, prefix="insn"): + pass + + def add_buffer(self, buffer, intent: Intent | None = None) -> str: + # TODO: This only works for np.ndarrays for development atm + + if isinstance(buffer, NullBuffer): + assert not buffer.nest_indices + + if buffer_key in self._kernel_names: + return self._kernel_names[buffer_key] + shape = self._temporary_shapes.get(buffer_key, (buffer.size,)) + assert isinstance(shape, tuple) and all(isinstance(s, numbers.Integral) for s in shape) + name_in_kernel = self.add_temporary("t", buffer.dtype, shape=shape) + else: + if intent is None: + raise ValueError("Global data must declare intent") + + if buffer_key in self._kernel_names: + if intent != self.global_buffer_intents[buffer_key]: + # We are accessing a buffer with different intents so have to + # pessimally claim RW access + self.global_buffer_intents[buffer_key] = RW + return self._kernel_names[buffer_key] + + if isinstance(buffer.handle, np.ndarray): + if isinstance(buffer.dtype, np.dtypes.IntDType): + name_in_kernel = self.unique_name("idat") + else: + name_in_kernel = self.unique_name("dat") + + # If the buffer is being passed straight through to a function then we + # have to make sure that the shapes match + shape = self._temporary_shapes.get(buffer_key, None) + # TODO: An equivalent of lp.GlobalArg is required here + # GlobalArg represents array, dtype, shape, address space (local or global variable) + iter_arg = (name_in_kernel, buffer.dtype, shape) + else: + assert isinstance(buffer, PetscMatBuffer) + assert buffer.mat_type not in {"nest", "python"} + + name_in_kernel = self.unique_name("mat") + iter_arg = (name_in_kernel, pyop3.dtypes.OpaqueType("mat")) + + self.global_buffers[buffer_key] = buffer + self.global_buffer_intents[buffer_key] = intent + self._arguments.append(iter_arg) + + self._kernel_names[buffer_key] = name_in_kernel + return name_in_kernel + + def add_subkernel(self, subkernel): + pass + + # NOTE: Temporary while we work with basic kernels + # Without petsc mats or standalone functions, this is just empty idict + def set_temporary_shapes(self, shapes): + self._temporary_shapes = shapes + + @staticmethod + def make_kernel(context: CodegenContext): + if not isinstance(context, MLIRCodegenContext): + return ValueError("Requires MLIRCodegenContext object") + + '''Making the kernel + Available arguments: + - context.domains - iteration domains + - context.instructions - expression operations + - context.arguments - in/out variables + - name - function name + + Each instruction holds an expression + i.e. + insn_0: p_0 <- dat_0[0] + dat_1[0] + insn_1: idat_0[0 + 1*i_0] <- 1 {dep=insn_0} + insn_2: CODE(idat_0, p_0, dat_1, dat_0) {dep=insn_1} + + where the dep informs us of any dependency. + Each InstructionContext may have local variables, such as insn_1->i_0, which can be used to infer that insn_1 will iterate with the temporary variable i_0. + ''' + + + return ModuleOp([]) diff --git a/sample_2.py b/sample_2.py new file mode 100644 index 0000000000..241d30a229 --- /dev/null +++ b/sample_2.py @@ -0,0 +1,65 @@ +from xdsl.builder import ImplicitBuilder +from xdsl.dialects import arith, func, scf, tensor +from xdsl.dialects.builtin import ( + FloatAttr, + IndexType, + ModuleOp, + TensorType, + f32, + i32, +) +from xdsl.ir import Block, Region + +index = IndexType() + +N = 128 +t_f = TensorType(f32, [N]) # dat_2, dat_3 (float data) +t_i = TensorType(i32, [N]) # idat_0, idat_2 (index data) + +fn_block = Block(arg_types=[t_f, t_f, t_i, t_i]) + +with ImplicitBuilder(fn_block) as (dat_2, dat_3, idat_0, idat_2): + c0 = arith.ConstantOp.from_int_and_width(0, index) + c1 = arith.ConstantOp.from_int_and_width(1, index) + n = arith.ConstantOp.from_int_and_width(N, index) + two = arith.ConstantOp(FloatAttr(2.0, f32)) + + body = Block(arg_types=[index, t_f]) + with ImplicitBuilder(body) as (i2, acc): + # k = idat_2[i_2] + k = tensor.ExtractOp(idat_2, [i2], i32) + k_idx = arith.IndexCastOp(k.result, index) + + # j = idat_0[k] + j = tensor.ExtractOp(idat_0, [k_idx.result], i32) + j_idx = arith.IndexCastOp(j.result, index) + + # v = dat_3[j] + v = tensor.ExtractOp(dat_3, [j_idx.result], f32) + + # r = 2.0 * v + r = arith.MulfOp(two.result, v.result) + + # dat_2[j] = r (value semantics -> produces new tensor) + new = tensor.InsertOp(r.result, acc, [j_idx.result]) + + scf.YieldOp(new.result) + + loop = scf.ForOp( + lb=c0.result, + ub=n.result, + step=c1.result, + iter_args=[dat_2], + body=Region(body), + ) + + func.ReturnOp(loop.results[0]) + +fn = func.FuncOp( + "indirect_scale", + ((t_f, t_f, t_i, t_i), (t_f,)), + Region(fn_block), +) + +module = ModuleOp([fn]) +print(module) From a081fa72a58338cd68bec0fa0ceb2f21395b6bcb Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Mon, 20 Jul 2026 13:13:02 +0100 Subject: [PATCH 05/30] incorporating mlir -> loopy flag and building pymbolic to mlir pipeline --- assign_local_size.mlir | 35 +++++ assign_local_size.txt | 29 ++++ ...ading_demo.py => assign_offloading_demo.py | 5 + assign_loop.txt => assign_size.txt | 0 indirect_offloading_demo.py | 15 ++ integral_loop.txt | 68 +++++++++ pyop3/debug_flags.py | 1 + pyop3/insn/exec.py | 7 +- pyop3/lower/context.py | 1 + pyop3/lower/loopy.py | 29 ++-- pyop3/lower/mlir.py | 136 ++++++++++++++---- 11 files changed, 286 insertions(+), 40 deletions(-) create mode 100644 assign_local_size.mlir create mode 100644 assign_local_size.txt rename gpu_offloading_demo.py => assign_offloading_demo.py (84%) rename assign_loop.txt => assign_size.txt (100%) create mode 100644 indirect_offloading_demo.py create mode 100644 integral_loop.txt create mode 100644 pyop3/debug_flags.py diff --git a/assign_local_size.mlir b/assign_local_size.mlir new file mode 100644 index 0000000000..0af0ddd9cd --- /dev/null +++ b/assign_local_size.mlir @@ -0,0 +1,35 @@ +builtin.module { + func.func @pyop3_loop( + %dat_0: tensor, + %dat_1: tensor, + %idat_0: tensor, + %idat_1: tensor, + %idat_2: tensor, + %idat_3: tensor + ) -> tensor { + + %c0 = arith.constant 0 : index // iter var + %c1 = arith.constant 1 : index // iter var + %c17 = arith.constant 17 : index // + %c15 = arith.constant 15 : index + %c32 = arith.constant 32 : index + %c2f = arith.constant 2.0 : f64 + + scf.for %i_0 = %c0 to %c17 step %c1 iter_args() -> () {} + + %f2 = scf.for %i_2 = %c0 to %c15 step %c1 iter_args(%dat_it = %dat_0) -> (tensor) { + %e1 = tensor.extract %idat_2[%i_2] : tensor + %ii_2 = arith.index_cast %e1 : i32 to index + %e2 = tensor.extract %idat_0[%ii_2] : tensor + %iii_2 = arith.index_cast %e2 : i32 to index + %v1 = tensor.extract %dat_1[%iii_2] : tensor + %v2 = arith.mulf %c2f, %v1 : f64 + %res = tensor.insert %v2 into %dat_it[%iii_2] : tensor + scf.yield %res : tensor + } + + scf.for %i_3 = %c0 to %c32 step %c1 iter_args() -> () {} + + func.return %f2 : tensor + } +} diff --git a/assign_local_size.txt b/assign_local_size.txt new file mode 100644 index 0000000000..d5e7205a71 --- /dev/null +++ b/assign_local_size.txt @@ -0,0 +1,29 @@ +******************************************************************************** +#include +#include +#include +#include + +void pyop3_loop(double *__restrict__ dat_0, double const *__restrict__ dat_1, int32_t const *__restrict__ idat_0, int32_t const *__restrict__ idat_1, int32_t const *__restrict__ idat_2, int32_t const *__restrict__ idat_3) +{ + for (int32_t i_0 = 0; i_0 <= 17; ++i_0) + { + } + for (int32_t i_2 = 0; i_2 <= 15; ++i_2) + dat_0[idat_0[idat_2[i_2]]] = 2.0 * dat_1[idat_0[idat_2[i_2]]]; + for (int32_t i_3 = 0; i_3 <= 32; ++i_3) + { + } + +} +******************************************************************************** +dat_0 (16) : [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] +dat_1 (16) : [10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10.] +idat_0 (67) : [ 0 0 0 0 0 1 2 3 3 3 3 4 4 4 4 5 5 5 5 6 6 6 6 7 + 7 7 7 8 8 8 8 9 9 9 9 9 9 10 10 10 10 11 11 11 11 12 12 12 + 12 13 13 13 13 13 13 13 13 14 14 14 14 15 15 15 15 15 15] +idat_1 (18) : [ 0 7 11 15 19 23 27 31 33 37 41 45 49 51 53 57 61 63] +idat_2 (16) : [ 4 5 6 10 14 18 22 26 30 36 40 44 48 56 60 66] +idat_3 (33) : [ 1 2 3 8 9 12 13 16 17 20 21 24 25 28 29 32 34 35 38 39 42 43 46 47 + 50 52 54 55 58 59 62 64 65] +******************************************************************************** diff --git a/gpu_offloading_demo.py b/assign_offloading_demo.py similarity index 84% rename from gpu_offloading_demo.py rename to assign_offloading_demo.py index f7133521cf..af720699b7 100644 --- a/gpu_offloading_demo.py +++ b/assign_offloading_demo.py @@ -2,14 +2,19 @@ import pyop3 as op3 import numpy as np, cupy as cp +import pyop3.debug_flags + mesh = UnitSquareMesh(3,3) + V = FunctionSpace(mesh, "CG", 1) f = Function(V).assign(10) g = Function(V) gpu = op3.CUDAGPU() +pyop3.debug_flags.hit_assign = True g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile") +pyop3.debug_flags.hit_assign = False # with op3.offloading(gpu): # g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile") # assert isinstance(g.dat.data_ro, cp.ndarray) # Device diff --git a/assign_loop.txt b/assign_size.txt similarity index 100% rename from assign_loop.txt rename to assign_size.txt diff --git a/indirect_offloading_demo.py b/indirect_offloading_demo.py new file mode 100644 index 0000000000..8b4e8e6148 --- /dev/null +++ b/indirect_offloading_demo.py @@ -0,0 +1,15 @@ +from firedrake import * +import pyop3 as op3 +import numpy as np, cupy as cp + +import pyop3.debug_flags + +mesh = UnitSquareMesh(3,3) + +V = FunctionSpace(mesh, "CG", 1) +v = TestFunction(V) + +pyop3.debug_flags.hit_assign = True +b = assemble(conj(v) * dx) +pyop3.debug_flags.hit_assign = False + diff --git a/integral_loop.txt b/integral_loop.txt new file mode 100644 index 0000000000..feac230fa3 --- /dev/null +++ b/integral_loop.txt @@ -0,0 +1,68 @@ +******************************************************************************** +#include +#include +#include +#include +#include + +static void form_cell_integral(double *__restrict__ A, double const *__restrict__ coords_0); +static void form_cell_integral(double *__restrict__ A, double const *__restrict__ coords_0) +{ + double t0; + double t1; + double t2; + double t3[3] = { 0.33333333333333337, 0.33333333333333326, 0.33333333333333326 }; + + t0 = -1.0 * coords_0[0]; + t1 = -1.0 * coords_0[1]; + t2 = 0.5 * fabs((t0 + coords_0[2]) * (t1 + coords_0[5]) + -1.0 * (t0 + coords_0[4]) * (t1 + coords_0[3])); + for (int32_t j = 0; j <= 2; ++j) + A[j] = A[j] + t3[j] * t2; + +} + +void pyop3_loop(double const *__restrict__ dat_0, double *__restrict__ dat_1, int32_t const *__restrict__ idat_0, int32_t const *__restrict__ idat_1) +{ + int32_t j_0; + int32_t j_1; + int32_t j_2; + double t_0[3l]; + double t_1[6l]; + + for (int32_t i_0 = 0; i_0 <= 17; ++i_0) + { + for (int32_t i_1 = 0; i_1 <= 2; ++i_1) + { + j_0 = 0 + 1 * (i_1 + 0); + t_0[j_0] = (double) (0.0); + } + for (int32_t i_2 = 0; i_2 <= 2; ++i_2) + for (int32_t i_3 = 0; i_3 <= 1; ++i_3) + { + j_1 = 0 + 1 * (i_2 * 2 + 0 + i_3); + t_1[j_1] = dat_0[idat_0[3 * i_0 + i_2] + i_3]; + } + form_cell_integral(&(t_0[0]), &(t_1[0])); + for (int32_t i_6 = 0; i_6 <= 2; ++i_6) + { + j_2 = 0 + 1 * (i_6 + 0); + dat_1[idat_1[3 * i_0 + i_6]] = dat_1[idat_1[3 * i_0 + i_6]] + t_0[j_2]; + } + } + +} +******************************************************************************** +dat_0 (32) : [0.33333333 0. 0. 0.33333333 0. 0. + 0.33333333 0.33333333 0. 0.66666667 0.66666667 0. + 0.33333333 0.66666667 0.66666667 0.33333333 0. 1. + 1. 0. 0.33333333 1. 0.66666667 0.66666667 + 1. 0.33333333 0.66666667 1. 1. 0.66666667 + 1. 1. ] +dat_1 (16) : [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] +idat_0 (54) : [ 0 2 4 0 2 6 2 6 8 0 6 10 6 8 12 6 10 14 8 12 16 6 12 14 + 10 14 18 12 16 20 12 14 22 14 18 24 12 20 22 14 22 24 20 22 26 22 24 28 + 22 26 28 26 28 30] +idat_1 (54) : [ 0 1 2 0 1 3 1 3 4 0 3 5 3 4 6 3 5 7 4 6 8 3 6 7 + 5 7 9 6 8 10 6 7 11 7 9 12 6 10 11 7 11 12 10 11 13 11 12 14 + 11 13 14 13 14 15] +******************************************************************************** diff --git a/pyop3/debug_flags.py b/pyop3/debug_flags.py new file mode 100644 index 0000000000..63c21bc953 --- /dev/null +++ b/pyop3/debug_flags.py @@ -0,0 +1 @@ +hit_assign = False diff --git a/pyop3/insn/exec.py b/pyop3/insn/exec.py index 3ddce126c9..bc0de9fdcd 100644 --- a/pyop3/insn/exec.py +++ b/pyop3/insn/exec.py @@ -27,6 +27,8 @@ from pyop3.cache import cached_method, memory_cache from pyop3.insn.base import READ, WRITE, RW, INC, MIN_RW, MIN_WRITE, MAX_RW, MAX_WRITE +import pyop3.debug_flags + @dataclasses.dataclass(frozen=True, kw_only=True) class CompilerParameters: @@ -55,6 +57,7 @@ class CompilerParameters: # TODO: handle these - need to build CompilerOptions + codegen: str = "loopy" # extra_cflags: tuple[str, ...] = () # extra_ldflags: tuple[str, ...] = () @@ -104,7 +107,7 @@ def parse_compiler_parameters(compiler_parameters: CompilerParametersT) -> Parse return compiler_parameters if compiler_parameters is None: - compiler_parameters = {} + compiler_parameters = {"codegen": "mlir"} else: # TODO: nice error message assert pyop3.collections.is_ordered_mapping(compiler_parameters) @@ -262,6 +265,7 @@ def _compile(self) -> CompiledCodeExecutor: assert num_buffers == len(self.preprocessed_buffers) compiler_parameters = parse_compiler_parameters(self.compiler_parameters) + loopy_code, buffer_index_map = _compile_static(self, compiler_parameters) extra_compiler_options = collect_compiler_options(self._preprocessed) @@ -278,6 +282,7 @@ def _compile(self) -> CompiledCodeExecutor: petsc_events=petsc_events, ) + # TODO: We don't do anything with nest indices yet because we have always already # unpacked things sorted_buffers = {} diff --git a/pyop3/lower/context.py b/pyop3/lower/context.py index 664bd66eff..42ed833b7c 100644 --- a/pyop3/lower/context.py +++ b/pyop3/lower/context.py @@ -13,6 +13,7 @@ def __init__(self, *, check_negatives): self._instructions = [] self._arguments = [] self._subkernels = [] + self._last_insn_id = None # determine dependence self._name_generator = utils.UniqueNameGenerator() diff --git a/pyop3/lower/loopy.py b/pyop3/lower/loopy.py index 178f4f00da..d7d81f8864 100644 --- a/pyop3/lower/loopy.py +++ b/pyop3/lower/loopy.py @@ -63,6 +63,7 @@ # TODO: import other way around? from pyop3.insn.exec import parse_compiler_parameters +import pyop3.debug_flags # Sam for debugging # FIXME this needs to be synchronised with TSFC, tricky # shared base package? or both set by Firedrake - better solution @@ -74,7 +75,6 @@ def __init__(self, *, check_negatives): super().__init__(check_negatives=check_negatives) self._within_inames = frozenset() - self._last_insn_id = None # initializer hash -> temporary name self._reusable_temporaries: dict[int, str] = {} @@ -141,6 +141,7 @@ def add_buffer(self, buffer, intent: Intent | None = None) -> str: # dat2[i] = dat1[2*i] # # is not. + if buffer.is_nested: raise NotImplementedError("Currently handle nesting outside the generated code") @@ -207,6 +208,8 @@ def add_buffer(self, buffer, intent: Intent | None = None) -> str: def add_temporary(self, prefix="t", dtype=IntType, *, shape=(), initializer: np.ndarray = None, read_only: bool = False) -> str: # If multiple temporaries with the same initializer are used then they # can be shared. + global mlir_context + can_reuse = initializer is not None and read_only if can_reuse: key = initializer.data.tobytes() @@ -397,8 +400,10 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed else: cs_expr = (insn,) - context = LoopyCodegenContext(check_negatives=compiler_parameters.check_negatives) - mlir_context = MLIRCodegenContext(check_negatives=compiler_parameters.check_negatives) + if compiler_parameters.codegen == "mlir": + context = MLIRCodegenContext(check_negatives=compiler_parameters.check_negatives) + else: + context = LoopyCodegenContext(check_negatives=compiler_parameters.check_negatives) # 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. @@ -420,6 +425,9 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed "The generated kernel does not modify any global data, this may indicate that something has gone wrong" ) + if pyop3.debug_flags.hit_assign: + breakpoint() + # add a no-op instruction touching all of the kernel arguments so they are # not silently dropped noop = lp.CInstruction( @@ -437,8 +445,6 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed ("30_petsc", "#include "), # perhaps only if petsc callable used? ] - breakpoint() - translation_unit = lp.make_kernel( context.domains, context.instructions, @@ -578,10 +584,11 @@ def parse_loop_properly_this_time( if axis_tree.linearize(path_, partial=True).size == 0: continue - elif component.size != 1: + elif component.local_size != 1: + iname = codegen_context.unique_name("i") domain_var = register_extent( - component.size, + component.local_size, iname_map, loop_indices, codegen_context, @@ -785,6 +792,7 @@ def compile_array_assignment( axis_tree=None, paths=None, ): + if paths is None: paths = [] if iname_replace_maps is None: @@ -818,10 +826,10 @@ def compile_array_assignment( # 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: + elif component.local_size != 1: iname = codegen_context.unique_name("i") extent_var = register_extent( - component.size, + component.local_size, iname_replace_maps[-1], loop_indices, codegen_context, @@ -865,7 +873,6 @@ def compile_array_assignment( loop_indices, ) - def add_leaf_assignment( assignment, paths, @@ -935,7 +942,7 @@ def _(mul: pyop3.expr.Mul, /, *args, **kwargs) -> pym.Expression: @_lower_expr.register(pyop3.expr.Modulo) -def _(mod: pyop3.expr.Mul, /, *args, **kwargs) -> pym.Expression: +def _(mod: pyop3.expr.Modulo, /, *args, **kwargs) -> pym.Expression: return _lower_expr(mod.a, *args, **kwargs) % _lower_expr(mod.b, *args, **kwargs) diff --git a/pyop3/lower/mlir.py b/pyop3/lower/mlir.py index 0d3a033f93..9f2473f023 100644 --- a/pyop3/lower/mlir.py +++ b/pyop3/lower/mlir.py @@ -1,10 +1,24 @@ -from xdsl.dialects import arith, func, tensor, linalg +import contextlib +import numpy as np + +from xdsl.dialects import arith, func, tensor, scf from xdsl.dialects.builtin import ( ModuleOp, - IntegerAttr + IntegerAttr, + FunctionType, + TensorType, + i32, + i64, + f64 +) + +from xdsl.dialects.func import ( + FuncOp ) +from xdsl.ir import SSAValue + import pyop3 from pyop3.buffer import AbstractBuffer, ConcreteBuffer, PetscMatBuffer, ArrayBuffer, NullBuffer from pyop3.dtypes import IntType @@ -15,11 +29,54 @@ Intent ) +class Argument: + def __init__(self, name, dtype, shape): + self.name = name + self.dtype = dtype + self.shape = shape + + def __str__(self): + return self.name + +class SymbolTable: + ''' + Symbol Table that acts as lookup for pymbol to MLIR SSAValue + + Context manager works around MLIR's region and block based system + ''' + + def __init__(self): + self._scopes: list[dict[str, SSAValue]] = [{}] + + @contextlib.contextmanager + def scope(self): + self._scopes.append({}) + try: + yield self + finally: + self._scopes.pop() + + def insert(self, name: str, value: SSAValue): + self._scopes[-1][name] = value + + def lookup(self, name: str) -> SSAValue: + for sc in reversed(self._scopes): + if name in sc: + return sc[name] + raise KeyError(f"Unknown variable: {name}") + + def __str__(self): + return str(self._scopes) + class MLIRCodegenContext(CodegenContext): def __init__(self, *, check_negatives): super().__init__(check_negatives=check_negatives) + self.symbol_table = SymbolTable() + + self._within_inames = frozenset() + def add_domain(self, iname, *args): nargs = len(args) if nargs == 1: @@ -30,6 +87,7 @@ def add_domain(self, iname, *args): self._domains.append((start, stop)) def add_assignment(self, assigneee, expression, prefix="insn"): + # Assignee and expression come in pymbolic expression pass def add_function_call(self, assignees, expression, prefix="insn"): @@ -38,6 +96,7 @@ def add_function_call(self, assignees, expression, prefix="insn"): def add_buffer(self, buffer, intent: Intent | None = None) -> str: # TODO: This only works for np.ndarrays for development atm + buffer_key = (buffer.name, buffer.nest_indices) if isinstance(buffer, NullBuffer): assert not buffer.nest_indices @@ -68,13 +127,13 @@ def add_buffer(self, buffer, intent: Intent | None = None) -> str: shape = self._temporary_shapes.get(buffer_key, None) # TODO: An equivalent of lp.GlobalArg is required here # GlobalArg represents array, dtype, shape, address space (local or global variable) - iter_arg = (name_in_kernel, buffer.dtype, shape) + iter_arg = Argument(name_in_kernel, buffer.dtype, shape) else: assert isinstance(buffer, PetscMatBuffer) assert buffer.mat_type not in {"nest", "python"} name_in_kernel = self.unique_name("mat") - iter_arg = (name_in_kernel, pyop3.dtypes.OpaqueType("mat")) + iter_arg = Argument(name_in_kernel, pyop3.dtypes.OpaqueType("mat")) self.global_buffers[buffer_key] = buffer self.global_buffer_intents[buffer_key] = intent @@ -86,32 +145,53 @@ def add_buffer(self, buffer, intent: Intent | None = None) -> str: def add_subkernel(self, subkernel): pass + # NOTE: Here while I port code, to be removed + @contextlib.contextmanager + def within_inames(self, inames) -> None: + orig_within_inames = self._within_inames + self._within_inames |= inames + yield + self._within_inames = orig_within_inames # NOTE: Temporary while we work with basic kernels # Without petsc mats or standalone functions, this is just empty idict def set_temporary_shapes(self, shapes): self._temporary_shapes = shapes - @staticmethod - def make_kernel(context: CodegenContext): - if not isinstance(context, MLIRCodegenContext): - return ValueError("Requires MLIRCodegenContext object") - - '''Making the kernel - Available arguments: - - context.domains - iteration domains - - context.instructions - expression operations - - context.arguments - in/out variables - - name - function name - - Each instruction holds an expression - i.e. - insn_0: p_0 <- dat_0[0] + dat_1[0] - insn_1: idat_0[0 + 1*i_0] <- 1 {dep=insn_0} - insn_2: CODE(idat_0, p_0, dat_1, dat_0) {dep=insn_1} - - where the dep informs us of any dependency. - Each InstructionContext may have local variables, such as insn_1->i_0, which can be used to infer that insn_1 will iterate with the temporary variable i_0. - ''' - - - return ModuleOp([]) +def make_kernel(context: MLIRCodegenContext): + if not isinstance(context, MLIRCodegenContext): + return ValueError("Requires MLIRCodegenContext object") + + '''Making the kernel + Available arguments: + - context.domains - iteration domains + - context.instructions - expression operations + - context.arguments - in/out variables + - name - function name + + Each instruction holds an expression + i.e. + insn_0: p_0 <- dat_0[0] + dat_1[0] + insn_1: idat_0[0 + 1*i_0] <- 1 {dep=insn_0} + insn_2: CODE(idat_0, p_0, dat_1, dat_0) {dep=insn_1} + + where the dep informs us of any dependency. + Each InstructionContext may have local variables, such as insn_1->i_0, which can be used to infer that insn_1 will iterate with the temporary variable i_0. + + The domain exists come in (start, stop) pairs, so we could develop our scf.for loops appropriately. + + + ''' + # Update when moving away from loopy-defined argument variables + arg_types = [ + TensorType(arg.dtype, [-1]) # -1 for dynamic, but vector length is known + for arg in self._arguments + ] + + func_op = FuncOp("pyop3_loop", FunctionType.from_lists(arg_types, [])) + block_args = self.parent_func.body.blocks[0].args + + for arg, ssa_value in zip(self._arguments, block_args): + self.symbol_table.insert(arg, ssa_value) + + + return ModuleOp([]) From 8b4f6d80ede6dfa988bef98b30472630f22da3f2 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Tue, 21 Jul 2026 16:36:02 +0100 Subject: [PATCH 06/30] MLIR generating stage - but MLIR typing is incorrect, need type inference --- assign_offloading_demo.py | 2 +- loopy_assign.txt | 38 ++++++ mlir_assign.mlir | 172 ++++++++++++++++++++++++ pyop3/insn/exec.py | 2 +- pyop3/lower/context.py | 2 +- pyop3/lower/loopy.py | 58 ++++---- pyop3/lower/mlir.py | 271 ++++++++++++++++++++++++++++++++------ 7 files changed, 480 insertions(+), 65 deletions(-) create mode 100644 loopy_assign.txt create mode 100644 mlir_assign.mlir diff --git a/assign_offloading_demo.py b/assign_offloading_demo.py index af720699b7..8c1555a45d 100644 --- a/assign_offloading_demo.py +++ b/assign_offloading_demo.py @@ -13,7 +13,7 @@ gpu = op3.CUDAGPU() pyop3.debug_flags.hit_assign = True -g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile") +g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile", compiler_parameters={"codegen": "mlir"}) pyop3.debug_flags.hit_assign = False # with op3.offloading(gpu): # g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile") diff --git a/loopy_assign.txt b/loopy_assign.txt new file mode 100644 index 0000000000..d14de3083d --- /dev/null +++ b/loopy_assign.txt @@ -0,0 +1,38 @@ +--------------------------------------------------------------------------- +KERNEL: pyop3_loop +--------------------------------------------------------------------------- +ARGUMENTS: +dat_0: ArrayArg, type: np:dtype('float64'), shape: unknown in/out aspace: global +dat_1: ArrayArg, type: np:dtype('float64'), shape: unknown in aspace: global +idat_0: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global +idat_1: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global +idat_2: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global +idat_3: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global +--------------------------------------------------------------------------- +DOMAINS: +{ [i_0] : 0 <= i_0 <= 17 } +{ [i_1] : 1 = 0 } +{ [i_2] : 0 <= i_2 <= 15 } +{ [i_3] : 0 <= i_3 <= 32 } +{ [i_4] : 1 = 0 } +--------------------------------------------------------------------------- +INAME TAGS: +i_0: None +i_1: None +i_2: None +i_3: None +i_4: None +--------------------------------------------------------------------------- +INSTRUCTIONS: + for i_1, i_0 +↱ dat_0[0 + 1*(idat_0[0 + 1*idat_1[0 + 1*i_0]] + i_1)] = 2*dat_1[0 + 1*(idat_0[0 + 1*idat_1[0 + 1*i_0]] + i_1)] {id=insn_0} +│ end i_1, i_0 +│ for i_2 +└↱ dat_0[0 + 1*(idat_0[0 + 1*idat_2[0 + 1*i_2]] + 0)] = 2*dat_1[0 + 1*(idat_0[0 + 1*idat_2[0 + 1*i_2]] + 0) {id=insn_1} + │ end i_2 + │ for i_4, i_3 +↱└ dat_0[0 + 1*(idat_0[0 + 1*idat_3[0 + 1*i_3]] + i_4) = 2*dat_1[0 + 1*(idat_0[0 + 1*idat_3[0 + 1*i_3]] + i_4) {id=insn_2} +│ end i_4, i_3 +└ CODE(idat_2, dat_0, idat_0, idat_1, dat_1, idat_3|) {id=insn} + +--------------------------------------------------------------------------- diff --git a/mlir_assign.mlir b/mlir_assign.mlir new file mode 100644 index 0000000000..4eab0e02b9 --- /dev/null +++ b/mlir_assign.mlir @@ -0,0 +1,172 @@ +builtin.module { + func.func @pyop3_loop(%0: tensor, %1: tensor, %2: tensor, %3: tensor, %4: tensor, %5: tensor) { + %6 = arith.constant 0 : i32 + %7 = arith.index_cast %6 : i32 to index + %8 = arith.constant 18 : i32 + %9 = arith.index_cast %8 : i32 to index + %10 = arith.constant 1 : i32 + %11 = arith.index_cast %10 : i32 to index + %12 = scf.for %13 = %7 to %9 step %11 iter_args(%14 = %0) -> (tensor) { + %15 = arith.constant 0 : i32 + %16 = arith.index_cast %15 : i32 to index + %17 = arith.constant 0 : i32 + %18 = arith.index_cast %17 : i32 to index + %19 = arith.constant 1 : i32 + %20 = arith.index_cast %19 : i32 to index + %21 = scf.for %22 = %16 to %18 step %20 iter_args(%23 = %14) -> (tensor) { + %24 = arith.constant 2 : i32 + %25 = arith.constant 0 : i32 + %26 = arith.constant 1 : i32 + %27 = arith.constant 0 : i32 + %28 = arith.constant 1 : i32 + %29 = arith.constant 0 : i32 + %30 = arith.constant 1 : i32 + %31 = arith.muli %30, %13 : i32 + %32 = arith.addi %29, %31 : i32 + %33 = arith.index_cast %32 : i32 to index + %34 = tensor.extract %2[%33] : tensor + %35 = arith.muli %28, %34 : i32 + %36 = arith.addi %27, %35 : i32 + %37 = arith.index_cast %36 : i32 to index + %38 = tensor.extract %1[%37] : tensor + %39 = arith.addi %38, %22 : i32 + %40 = arith.muli %26, %39 : i32 + %41 = arith.addi %25, %40 : i32 + %42 = arith.index_cast %41 : i32 to index + %43 = tensor.extract %3[%42] : tensor + %44 = arith.muli %24, %43 : i32 + %45 = arith.constant 0 : i32 + %46 = arith.constant 1 : i32 + %47 = arith.constant 0 : i32 + %48 = arith.constant 1 : i32 + %49 = arith.constant 0 : i32 + %50 = arith.constant 1 : i32 + %51 = arith.muli %50, %13 : i32 + %52 = arith.addi %49, %51 : i32 + %53 = arith.index_cast %52 : i32 to index + %54 = tensor.extract %2[%53] : tensor + %55 = arith.muli %48, %54 : i32 + %56 = arith.addi %47, %55 : i32 + %57 = arith.index_cast %56 : i32 to index + %58 = tensor.extract %1[%57] : tensor + %59 = arith.addi %58, %22 : i32 + %60 = arith.muli %46, %59 : i32 + %61 = arith.addi %45, %60 : i32 + %62 = arith.index_cast %61 : i32 to index + %63 = tensor.insert %44 into %23[%62] : tensor + scf.yield %63 : tensor + } + scf.yield %21 : tensor + } + %64 = arith.constant 0 : i32 + %65 = arith.index_cast %64 : i32 to index + %66 = arith.constant 16 : i32 + %67 = arith.index_cast %66 : i32 to index + %68 = arith.constant 1 : i32 + %69 = arith.index_cast %68 : i32 to index + %70 = scf.for %71 = %65 to %67 step %69 iter_args(%72 = %12) -> (tensor) { + %73 = arith.constant 2 : i32 + %74 = arith.constant 0 : i32 + %75 = arith.constant 1 : i32 + %76 = arith.constant 0 : i32 + %77 = arith.constant 1 : i32 + %78 = arith.constant 0 : i32 + %79 = arith.constant 1 : i32 + %80 = arith.muli %79, %71 : i32 + %81 = arith.addi %78, %80 : i32 + %82 = arith.index_cast %81 : i32 to index + %83 = tensor.extract %4[%82] : tensor + %84 = arith.muli %77, %83 : i32 + %85 = arith.addi %76, %84 : i32 + %86 = arith.index_cast %85 : i32 to index + %87 = tensor.extract %1[%86] : tensor + %88 = arith.constant 0 : i32 + %89 = arith.addi %87, %88 : i32 + %90 = arith.muli %75, %89 : i32 + %91 = arith.addi %74, %90 : i32 + %92 = arith.index_cast %91 : i32 to index + %93 = tensor.extract %3[%92] : tensor + %94 = arith.muli %73, %93 : i32 + %95 = arith.constant 0 : i32 + %96 = arith.constant 1 : i32 + %97 = arith.constant 0 : i32 + %98 = arith.constant 1 : i32 + %99 = arith.constant 0 : i32 + %100 = arith.constant 1 : i32 + %101 = arith.muli %100, %71 : i32 + %102 = arith.addi %99, %101 : i32 + %103 = arith.index_cast %102 : i32 to index + %104 = tensor.extract %4[%103] : tensor + %105 = arith.muli %98, %104 : i32 + %106 = arith.addi %97, %105 : i32 + %107 = arith.index_cast %106 : i32 to index + %108 = tensor.extract %1[%107] : tensor + %109 = arith.constant 0 : i32 + %110 = arith.addi %108, %109 : i32 + %111 = arith.muli %96, %110 : i32 + %112 = arith.addi %95, %111 : i32 + %113 = arith.index_cast %112 : i32 to index + %114 = tensor.insert %94 into %72[%113] : tensor + scf.yield %114 : tensor + } + %115 = arith.constant 0 : i32 + %116 = arith.index_cast %115 : i32 to index + %117 = arith.constant 33 : i32 + %118 = arith.index_cast %117 : i32 to index + %119 = arith.constant 1 : i32 + %120 = arith.index_cast %119 : i32 to index + %121 = scf.for %122 = %116 to %118 step %120 iter_args(%123 = %70) -> (tensor) { + %124 = arith.constant 0 : i32 + %125 = arith.index_cast %124 : i32 to index + %126 = arith.constant 0 : i32 + %127 = arith.index_cast %126 : i32 to index + %128 = arith.constant 1 : i32 + %129 = arith.index_cast %128 : i32 to index + %130 = scf.for %131 = %125 to %127 step %129 iter_args(%132 = %123) -> (tensor) { + %133 = arith.constant 2 : i32 + %134 = arith.constant 0 : i32 + %135 = arith.constant 1 : i32 + %136 = arith.constant 0 : i32 + %137 = arith.constant 1 : i32 + %138 = arith.constant 0 : i32 + %139 = arith.constant 1 : i32 + %140 = arith.muli %139, %122 : i32 + %141 = arith.addi %138, %140 : i32 + %142 = arith.index_cast %141 : i32 to index + %143 = tensor.extract %5[%142] : tensor + %144 = arith.muli %137, %143 : i32 + %145 = arith.addi %136, %144 : i32 + %146 = arith.index_cast %145 : i32 to index + %147 = tensor.extract %1[%146] : tensor + %148 = arith.addi %147, %131 : i32 + %149 = arith.muli %135, %148 : i32 + %150 = arith.addi %134, %149 : i32 + %151 = arith.index_cast %150 : i32 to index + %152 = tensor.extract %3[%151] : tensor + %153 = arith.muli %133, %152 : i32 + %154 = arith.constant 0 : i32 + %155 = arith.constant 1 : i32 + %156 = arith.constant 0 : i32 + %157 = arith.constant 1 : i32 + %158 = arith.constant 0 : i32 + %159 = arith.constant 1 : i32 + %160 = arith.muli %159, %122 : i32 + %161 = arith.addi %158, %160 : i32 + %162 = arith.index_cast %161 : i32 to index + %163 = tensor.extract %5[%162] : tensor + %164 = arith.muli %157, %163 : i32 + %165 = arith.addi %156, %164 : i32 + %166 = arith.index_cast %165 : i32 to index + %167 = tensor.extract %1[%166] : tensor + %168 = arith.addi %167, %131 : i32 + %169 = arith.muli %155, %168 : i32 + %170 = arith.addi %154, %169 : i32 + %171 = arith.index_cast %170 : i32 to index + %172 = tensor.insert %153 into %132[%171] : tensor + scf.yield %172 : tensor + } + scf.yield %130 : tensor + } + func.return + } +} diff --git a/pyop3/insn/exec.py b/pyop3/insn/exec.py index bc0de9fdcd..3064b278c3 100644 --- a/pyop3/insn/exec.py +++ b/pyop3/insn/exec.py @@ -107,7 +107,7 @@ def parse_compiler_parameters(compiler_parameters: CompilerParametersT) -> Parse return compiler_parameters if compiler_parameters is None: - compiler_parameters = {"codegen": "mlir"} + compiler_parameters = {} else: # TODO: nice error message assert pyop3.collections.is_ordered_mapping(compiler_parameters) diff --git a/pyop3/lower/context.py b/pyop3/lower/context.py index 42ed833b7c..a4f4fc5228 100644 --- a/pyop3/lower/context.py +++ b/pyop3/lower/context.py @@ -9,7 +9,7 @@ class CodegenContext(ABC): def __init__(self, *, check_negatives): self.check_negatives = check_negatives - self._domains = [] + self._domains = [] self._instructions = [] self._arguments = [] self._subkernels = [] diff --git a/pyop3/lower/loopy.py b/pyop3/lower/loopy.py index d7d81f8864..3af8da952d 100644 --- a/pyop3/lower/loopy.py +++ b/pyop3/lower/loopy.py @@ -98,6 +98,7 @@ def add_assignment(self, assignee, expression, prefix="insn"): depends_on=self._depends_on, depends_on_is_final=True, ) + self._add_instruction(insn) def add_cinstruction(self, insn_str, read_variables=frozenset()): @@ -208,8 +209,6 @@ def add_buffer(self, buffer, intent: Intent | None = None) -> str: def add_temporary(self, prefix="t", dtype=IntType, *, shape=(), initializer: np.ndarray = None, read_only: bool = False) -> str: # If multiple temporaries with the same initializer are used then they # can be shared. - global mlir_context - can_reuse = initializer is not None and read_only if can_reuse: key = initializer.data.tobytes() @@ -425,36 +424,43 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed "The generated kernel does not modify any global data, this may indicate that something has gone wrong" ) - if pyop3.debug_flags.hit_assign: - breakpoint() - # 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) + # NOTE: Obviously to improve this ugly if-operation + if compiler_parameters.codegen == "loopy": + 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)) + if compiler_parameters.codegen == "mlir": + mlir_unit = context.make_kernel() + else: + 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)) + + if compiler_parameters.codegen == "mlir": + breakpoint() + raise NotImplementedError("Still at generation stage") + entrypoint = translation_unit.default_entrypoint if compiler_parameters.add_likwid_markers: @@ -473,11 +479,13 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed buffer_index = op.preprocessed_buffers.index(buffer_ref) intent = context.global_buffer_intents[buffer_key] buffer_index_map[kernel_arg.name] = (buffer_index, buffer_ref.nest_indices, intent) + + if pyop3.debug_flags.hit_assign: + breakpoint() return translation_unit, buffer_index_map - # put into a class in transform.py? @functools.singledispatch def _collect_temporary_shapes(expr): diff --git a/pyop3/lower/mlir.py b/pyop3/lower/mlir.py index 9f2473f023..0ad06ba1d1 100644 --- a/pyop3/lower/mlir.py +++ b/pyop3/lower/mlir.py @@ -1,11 +1,20 @@ import contextlib +import functools +import numbers +import dataclasses import numpy as np +import pymbolic as pym +import loopy as lp # NOTE: For typing, temporary until fully separated + from xdsl.dialects import arith, func, tensor, scf from xdsl.dialects.builtin import ( + DYNAMIC_INDEX, ModuleOp, IntegerAttr, + IntegerType, + IndexType, FunctionType, TensorType, i32, @@ -17,7 +26,8 @@ FuncOp ) -from xdsl.ir import SSAValue +from xdsl.builder import Builder, InsertPoint +from xdsl.ir import SSAValue, Block, Region import pyop3 from pyop3.buffer import AbstractBuffer, ConcreteBuffer, PetscMatBuffer, ArrayBuffer, NullBuffer @@ -28,6 +38,22 @@ from pyop3.insn.base import ( Intent ) +NUMPY_TO_XDSL = { + np.dtype(np.float64): f64, + np.dtype(np.int32): i32, + np.dtype(np.int64): i64, +} + +@dataclasses.dataclass +class Assignment: + assignee: object + expression: object + within_inames: object + id: str + + def __str__(self): + return f"{self.assignee} = {self.expression}" + class Argument: def __init__(self, name, dtype, shape): @@ -38,6 +64,9 @@ def __init__(self, name, dtype, shape): def __str__(self): return self.name + def __repr__(self): + return f"<{self.name}, dtype: {self.dtype}, shape: {self.shape if self.shape else '?'}>" + class SymbolTable: ''' Symbol Table that acts as lookup for pymbol to MLIR SSAValue @@ -75,7 +104,9 @@ def __init__(self, *, check_negatives): self.symbol_table = SymbolTable() + # NOTE: Temporary & unused while I rewrite lower/ self._within_inames = frozenset() + self._domains = dict() def add_domain(self, iname, *args): nargs = len(args) @@ -84,11 +115,17 @@ def add_domain(self, iname, *args): else: assert nargs == 2 start, stop = args[0], args[1] - self._domains.append((start, stop)) + self._domains[iname] = (start, stop) - def add_assignment(self, assigneee, expression, prefix="insn"): + def add_assignment(self, assignee, expression, prefix="insn"): # Assignee and expression come in pymbolic expression - pass + insn = Assignment( + assignee=assignee, + expression=expression, + within_inames=self._within_inames, + id=self.unique_name(prefix) + ) + self._instructions.append(insn) def add_function_call(self, assignees, expression, prefix="insn"): pass @@ -145,6 +182,15 @@ def add_buffer(self, buffer, intent: Intent | None = None) -> str: def add_subkernel(self, subkernel): pass + def add_instruction(self, insn): + # TODO: Ignoring CInstruction for now because no MLIR equivalent built + if isinstance(insn, lp.CInstruction): + raise ValueError("Cannot deal with Loopy CInstructions") + + self._instructions.append(insn) + self._last_insn_id = insn.id + + # NOTE: Here while I port code, to be removed @contextlib.contextmanager def within_inames(self, inames) -> None: @@ -157,41 +203,192 @@ def within_inames(self, inames) -> None: def set_temporary_shapes(self, shapes): self._temporary_shapes = shapes -def make_kernel(context: MLIRCodegenContext): - if not isinstance(context, MLIRCodegenContext): - return ValueError("Requires MLIRCodegenContext object") - - '''Making the kernel - Available arguments: - - context.domains - iteration domains - - context.instructions - expression operations - - context.arguments - in/out variables - - name - function name - - Each instruction holds an expression - i.e. - insn_0: p_0 <- dat_0[0] + dat_1[0] - insn_1: idat_0[0 + 1*i_0] <- 1 {dep=insn_0} - insn_2: CODE(idat_0, p_0, dat_1, dat_0) {dep=insn_1} + + @functools.singledispatchmethod + def translate_expr(self, expr) -> SSAValue: + if isinstance(expr, tuple): + breakpoint() + raise ValueError(f"{type(expr)} not implemented yet.") - where the dep informs us of any dependency. - Each InstructionContext may have local variables, such as insn_1->i_0, which can be used to infer that insn_1 will iterate with the temporary variable i_0. + @translate_expr.register(pym.primitives.Subscript) + def _(self, expr: pym.primitives.Subscript) -> SSAValue: + array, index_ssa = self._resolve_subscript(expr) + extract = tensor.ExtractOp.build( + operands=[array, index_ssa], + result_types=[array.type.element_type], + ) + self.builder.insert(extract) + return extract.result + + # NOTE: Could maybe clean up Sum and Product as they are essentially same, just reduction ops. + # TODO: Need to figure out how to solve the dtype inference. Fine to assume i32 for indexing but not for compute generally + @translate_expr.register(pym.primitives.Sum) + def _(self, expr): + children = [self.translate_expr(c) for c in expr.children] + result = children[0] + for child in children[1:]: + result = self._arith_op(arith.AddiOp, arith.AddfOp, result, child) + return result + + @translate_expr.register(pym.primitives.Product) + def _(self, expr): + children = [self.translate_expr(c) for c in expr.children] + result = children[0] + for child in children[1:]: + result = self._arith_op(arith.MuliOp, arith.MulfOp, result, child) + return result + + @translate_expr.register(pym.primitives.Variable) + def _(self, expr: pym.primitives.Variable) -> SSAValue: + return self.symbol_table.lookup(expr.name) + + @translate_expr.register(numbers.Number) + def _(self, expr: numbers.Number) -> SSAValue: + if isinstance(expr, int): + attr = IntegerAttr.from_int_and_width(expr, 32) + else: + attr = FloatAttr(float(expr), f64) - The domain exists come in (start, stop) pairs, so we could develop our scf.for loops appropriately. + const = arith.ConstantOp(attr) + self.builder.insert(const) + return const.result + def _translate_assignment(self, ins): + + match ins.assignee: + case pym.primitives.Variable(): + value = self.translate_expr(ins.expression) + self.symbol_table.insert(ins.assignee.name, value) - ''' - # Update when moving away from loopy-defined argument variables - arg_types = [ - TensorType(arg.dtype, [-1]) # -1 for dynamic, but vector length is known - for arg in self._arguments - ] - - func_op = FuncOp("pyop3_loop", FunctionType.from_lists(arg_types, [])) - block_args = self.parent_func.body.blocks[0].args + case pym.primitives.Subscript(): + value = self.translate_expr(ins.expression) + array, index_ssa = self._resolve_subscript(ins.assignee) + insert = tensor.InsertOp.build( + operands=[value, array, index_ssa], + result_types=[array.type], + ) + self.builder.insert(insert) + self.symbol_table.insert(ins.assignee.aggregate.name, insert.result) - for arg, ssa_value in zip(self._arguments, block_args): - self.symbol_table.insert(arg, ssa_value) - - - return ModuleOp([]) + + def make_kernel(self): + # TODO: Update when moving away from loopy-style-defined argument variables + arg_types = [TensorType(NUMPY_TO_XDSL[arg.dtype], [DYNAMIC_INDEX]) for arg in self._arguments] + + func_op = FuncOp("pyop3_loop", FunctionType.from_lists(arg_types, [])) + + entry = func_op.body.blocks[0] + for arg, ssa in zip(self._arguments, entry.args): + self.symbol_table.insert(arg.name, ssa) + + self.builder = Builder(InsertPoint.at_end(entry)) + + self._build_nest(self._instructions, frozenset()) + + self.builder.insert(func.ReturnOp()) + self.module = ModuleOp([func_op]) + return self.module + + + def _build_nest(self, instructions, entered): + ''' Building nesting order to deal with instructions within loops ''' + + for ins in (i for i in instructions if i.within_inames == entered): + self._translate_assignment(ins) + + deeper = [i for i in instructions if i.within_inames != entered] + if not deeper: + return + + def next_iname(ins): + needed = ins.within_inames - entered + for iname in self._domains: + if iname in needed: + return iname + raise RuntimeError("inconsistent iname state") + + groups: dict[str, list] = {} + for ins in deeper: + groups.setdefault(next_iname(ins), []).append(ins) + + for iname in self._domains: + if iname in groups: + self._build_loop(iname, groups[iname], entered) + + + def _build_loop(self, iname, instructions, entered): + ''' Build scf for loops, this is where change would happen if we want to switch from scf ''' + + start, stop = self._domains[iname] + + lb = self._to_index(self.translate_expr(start)) + ub = self._to_index(self.translate_expr(stop)) + step = self._to_index(self.translate_expr(1)) + + carried = self._written_arrays(instructions) + init_values = [self.symbol_table.lookup(name) for name in carried] + + block_arg_types = [IndexType()] + [v.type for v in init_values] + body = Block(arg_types=block_arg_types) + + for_op = scf.ForOp(lb, ub, step, init_values, Region(body)) + self.builder.insert(for_op) + + with self.symbol_table.scope(): + self.symbol_table.insert(iname, body.args[0]) + # bind carried names to this loop's block args, not the outer values + for name, block_arg in zip(carried, body.args[1:]): + self.symbol_table.insert(name, block_arg) + + old = self.builder + self.builder = Builder(InsertPoint.at_end(body)) + + self._build_nest(instructions, entered | {iname}) + + yielded = [self.symbol_table.lookup(name) for name in carried] + self.builder.insert(scf.YieldOp(*yielded)) + self.builder = old + + # send result/yield back to outer scope + for name, result in zip(carried, for_op.results): + + self.symbol_table.insert(name, result) + + def _to_index(self, value): + if isinstance(value.type, IndexType): + return value + cast = arith.IndexCastOp.build(operands=[value], result_types=[IndexType()]) + self.builder.insert(cast) + return cast.result + + def _written_arrays(self, instructions): + ''' Finding all arrays that are written to within a loop for iterator arguments ''' + written = [] + seen = set() + for ins in instructions: + if isinstance(ins.assignee, pym.primitives.Subscript): + name = ins.assignee.aggregate.name + if name not in seen: + seen.add(name) + written.append(name) + return written + + def _resolve_subscript(self, subscript): + ''' Helper function for indices in tuples ''' + array = self.symbol_table.lookup(subscript.aggregate.name) + indices = subscript.index if isinstance(subscript.index, tuple) else (subscript.index,) + index_ssa = [self._to_index(self.translate_expr(i)) for i in indices] + return array, index_ssa + + def _arith_op(self, int_op, float_op, lhs, rhs): + ''' + Helper function to resolve typing between int and float ops + This is both super ugly and makes the assumption that type is determined by one side + ''' + t = lhs.type + if isinstance(t, (IntegerType, IndexType)): + op = int_op(lhs, rhs) + else: # float type + op = float_op(lhs, rhs) + self.builder.insert(op) + return op.result From 8a9409433d17180a8c06f8c2488391da78d7e50a Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Mon, 24 Aug 2026 15:04:24 +0100 Subject: [PATCH 07/30] refactoring pyop3/lower && integrating MLIR - Goal: Create an interface to a code generation context (MLIR or Loopy) - Status: Interface created, battling PETSc bug before cleaning more. - Goal: Integrate MLIR for auto-generation - Status: Was working but transitioning to pyop3->mlir pipeline as opposed to pym->mlir. Refactoring process is ongoing. --- pyop3/insn/exec.py | 6 +- pyop3/lower/codegen.py | 359 +++++++++-- pyop3/lower/context.py | 100 +++- pyop3/lower/loopy.py | 1296 +++++++++++++++------------------------- pyop3/lower/mlir.py | 16 +- 5 files changed, 860 insertions(+), 917 deletions(-) diff --git a/pyop3/insn/exec.py b/pyop3/insn/exec.py index 3064b278c3..1b1a06928e 100644 --- a/pyop3/insn/exec.py +++ b/pyop3/insn/exec.py @@ -57,7 +57,7 @@ class CompilerParameters: # TODO: handle these - need to build CompilerOptions - codegen: str = "loopy" + codegen: str = "loopy" # extra_cflags: tuple[str, ...] = () # extra_ldflags: tuple[str, ...] = () @@ -237,8 +237,8 @@ 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 + # 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. diff --git a/pyop3/lower/codegen.py b/pyop3/lower/codegen.py index 17a6d89760..9ccf7f9617 100644 --- a/pyop3/lower/codegen.py +++ b/pyop3/lower/codegen.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import abc import collections import contextlib @@ -15,14 +17,62 @@ from typing import Any from weakref import WeakValueDictionary -# 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, *args, **kwargs: op.comm, +from cachetools import cachedmethod +from petsc4py import PETSc + +import loopy as lp +import numpy as np +import pymbolic as pym +from immutabledict import immutabledict as idict + +import pyop3.axis_tree +import pyop3.cache +import pyop3.config +import pyop3.dtypes +import pyop3.expr +from pyop3 import utils, mpi +from pyop3.cache import memory_and_disk_cache +from pyop3.expr import NonlinearDatBufferExpression +from pyop3.expr.visitors import collect_axis_vars, replace +from pyop3.axis_tree.tree import UNIT_AXIS_TREE, IndexedAxisTree, AxisComponent, relabel_path +from pyop3.buffer import AbstractBuffer, ConcreteBuffer, PetscMatBuffer, ArrayBuffer, NullBuffer +from pyop3.dtypes import IntType +from pyop3.lower.transform import with_likwid_markers, with_petsc_event, with_attach_debugger +from pyop3.lower.context import CodegenContext +from pyop3.lower.mlir import MLIRCodegenContext # to remove +from pyop3.lower.loopy import LoopyCodegenContext +from pyop3.insn.base import ( + Intent, + INC, + MAX_RW, + MAX_WRITE, + MIN_RW, + MIN_WRITE, + READ, + RW, + AbstractAssignment, + Exscan, + NullInstruction, + assignment_type_as_intent, + WRITE, + AssignmentType, + ConcretizedNonEmptyArrayAssignment, + StandaloneCalledFunction, + Loop, + InstructionList, ) -def _compile_static(op: InstructionExecutionContext, compiler_parameters: ParsedCompilerParameters) -> tuple: +# TODO: import other way around? +from pyop3.insn.exec import parse_compiler_parameters + +# def _compile_static_hashkey(op: PreprocessedOperation, compiler_parameters: ParsedCompilerParameters) -> Hashable: +# # NOTE: is config valid to include here? +# return (op.disk_cache_key, compiler_parameters, pyop3.config) + +# @pyop3.cache.memory_and_disk_cache( +# hashkey=_compile_static_hashkey, +# get_comm=lambda op, *args, **kwargs: op.comm, +# ) +def _compile_static(op, compiler_parameters) -> Tuple: """Compile the operation without regard for specific data values. This function is therefore suitable for disk caching. @@ -33,79 +83,272 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed datamap """ + insn = op.preprocess() - function_name = "pyop3_loop" # TODO: Provide as kwarg + function_name = "pyop3_loop" if isinstance(insn, InstructionList): cs_expr = insn.instructions else: cs_expr = (insn,) - context = LoopyCodegenContext(check_negatives=compiler_parameters.check_negatives) + # Default to loopy codegen backend + target = getattr(compiler_parameters, 'codegen', 'loopy') + if target == "loopy": + codegen_context = LoopyCodegenContext(check_negatives=compiler_parameters.check_negatives) + elif target == "mlir": + codegen_context = MLIRCodegenContext(check_negatives=compiler_parameters.check_negatives) + # 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: + # explicitly dealing with codegen_contexts at this point. Can always sniff them out again. + # for codegen_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 + # FIXME: removed because cs_expr needs to sniff the codegen_context now loop_indices = {} for e in utils.as_tuple(ex): # TODO: get rid of this loop - # context manager? - context.set_temporary_shapes(_collect_temporary_shapes(e)) - _compile(e, loop_indices, context) + # codegen_context manager? + codegen_context.set_temporary_shapes(_collect_temporary_shapes(e)) + _compile(e, loop_indices, codegen_context) - if not context.global_buffers: + if not codegen_context.global_buffers: + import pyop3.exceptions 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) - - kernel_to_buffer_names = utils.invert_mapping(context._kernel_names) + translation_unit, final_context = codegen_context.finalize_kernel(function_name, compiler_parameters) + + kernel_to_buffer_names = utils.invert_mapping(final_context._kernel_names) buffer_index_map = {} - for kernel_arg in entrypoint.args: + for kernel_arg in translation_unit.default_entrypoint.args: buffer_key = kernel_to_buffer_names[kernel_arg.name] - buffer_ref = context.global_buffers[buffer_key] + buffer_ref = final_context.global_buffers[buffer_key] buffer_index = op.preprocessed_buffers.index(buffer_ref) - intent = context.global_buffer_intents[buffer_key] + intent = final_context.global_buffer_intents[buffer_key] buffer_index_map[kernel_arg.name] = (buffer_index, buffer_ref.nest_indices, intent) - + return translation_unit, buffer_index_map + +@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): + return idict() + +@_collect_temporary_shapes.register(StandaloneCalledFunction) +def _(call): + import loopy as lp # TODO: Remove once StandaloneCalledFunction integrated with MLIR + return idict( + { + (arg.buffer.name, arg.buffer.nest_indices): 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: 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: + _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: + for stmt in loop.statements: + _compile( + stmt, + 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 + + if component.local_size != 1: + iname = codegen_context.unique_name("i") + domain_var = codegen_context.register_extent(component.local_size, iname_map, loop_indices) + codegen_context.add_domain(iname, domain_var) + iname_replace_map_ = iname_map | {axis.label: pym.var(iname)} + within = frozenset({iname}) + else: + iname_replace_map_ = iname_map | {axis.label: 0} + within = set() + + with codegen_context.within_inames(within): + 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 stmt in loop.statements: _compile(stmt, loop_indices, codegen_context) + +@_compile.register(StandaloneCalledFunction) +def _(call, loop_indices, codegen_context): + codegen_context.compile_standalone_function(call, loop_indices) + +@_compile.register(ConcretizedNonEmptyArrayAssignment) +def _(assignment, loop_indices, codegen_context): + if any(isinstance(arg, pyop3.expr.MatPetscMatBufferExpression) for arg in assignment.arguments): + codegen_context.compile_petsc_mat(assignment, loop_indices) + else: + _compile_array_assignment(assignment, loop_indices, codegen_context, assignment.axis_trees) + +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("Refactor needed") + + codegen_context.add_leaf_assignment( + assignment, + 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.local_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: pym.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, + codegen_context, + axis_trees, + iname_replace_maps=new_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_maps, + axis_tree=None, + paths=new_paths + ) + else: + codegen_context.add_leaf_assignment( + assignment, + new_paths, + new_maps, + loop_indices + ) + +@_compile.register(Exscan) +def _(exscan, loop_indices, codegen_context): + codegen_context.compile_exscan(exscan, loop_indices) diff --git a/pyop3/lower/context.py b/pyop3/lower/context.py index a4f4fc5228..85f655bf5a 100644 --- a/pyop3/lower/context.py +++ b/pyop3/lower/context.py @@ -1,12 +1,18 @@ from abc import ABC, abstractmethod +from typing import Any, List, Dict, Tuple +import numbers from pyop3 import utils -from pyop3.insn.base import ( - Intent -) +from pyop3.insn.base import Intent, READ, assignment_type_as_intent class CodegenContext(ABC): - def __init__(self, *, check_negatives): + """ + Abstract base class for code generation contexts. + + Abstract methods required for auto-generating based on _compile_static in codegen.py + """ + + def __init__(self, *, check_negatives: bool): self.check_negatives = check_negatives self._domains = [] @@ -28,59 +34,103 @@ def __init__(self, *, check_negatives): self._assignees = {} @property - def domains(self) -> tuple: + def domains(self) -> Tuple: return tuple(self._domains) @property - def instructions(self) -> tuple: + def instructions(self) -> Tuple: return tuple(self._instructions) @property - def arguments(self) -> tuple: - return tuple(sorted(self._arguments, key=lambda arg: arg.name)) + def arguments(self) -> Tuple: + return tuple(sorted(self._arguments, key=lambda arg: getattr(arg, 'name', ''))) @property - def subkernels(self) -> tuple: + def subkernels(self) -> Tuple: return tuple(self._subkernels) - def __str__(self) -> str: - 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 + # {{{ abstract methods @abstractmethod - def add_domain(self, iname, *args): + def add_domain(self, iname: str, *args) -> None: pass @abstractmethod - def add_assignment(self, assigneee, expression, prefix="insn"): + def add_assignment(self, assignee, expression, prefix: str = "insn") -> None: pass @abstractmethod - def add_function_call(self, assignees, expression, prefix="insn"): + def add_function_call(self, assignees, expression, prefix: str = "insn") -> None: pass @abstractmethod - def add_buffer(self, buffer, intent: Intent | None = None) -> str: + def add_buffer(self, buffer: AbstractBuffer, intent: Intent | None = None) -> str: pass @abstractmethod - def add_subkernel(self, subkernel): + def add_subkernel(self, subkernel) -> None: pass @abstractmethod - def set_temporary_shapes(self, shapes): + def set_temporary_shapes(self, shapes) -> None: pass - def unique_name(self, prefix): + ''' Lowering passes for respective codegen context ''' + @abstractmethod + def lower_expr(self, expr, iname_maps, loop_indices, + intent: Intent = READ, 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: AbstractBuffer, layouts, iname_maps, + loop_indices, intent: Intent): + pass + + @abstractmethod + def add_leaf_assignment(self, assignment, paths, iname_maps, loop_indices): + pass + + @abstractmethod + def register_extent(self, obj: Any, inames, loop_indices): + pass + + @abstractmethod + def compile_standalone_function(self, call, loop_indices): + pass + + @abstractmethod + def compile_petsc_mat(self, assignment, loop_indices): + pass + + @abstractmethod + def compile_exscan(self, call, loop_indices): + pass + + # }}} + + def unique_name(self, prefix: str) -> str: return self._name_generator(prefix) - def _add_instruction(self, insn): + def _add_instruction(self, insn: Any) -> None: self._instructions.append(insn) self._last_insn_id = insn.id @property - def _depends_on(self): - return frozenset({self._last_insn_id}) - {None} + def _depends_on(self) -> frozenset: + if self._last_insn_id is None: + return frozenset() + return frozenset({self._last_insn_id}) + + def __str__(self) -> str: + 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 3af8da952d..0749651273 100644 --- a/pyop3/lower/loopy.py +++ b/pyop3/lower/loopy.py @@ -37,9 +37,8 @@ from pyop3.axis_tree.tree import UNIT_AXIS_TREE, IndexedAxisTree, AxisComponent, relabel_path from pyop3.buffer import AbstractBuffer, ConcreteBuffer, PetscMatBuffer, ArrayBuffer, NullBuffer from pyop3.dtypes import IntType -from pyop3.lower.transform import with_likwid_markers, with_petsc_event, with_attach_debugger from pyop3.lower.context import CodegenContext -from pyop3.lower.mlir import MLIRCodegenContext # to remove +from pyop3.lower.transform import with_likwid_markers, with_petsc_event, with_attach_debugger from pyop3.insn.base import ( Intent, INC, @@ -60,25 +59,307 @@ Loop, InstructionList, ) + # TODO: import other way around? from pyop3.insn.exec import parse_compiler_parameters import pyop3.debug_flags # Sam for debugging - -# 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 LACallable(lp.ScalarCallable, metaclass=abc.ABCMeta): + """ + The LACallable (Linear algebra callable) + replaces loopy.CallInstructions to linear algebra functions + like solve or inverse by LAPACK calls. + """ + 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, + name_in_target=name_in_target) + + @abc.abstractproperty + def name(self): + pass + + @abc.abstractmethod + def generate_preambles(self, target): + pass + + def with_types(self, arg_id_to_dtype, callables_table): + dtypes = {} + for i in range(len(arg_id_to_dtype)): + if arg_id_to_dtype.get(i) is None: + # the types provided aren't mature enough to specialize the + # callable + return (self.copy(arg_id_to_dtype=arg_id_to_dtype), + callables_table) + else: + mat_dtype = arg_id_to_dtype[i].numpy_dtype + dtypes[i] = lp.types.NumpyType(mat_dtype) + dtypes[-1] = lp.types.NumpyType(dtypes[0].dtype) + + return (self.copy(name_in_target=self.name_in_target, + arg_id_to_dtype=idict(dtypes)), + callables_table) + + def emit_call_insn(self, insn, target, expression_to_code_mapper): + assert self.is_ready_for_codegen() + assert isinstance(insn, lp.CallInstruction) + + parameters = insn.expression.parameters + + parameters = list(parameters) + par_dtypes = [self.arg_id_to_dtype[i] for i, _ in enumerate(parameters)] + + parameters.append(insn.assignees[-1]) + par_dtypes.append(self.arg_id_to_dtype[0]) + + mat_descr = self.arg_id_to_descr[0] + arg_c_parameters = [ + expression_to_code_mapper( + par, + pym.mapper.stringifier.PREC_NONE, + lp.expression.dtype_to_type_context(target, par_dtype), + par_dtype + ).expr + for par, par_dtype in zip(parameters, par_dtypes) + ] + c_parameters = [arg_c_parameters[-1]] + c_parameters.extend([arg for arg in arg_c_parameters[:-1]]) + c_parameters.append(np.int32(mat_descr.shape[1])) # n + return pym.var(self.name_in_target)(*c_parameters), False + + +# Read c files for linear algebra callables in on import +if mpi.COMM_WORLD.rank == 0: + with open(os.path.dirname(__file__)+"/inverse.c", "r") as myfile: + inverse_preamble = myfile.read() + with open(os.path.dirname(__file__)+"/solve.c", "r") as myfile: + solve_preamble = myfile.read() +else: + solve_preamble = None + inverse_preamble = None + +inverse_preamble = mpi.COMM_WORLD.bcast(inverse_preamble, root=0) +solve_preamble = mpi.COMM_WORLD.bcast(solve_preamble, root=0) + + +class INVCallable(LACallable): + """ + The InverseCallable replaces loopy.CallInstructions to "inverse" + functions by LAPACK getri. + """ + name = "inverse" + + def generate_preambles(self, target): + assert isinstance(target, type(target)) + yield ("inverse", inverse_preamble) + + +class SolveCallable(LACallable): + """ + The SolveCallable replaces loopy.CallInstructions to "solve" + functions by LAPACK getrs. + """ + name = "solve" + + def generate_preambles(self, target): + assert isinstance(target, type(target)) + yield ("solve", solve_preamble) + class LoopyCodegenContext(CodegenContext): - def __init__(self, *, check_negatives): + def __init__(self, *, check_negatives: bool): super().__init__(check_negatives=check_negatives) self._within_inames = frozenset() - # initializer hash -> temporary name + # initializer hash -> temporary name self._reusable_temporaries: dict[int, str] = {} + ''' Lowering path from PyOP3 to pymbolic ''' + def lower_expr(self, expr, iname_maps, loop_indices, intent=READ, paths=None): + return self._lower_expr(expr, iname_maps, loop_indices, intent=intent, paths=paths) + + @functools.singledispatchmethod + def _lower_expr(self, obj, iname_maps, loop_indices, *, intent, paths): + raise TypeError(f"No Loopy handler defined for {type(obj).__name__}") + + @_lower_expr.register(numbers.Number) + def _(self, num, *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.Modulo, /, *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, iname_maps, *args, **kwargs): + 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.Scalar) + def _(self, scalar, iname_maps, loop_indices, *, intent, **kwargs): + name = self.add_buffer(scalar.buffer, intent) + return pym.subscript(pym.var(name), (0,)) + + @_lower_expr.register(pyop3.expr.ScalarBufferExpression) + def _(self, expr: pyop3.expr.ScalarBufferExpression, /, iname_maps, loop_indices, *, intent, **kwargs) -> pym.Expression: + return self.lower_buffer_access(expr.buffer, [0], iname_maps, loop_indices, intent=intent) + + @_lower_expr.register(pyop3.expr.LinearDatBufferExpression) + def _(self, expr, iname_maps, loop_indices, *, intent, **kwargs): + return self.lower_buffer_access(expr.buffer, [expr.layout], iname_maps, loop_indices, 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, [expr.layouts[path]], iname_maps, loop_indices, intent=intent) + + @_lower_expr.register(pyop3.expr.MatPetscMatBufferExpression) + def _(self, mat_expr, iname_maps, loop_indices, *, intent, paths, **kwargs): + row_path, col_path = paths + layouts = (mat_expr.row_layout.linearize(row_path), mat_expr.column_layout.linearize(col_path)) + return self.lower_buffer_access(mat_expr.buffer, layouts, iname_maps, loop_indices, intent) + + @_lower_expr.register(pyop3.expr.Conditional) + def _(self, cond, iname_maps, loop_indices, **kwargs): + return pym.primitives.If( + self._lower_expr(cond.a, iname_maps, loop_indices, **kwargs), + self._lower_expr(cond.b, iname_maps, loop_indices, **kwargs), + self._lower_expr(cond.c, iname_maps, loop_indices, **kwargs) + ) + + @_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, layouts, iname_maps, loop_indices, intent=intent) + + def lower_buffer_access(self, buffer: AbstractBuffer, layouts, iname_maps, loop_indices, intent): + name_in_kernel = self.add_buffer(buffer, intent) + + # 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, offset_expr) + + subscript = pym.subscript(pym.var(name_in_kernel), indices) + if self.check_negatives and intent == 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 add_leaf_assignment(self, assignment, paths, iname_maps, loop_indices): + intent = assignment_type_as_intent(assignment.assignment_type) + lexpr = self.lower_expr(assignment.assignee, iname_maps, loop_indices, intent=intent, paths=paths) + rexpr = self.lower_expr(assignment.expression, iname_maps, loop_indices, paths=paths) + if assignment.assignment_type == AssignmentType.INC: + rexpr = lexpr + rexpr + self.add_assignment(lexpr, rexpr) + + def maybe_multiindex(self, buffer_ref, offset_expr): + # hack to handle the facbuffer.t that temporaries can have shape but we want to + # linearly index it here + buffer_key = (buffer_ref.name, buffer_ref.nest_indices) + if buffer_key in self._temporary_shapes: + shape = self._temporary_shapes[buffer_key] + 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 + + @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 + def add_domain(self, iname, *args): nargs = len(args) if nargs == 1: @@ -206,7 +487,7 @@ def add_buffer(self, buffer, intent: Intent | None = None) -> str: self._kernel_names[buffer_key] = name_in_kernel return name_in_kernel - def add_temporary(self, prefix="t", dtype=IntType, *, shape=(), initializer: np.ndarray = None, read_only: bool = False) -> str: + def add_temporary(self, prefix="t", dtype=IntType, *, shape=(), initializer: np.ndarray =None, read_only=False) -> str: # If multiple temporaries with the same initializer are used then they # can be shared. can_reuse = initializer is not None and read_only @@ -244,861 +525,218 @@ def add_opaque(self, opaque: OpaqueTerminal, intent) -> str: self._kernel_names[opaque] = name_in_kernel return name_in_kernel - def add_subkernel(self, subkernel): + def add_subkernel(self, subkernel): self._subkernels.append(subkernel) @contextlib.contextmanager - def within_inames(self, inames) -> None: - orig_within_inames = self._within_inames + def within_inames(self, inames): + orig = self._within_inames self._within_inames |= inames yield - self._within_inames = orig_within_inames + self._within_inames = orig - # FIXME, bad API but it is context-dependent def set_temporary_shapes(self, shapes): self._temporary_shapes = shapes - - -class LACallable(lp.ScalarCallable, metaclass=abc.ABCMeta): - """ - The LACallable (Linear algebra callable) - replaces loopy.CallInstructions to linear algebra functions - like solve or inverse by LAPACK calls. - """ - 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, - name_in_target=name_in_target) - - @abc.abstractproperty - def name(self): - pass - - @abc.abstractmethod - def generate_preambles(self, target): - pass - - def with_types(self, arg_id_to_dtype, callables_table): - dtypes = {} - for i in range(len(arg_id_to_dtype)): - if arg_id_to_dtype.get(i) is None: - # the types provided aren't mature enough to specialize the - # callable - return (self.copy(arg_id_to_dtype=arg_id_to_dtype), - callables_table) + def compile_standalone_function( + self, + call: StandaloneCalledFunction, + loop_indices + ) -> 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 = self.add_buffer(arg.buffer, 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: - mat_dtype = arg_id_to_dtype[i].numpy_dtype - dtypes[i] = lp.types.NumpyType(mat_dtype) - dtypes[-1] = lp.types.NumpyType(dtypes[0].dtype) - - return (self.copy(name_in_target=self.name_in_target, - arg_id_to_dtype=idict(dtypes)), - callables_table) - - def emit_call_insn(self, insn, target, expression_to_code_mapper): - assert self.is_ready_for_codegen() - assert isinstance(insn, lp.CallInstruction) - - parameters = insn.expression.parameters - - parameters = list(parameters) - par_dtypes = [self.arg_id_to_dtype[i] for i, _ in enumerate(parameters)] - - parameters.append(insn.assignees[-1]) - par_dtypes.append(self.arg_id_to_dtype[0]) - - mat_descr = self.arg_id_to_descr[0] - arg_c_parameters = [ - expression_to_code_mapper( - par, - pym.mapper.stringifier.PREC_NONE, - lp.expression.dtype_to_type_context(target, par_dtype), - par_dtype - ).expr - for par, par_dtype in zip(parameters, par_dtypes) - ] - c_parameters = [arg_c_parameters[-1]] - c_parameters.extend([arg for arg in arg_c_parameters[:-1]]) - c_parameters.append(np.int32(mat_descr.shape[1])) # n - return pym.var(self.name_in_target)(*c_parameters), False - - -# Read c files for linear algebra callables in on import -if mpi.COMM_WORLD.rank == 0: - with open(os.path.dirname(__file__)+"/inverse.c", "r") as myfile: - inverse_preamble = myfile.read() - with open(os.path.dirname(__file__)+"/solve.c", "r") as myfile: - solve_preamble = myfile.read() -else: - solve_preamble = None - inverse_preamble = None - -inverse_preamble = mpi.COMM_WORLD.bcast(inverse_preamble, root=0) -solve_preamble = mpi.COMM_WORLD.bcast(solve_preamble, root=0) - - -class INVCallable(LACallable): - """ - The InverseCallable replaces loopy.CallInstructions to "inverse" - functions by LAPACK getri. - """ - name = "inverse" - - def generate_preambles(self, target): - assert isinstance(target, type(target)) - yield ("inverse", inverse_preamble) - - -class SolveCallable(LACallable): - """ - The SolveCallable replaces loopy.CallInstructions to "solve" - functions by LAPACK getrs. - """ - name = "solve" - - def generate_preambles(self, target): - assert isinstance(target, type(target)) - yield ("solve", solve_preamble) - - -def _compile_static_hashkey(op: PreprocessedOperation, compiler_parameters: ParsedCompilerParameters) -> Hashable: - # NOTE: is config valid to include here? - return (op.disk_cache_key, compiler_parameters, pyop3.config) - + assert isinstance(loopy_arg, lp.ValueArg) + subarrayrefs[arg] = pym.var(name_in_kernel) -# 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, *args, **kwargs: 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,) - - if compiler_parameters.codegen == "mlir": - context = MLIRCodegenContext(check_negatives=compiler_parameters.check_negatives) - else: - context = LoopyCodegenContext(check_negatives=compiler_parameters.check_negatives) - - # 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 utils.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.global_buffers: - raise pyop3.exceptions.EffectlessComputationException( - "The generated kernel does not modify any global data, this may indicate that something has gone wrong" + 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} ) - - # add a no-op instruction touching all of the kernel arguments so they are - # not silently dropped - # NOTE: Obviously to improve this ugly if-operation - if compiler_parameters.codegen == "loopy": - 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, + 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._instructions.append(noop) - - preambles = [ - ("20_debug", "#include "), # dont always inject - ("30_petsc", "#include "), # perhaps only if petsc callable used? - ] - - if compiler_parameters.codegen == "mlir": - mlir_unit = context.make_kernel() - else: - 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)) - - if compiler_parameters.codegen == "mlir": - breakpoint() - raise NotImplementedError("Still at generation stage") - - - 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) - - kernel_to_buffer_names = utils.invert_mapping(context._kernel_names) - buffer_index_map = {} - for kernel_arg in entrypoint.args: - buffer_key = kernel_to_buffer_names[kernel_arg.name] - buffer_ref = context.global_buffers[buffer_key] - buffer_index = op.preprocessed_buffers.index(buffer_ref) - intent = context.global_buffer_intents[buffer_key] - buffer_index_map[kernel_arg.name] = (buffer_index, buffer_ref.nest_indices, intent) - - if pyop3.debug_flags.hit_assign: - breakpoint() - - return translation_unit, buffer_index_map - - -# put into a class in transform.py? -@functools.singledispatch -def _collect_temporary_shapes(expr): - raise TypeError(f"No handler defined for {type(expr).__name__}") + self.add_function_call(assignees, expression) + subkernel = call.function.code.with_entrypoints(frozenset()) + self.add_subkernel(subkernel) + + def compile_petsc_mat( + self, + assignment: ConcretizedNonEmptyArrayAssignment, + loop_indices + ) -> 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, PetscMatBuffer): + mat = assignment.assignee + expr = assignment.expression + setting_mat_values = True + else: + mat = assignment.expression + expr = assignment.assignee + setting_mat_values = False -@_collect_temporary_shapes.register(InstructionList) -def _(insn_list: InstructionList, /) -> idict: - return utils.merge_dicts(_collect_temporary_shapes(insn) for insn in insn_list) + row_axis_tree, column_axis_tree = assignment.axis_trees -@_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 + assert isinstance(expr, pyop3.expr.BufferExpression) + array_buffer = expr.buffer + # now emit the right line of code, this should properly be a lp.ScalarCallable + # https://petsc.org/release/manualpages/Mat/MatGetValuesLocal/ + mat_name = self.add_buffer(mat.buffer, assignment_type_as_intent(assignment.assignment_type)) -@_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: Is this always correct? It is for now. + array_name = self.add_buffer(array_buffer, READ) + rsize = row_axis_tree.size + csize = column_axis_tree.size -@_collect_temporary_shapes.register -def _(call: StandaloneCalledFunction): - return idict( - { - (arg.buffer.name, arg.buffer.nest_indices): 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.local_size != 1: - - iname = codegen_context.unique_name("i") - domain_var = register_extent( - component.local_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, 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(ConcretizedNonEmptyArrayAssignment) -def parse_assignment(assignment: ConcretizedNonEmptyArrayAssignment, 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, + # these sizes can be expressions that need evaluating + rsize_var = self.register_extent( + rsize, + {}, loop_indices, - context, - assignment.axis_trees, ) + csize_var = self.register_extent( + csize, + {}, + loop_indices, + ) -def _compile_petsc_mat(assignment: ConcretizedNonEmptyArrayAssignment, 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, 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) - array_buffer = expr.buffer - - # 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, assignment_type_as_intent(assignment.assignment_type)) - - # NOTE: Is this always correct? It is for now. - array_name = context.add_buffer(array_buffer, 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.local_size != 1: - iname = codegen_context.unique_name("i") - extent_var = register_extent( - component.local_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} + # 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 = self.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 = self._petsc_mat_store(*myargs) + case AssignmentType.INC: + call_str = self._petsc_mat_add(*myargs) + case _: + raise AssertionError 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) - - if assignment.assignment_type == AssignmentType.INC: - rexpr = 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) - - lexpr = lower_expr(exscan.assignee, [{exscan.scan_axis.label: pym.var(iname)+1}], loop_indices, context, intent=WRITE) - lexpr2 = lower_expr(exscan.assignee, [{exscan.scan_axis.label: pym.var(iname)}], loop_indices, context) - rexpr = lower_expr(exscan.expression, [{exscan.scan_axis.label: pym.var(iname)}], loop_indices, context) - - rexpr = lexpr2 + rexpr - 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) - + call_str = self._petsc_mat_load(*myargs) -@_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.Modulo, /, *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.Scalar) -def _(scalar: pyop3.expr.Scalar, /, iname_maps, loop_indices, context, *, intent, **kwargs) -> pym.Expression: - # TODO: Need a ScalarBufferExpression or similar to encode nested-ness - buffer_ref = scalar.buffer - name_in_kernel = context.add_buffer(buffer_ref, intent) - return pym.subscript(pym.var(name_in_kernel), (0,)) - - -@_lower_expr.register(pyop3.expr.ScalarBufferExpression) -def _(expr: pyop3.expr.ScalarBufferExpression, /, iname_maps, loop_indices, context, *, intent, **kwargs) -> pym.Expression: - return lower_buffer_access(expr.buffer, [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, [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, [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, layouts, iname_maps, loop_indices, context, intent=intent) + self.add_cinstruction(call_str) + def _petsc_mat_load(self, 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]));" -@_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, layouts, iname_maps, loop_indices, context, intent=intent) + def _petsc_mat_store(self, 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 lower_buffer_access(buffer: AbstractBuffer, layouts, iname_maps, loop_indices, context, *, intent) -> pym.Expression: - name_in_kernel = context.add_buffer(buffer, intent) - # 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 + def _petsc_mat_add(self, 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);" + + + def compile_exscan( + self, + exscan: Exscan, + loop_indices + ) -> None: + assert isinstance(exscan, Exscan) + + if exscan.scan_type != "+": + raise NotImplementedError + domain_var = self.register_extent( + exscan.extent, + {}, + loop_indices ) - ) - - # Add some leading zeros to make loopy happy - indices = maybe_multiindex(buffer, offset_expr, context) - - subscript = pym.subscript(pym.var(name_in_kernel), indices) - if context.check_negatives and intent == 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_ref, offset_expr, context): - # hack to handle the facbuffer.t that temporaries can have shape but we want to - # linearly index it here - buffer_key = (buffer_ref.name, buffer_ref.nest_indices) - if buffer_key in context._temporary_shapes: - shape = context._temporary_shapes[buffer_key] - 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,) + iname = self.unique_name("i") + self.add_domain(iname, domain_var) - return indices + lexpr = self.lower_expr(exscan.assignee, [{exscan.scan_axis.label: pym.var(iname)+1}], loop_indices, intent=WRITE) + lexpr2 = self.lower_expr(exscan.assignee, [{exscan.scan_axis.label: pym.var(iname)}], loop_indices) + rexpr = self.lower_expr(exscan.expression, [{exscan.scan_axis.label: pym.var(iname)}], loop_indices) + rexpr = lexpr2 + rexpr + self.add_assignment(lexpr, rexpr) -@_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)) - - -@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 - + def finalize_kernel(self, function_name, compiler_parameters): + preambles = [ + ("20_debug", "#include "), # dont always inject + ("30_petsc", "#include "), # perhaps only inject if petsc callable used + ] -@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 + # Add noop + 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) + + return translation_unit.with_kernel(entrypoint), self diff --git a/pyop3/lower/mlir.py b/pyop3/lower/mlir.py index 0ad06ba1d1..6e3403c43b 100644 --- a/pyop3/lower/mlir.py +++ b/pyop3/lower/mlir.py @@ -1,3 +1,9 @@ +''' + SIGNIFICANT REWRITE OF THIS CLASS DUE TO FOLLOWING CHANGE: + FROM PYM->MLIR + TO PYOP3->MLIR +''' + import contextlib import functools import numbers @@ -209,6 +215,14 @@ def translate_expr(self, expr) -> SSAValue: if isinstance(expr, tuple): breakpoint() raise ValueError(f"{type(expr)} not implemented yet.") + + + @translate_expr.register(pyop3.expr.Scalar) + def _(self, scalar: pyop3.expr.Scalar): + buffer_ref = scalar.buffer + name_in_kernel = context.add_buffer(buffer_ref) + return buffer_ref + @translate_expr.register(pym.primitives.Subscript) def _(self, expr: pym.primitives.Subscript) -> SSAValue: @@ -342,7 +356,6 @@ def _build_loop(self, iname, instructions, entered): old = self.builder self.builder = Builder(InsertPoint.at_end(body)) - self._build_nest(instructions, entered | {iname}) yielded = [self.symbol_table.lookup(name) for name in carried] @@ -351,7 +364,6 @@ def _build_loop(self, iname, instructions, entered): # send result/yield back to outer scope for name, result in zip(carried, for_op.results): - self.symbol_table.insert(name, result) def _to_index(self, value): From f50a0f756c9117e8ef30cb730bfceadf2d347e2c Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Mon, 24 Aug 2026 15:09:08 +0100 Subject: [PATCH 08/30] introducing dtypes for type-inference and testing case --- assign_offloading_demo.py | 5 +++-- pyop3/expr/base.py | 47 +++++++++++++++++++++++++++++++++++++-- 2 files changed, 48 insertions(+), 4 deletions(-) diff --git a/assign_offloading_demo.py b/assign_offloading_demo.py index 8c1555a45d..ead9195f58 100644 --- a/assign_offloading_demo.py +++ b/assign_offloading_demo.py @@ -12,9 +12,10 @@ gpu = op3.CUDAGPU() -pyop3.debug_flags.hit_assign = True -g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile", compiler_parameters={"codegen": "mlir"}) pyop3.debug_flags.hit_assign = False +g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile", compiler_parameters={"codegen": "loopy"}) +pyop3.debug_flags.hit_assign = False + # with op3.offloading(gpu): # g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile") # assert isinstance(g.dat.data_ro, cp.ndarray) # Device diff --git a/pyop3/expr/base.py b/pyop3/expr/base.py index 752a380ad3..9f80ac5dec 100644 --- a/pyop3/expr/base.py +++ b/pyop3/expr/base.py @@ -17,7 +17,6 @@ from pyop3.axis_tree import UNIT_AXIS_TREE, AxisTree, merge_axis_trees from pyop3.axis_tree.tree import MissingVariableException - class Expression(Node, abc.ABC): # {{{ abstract methods @@ -30,6 +29,11 @@ def local_max(self) -> numbers.Number: def local_min(self) -> numbers.Number: raise NotImplementedError + @property + @abc.abstractmethod + def dtype(self) -> np.dtype: + pass + @property @abc.abstractmethod def _full_str(self) -> str: @@ -257,6 +261,33 @@ def get_disk_cache_key(self, visitor): def operands(self) -> tuple[ExpressionT, ExpressionT]: return (self.a, self.b) + @property + def dtype(self) -> np.dtype: + + if isinstance(self.a, Expression): + a_dtype = np.dtype(self.a.dtype) + else: + a_dtype = np.dtype(type(self.a)) + + if isinstance(self.b, Expression): + b_dtype = np.dtype(self.b.dtype) + else: + b_dtype = np.dtype(type(self.b)) + + is_a_float = np.issubdtype(a_dtype, np.floating) + is_b_float = np.issubdtype(b_dtype, np.floating) + + if is_a_float or is_b_float: + # Keep only float dtypes and pick the highest precision one + float_dtypes = [ + dt for dt, is_float in [(a_dtype, is_a_float), (b_dtype, is_b_float)] + if is_float + ] + return max(float_dtypes, key=lambda dt: dt.itemsize) + + return max(a_dtype, b_dtype, key=lambda dt: dt.itemsize) + + # }}} # {{{ abstract methods @@ -294,7 +325,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) - + # }}} @@ -558,6 +589,10 @@ def local_max(self) -> numbers.Number: def local_min(self) -> numbers.Number: raise TypeError("not sure that this makes sense") + @property + def dtype(self) -> np.dtype: + raise TypeError("Not sure this makes sense") + @property def _full_str(self) -> str: return f"i_{{{self.axis.label}}}" @@ -584,6 +619,10 @@ def local_max(self) -> NoReturn: def local_min(self) -> NoReturn: raise TypeError + @property + def dtype(self) -> np.dtype: + return None + _full_str = "NaN" # }}} @@ -639,6 +678,10 @@ def local_max(self) -> numbers.Number: def local_min(self) -> numbers.Number: raise TypeError("not sure that this makes sense") + @property + def dtype(self) -> np.dtype: + return TypeError("not sure that this makes sense") # possibly int32/64? + @property def _full_str(self) -> str: return f"L_{{{self.loop_index.id}, {self.axis.label}}}" From b573e660e5250d0640f02e0b0e72c3edefa287f6 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Mon, 24 Aug 2026 16:17:20 +0100 Subject: [PATCH 09/30] push before connorjward/pyop3 merge update attempt requirements tracking pyop3 && typing hints --- pyop3/lower/context.py | 1 + requirements-build.txt | 6 +++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/pyop3/lower/context.py b/pyop3/lower/context.py index 85f655bf5a..66230dc6ae 100644 --- a/pyop3/lower/context.py +++ b/pyop3/lower/context.py @@ -3,6 +3,7 @@ import numbers from pyop3 import utils +from pyop3.buffer import AbstractBuffer from pyop3.insn.base import Intent, READ, assignment_type_as_intent class CodegenContext(ABC): diff --git a/requirements-build.txt b/requirements-build.txt index f202ca3776..2d2cdf2044 100644 --- a/requirements-build.txt +++ b/requirements-build.txt @@ -1,16 +1,16 @@ # Core build dependencies (adapted from pyproject.toml) Cython>=3.0 -firedrake-rtree +firedrake-rtree>=2026.2.0 libsupermesh>=2026.0 mpi4py>3; python_version >= '3.13' mpi4py; python_version < '3.13' numpy pkgconfig -petsctools @ git+https://github.com/firedrakeproject/petsctools.git@connorjward/cpetsc +petsctools @ git+https://github.com/firedrakeproject/petsctools.git@main pybind11 setuptools>=77.0.3 # Transitive build dependencies hatchling meson-python -scikit_build_core +scikit_build_core \ No newline at end of file From 3358de0bb788c93dcaeb456f20e292a23809ed53 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Tue, 25 Aug 2026 13:11:22 +0100 Subject: [PATCH 10/30] Updating lowering paths for new context class PyOP3 updates and path traversals reflected in the new lowering structure --- pyop3/insn/exec.py | 2 +- pyop3/lower/codegen.py | 233 ++++++++++++++++++++++------------------- pyop3/lower/context.py | 53 ++++++---- 3 files changed, 159 insertions(+), 129 deletions(-) diff --git a/pyop3/insn/exec.py b/pyop3/insn/exec.py index 7353bd8632..ed3210f6ec 100644 --- a/pyop3/insn/exec.py +++ b/pyop3/insn/exec.py @@ -687,7 +687,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 index 9ccf7f9617..fc2ed036ab 100644 --- a/pyop3/lower/codegen.py +++ b/pyop3/lower/codegen.py @@ -1,78 +1,66 @@ from __future__ import annotations import abc -import collections import contextlib -import ctypes -import dataclasses -import enum import functools -import os import numbers -import textwrap -import warnings -import weakref -from collections.abc import Mapping -from functools import cached_property +import os from typing import Any -from weakref import WeakValueDictionary - -from cachetools import cachedmethod -from petsc4py import PETSc 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 utils, mpi -from pyop3.cache import memory_and_disk_cache -from pyop3.expr import NonlinearDatBufferExpression -from pyop3.expr.visitors import collect_axis_vars, replace -from pyop3.axis_tree.tree import UNIT_AXIS_TREE, IndexedAxisTree, AxisComponent, relabel_path -from pyop3.buffer import AbstractBuffer, ConcreteBuffer, PetscMatBuffer, ArrayBuffer, NullBuffer +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.lower.transform import with_likwid_markers, with_petsc_event, with_attach_debugger -from pyop3.lower.context import CodegenContext -from pyop3.lower.mlir import MLIRCodegenContext # to remove -from pyop3.lower.loopy import LoopyCodegenContext from pyop3.insn.base import ( - Intent, - INC, - MAX_RW, - MAX_WRITE, - MIN_RW, - MIN_WRITE, - READ, - RW, AbstractAssignment, + AssignmentType, Exscan, + InstructionList, + Loop, + NonEmptyArrayAssignment, NullInstruction, - assignment_type_as_intent, - WRITE, - AssignmentType, - ConcretizedNonEmptyArrayAssignment, StandaloneCalledFunction, - Loop, - InstructionList, + assignment_type_as_intent, ) + +from pyop3.lower.loopy import LoopyCodegenContext + # TODO: import other way around? -from pyop3.insn.exec import parse_compiler_parameters +from pyop3.lower.transform import ( + with_attach_debugger, + with_likwid_markers, + with_petsc_event, +) -# def _compile_static_hashkey(op: PreprocessedOperation, compiler_parameters: ParsedCompilerParameters) -> Hashable: -# # NOTE: is config valid to include here? -# return (op.disk_cache_key, compiler_parameters, pyop3.config) +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, *args, **kwargs: op.comm, -# ) -def _compile_static(op, compiler_parameters) -> Tuple: +@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. @@ -83,55 +71,60 @@ def _compile_static(op, compiler_parameters) -> Tuple: datamap """ - insn = op.preprocess() - function_name = "pyop3_loop" + function_name = "pyop3_loop" # TODO: Provide as kwarg if isinstance(insn, InstructionList): cs_expr = insn.instructions else: cs_expr = (insn,) - # Default to loopy codegen backend - target = getattr(compiler_parameters, 'codegen', 'loopy') - if target == "loopy": - codegen_context = LoopyCodegenContext(check_negatives=compiler_parameters.check_negatives) - elif target == "mlir": - codegen_context = MLIRCodegenContext(check_negatives=compiler_parameters.check_negatives) - + if compiler_parameters.codegen == "loopy": + ContextClass = LoopyCodegenContext + elif compiler_parameters.codegen == "mlir": + raise NotImplementedError("Still implementing this class") + + context = ContextClass( + 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 codegen_contexts at this point. Can always sniff them out again. - # for codegen_context, ex in cs_expr: + # 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 codegen_context now + # FIXME: removed because cs_expr needs to sniff the context now loop_indices = {} - for e in utils.as_tuple(ex): # TODO: get rid of this loop - # codegen_context manager? - codegen_context.set_temporary_shapes(_collect_temporary_shapes(e)) - _compile(e, loop_indices, codegen_context) + 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 codegen_context.global_buffers: - import pyop3.exceptions + 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, final_context = codegen_context.finalize_kernel(function_name, compiler_parameters) - - kernel_to_buffer_names = utils.invert_mapping(final_context._kernel_names) - buffer_index_map = {} + 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: - buffer_key = kernel_to_buffer_names[kernel_arg.name] - buffer_ref = final_context.global_buffers[buffer_key] - buffer_index = op.preprocessed_buffers.index(buffer_ref) - intent = final_context.global_buffer_intents[buffer_key] - buffer_index_map[kernel_arg.name] = (buffer_index, buffer_ref.nest_indices, intent) - - return translation_unit, buffer_index_map + 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 _collect_temporary_shapes(expr): @@ -162,7 +155,7 @@ def _(assignment): @_collect_temporary_shapes.register(StandaloneCalledFunction) def _(call): - import loopy as lp # TODO: Remove once StandaloneCalledFunction integrated with MLIR + import loopy as lp # TODO: Remove once StandaloneCalledFunction/similar integrated with MLIR return idict( { (arg.buffer.name, arg.buffer.nest_indices): lp_arg.shape @@ -197,28 +190,31 @@ def _( loop_indices, codegen_context ) -> None: - _parse_loop_properly_this_time( + 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 +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: - for stmt in loop.statements: + # NOTE: might need an expression here sometimes + for statement in loop.statements: _compile( - stmt, - loop_indices, - codegen_context + statement, + # loop_indices | dict(loop_exprs), + loop_indices, + codegen_context, ) return @@ -229,40 +225,65 @@ def _parse_loop_properly_this_time( for component in axis.components: path_ = path | {axis.label: component.label} - if axis_tree.linearize(path_, partial=True).size == 0: continue - - if component.local_size != 1: + + 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.local_size, iname_map, loop_indices) + 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: pym.var(iname)} - within = frozenset({iname}) + within_inames = frozenset({iname}) else: iname_replace_map_ = iname_map | {axis.label: 0} - within = set() + within_inames = set() - with codegen_context.within_inames(within): + 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_) + 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 stmt in loop.statements: _compile(stmt, loop_indices, codegen_context) + for statement in loop.statements: + _compile( + statement, + loop_indices, + codegen_context, + ) @_compile.register(StandaloneCalledFunction) def _(call, loop_indices, codegen_context): codegen_context.compile_standalone_function(call, loop_indices) -@_compile.register(ConcretizedNonEmptyArrayAssignment) -def _(assignment, loop_indices, codegen_context): +@_compile.register(NonEmptyArrayAssignment) +def parse_assignment(assignment: NonEmptyArrayAssignment, loop_indices, context: CodegenContext): if any(isinstance(arg, pyop3.expr.MatPetscMatBufferExpression) for arg in assignment.arguments): - codegen_context.compile_petsc_mat(assignment, loop_indices) + context.compile_petsc_mat(assignment, loop_indices) else: - _compile_array_assignment(assignment, loop_indices, codegen_context, assignment.axis_trees) + compile_array_assignment( + assignment, + loop_indices, + context, + assignment.axis_trees, + ) -def _compile_array_assignment( +# NOTE: Move this? Weird to have this one here and rest in context classes. +# Probably move to context.py if I can remove pym references. +def compile_array_assignment( assignment, loop_indices, codegen_context, @@ -322,7 +343,7 @@ def _compile_array_assignment( with codegen_context.within_inames(within_inames): if axis_tree.node_map[new_paths[-1]]: - _compile_array_assignment( + compile_array_assignment( assignment, loop_indices, codegen_context, @@ -332,7 +353,7 @@ def _compile_array_assignment( paths=new_paths ) elif axis_trees: - _compile_array_assignment( + compile_array_assignment( assignment, loop_indices, codegen_context, diff --git a/pyop3/lower/context.py b/pyop3/lower/context.py index 66230dc6ae..f05302f915 100644 --- a/pyop3/lower/context.py +++ b/pyop3/lower/context.py @@ -3,8 +3,9 @@ import numbers from pyop3 import utils -from pyop3.buffer import AbstractBuffer -from pyop3.insn.base import Intent, READ, assignment_type_as_intent +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 class CodegenContext(ABC): """ @@ -13,8 +14,9 @@ class CodegenContext(ABC): Abstract methods required for auto-generating based on _compile_static in codegen.py """ - def __init__(self, *, check_negatives: bool): - self.check_negatives = check_negatives + 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 = [] @@ -24,12 +26,11 @@ def __init__(self, *, check_negatives: bool): self._name_generator = utils.UniqueNameGenerator() - # buffer name -> name in kernel - self._kernel_names = {} + # (buffer, nest_indices) -> name in kernel + self.kernel_names = {} # buffer name -> buffer - self.global_buffers = {} - self.global_buffer_intents = {} + self.buffer_intents = {} # assignee name -> indirection expression self._assignees = {} @@ -50,6 +51,14 @@ def arguments(self) -> Tuple: 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 @@ -65,7 +74,7 @@ def add_function_call(self, assignees, expression, prefix: str = "insn") -> None pass @abstractmethod - def add_buffer(self, buffer: AbstractBuffer, intent: Intent | None = None) -> str: + def add_buffer(self, buffer_view: IndexedBuffer, intent: Intent | None = None) -> str: pass @abstractmethod @@ -79,7 +88,7 @@ def set_temporary_shapes(self, shapes) -> None: ''' Lowering passes for respective codegen context ''' @abstractmethod def lower_expr(self, expr, iname_maps, loop_indices, - intent: Intent = READ, paths = None): + intent: Intent | None = None, paths = None): """ Lower a PyOP3 expression to the target's IR representation. @@ -90,8 +99,15 @@ def lower_expr(self, expr, iname_maps, loop_indices, pass @abstractmethod - def lower_buffer_access(self, buffer: AbstractBuffer, layouts, iname_maps, - loop_indices, intent: Intent): + def lower_buffer_access( + self, + buffer: IndexedBuffer, + layouts, + iname_maps, + loop_indices, + *, + intent + ): pass @abstractmethod @@ -116,19 +132,12 @@ def compile_exscan(self, call, loop_indices): # }}} + def add_subkernel(self, subkernel): + self._subkernels.append(subkernel) + def unique_name(self, prefix: str) -> str: return self._name_generator(prefix) - def _add_instruction(self, insn: Any) -> None: - self._instructions.append(insn) - self._last_insn_id = insn.id - - @property - def _depends_on(self) -> frozenset: - if self._last_insn_id is None: - return frozenset() - return frozenset({self._last_insn_id}) - def __str__(self) -> str: ctx = f"Domain: {str(self.domains)}\n\n" ctx += f"Instructions: {str(self.instructions)}\n\n" From 5a1f473d0c631da87b556c2d80d868904eaaa048 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Tue, 25 Aug 2026 13:44:03 +0100 Subject: [PATCH 11/30] folder for testing with mlir files --- mlir-testing/assign_local_size.mlir | 35 ++++++ mlir-testing/assign_local_size.txt | 29 +++++ mlir-testing/assign_mlir_demo.py | 17 +++ mlir-testing/assign_size.txt | 42 +++++++ mlir-testing/fixed_output.mlir | 159 +++++++++++++++++++++++++ mlir-testing/integral_loop.txt | 68 +++++++++++ mlir-testing/loopy_assign.txt | 38 ++++++ mlir-testing/mlir_assign.mlir | 172 ++++++++++++++++++++++++++++ mlir-testing/mlir_assign_fixed.mlir | 152 ++++++++++++++++++++++++ mlir-testing/mlir_assign_opt.mlir | 20 ++++ mlir-testing/output.ll | 153 +++++++++++++++++++++++++ mlir-testing/run_compiled_mlir.py | 52 +++++++++ mlir-testing/sample.py | 98 ++++++++++++++++ mlir-testing/sample_2.py | 65 +++++++++++ 14 files changed, 1100 insertions(+) create mode 100644 mlir-testing/assign_local_size.mlir create mode 100644 mlir-testing/assign_local_size.txt create mode 100644 mlir-testing/assign_mlir_demo.py create mode 100644 mlir-testing/assign_size.txt create mode 100644 mlir-testing/fixed_output.mlir create mode 100644 mlir-testing/integral_loop.txt create mode 100644 mlir-testing/loopy_assign.txt create mode 100644 mlir-testing/mlir_assign.mlir create mode 100644 mlir-testing/mlir_assign_fixed.mlir create mode 100644 mlir-testing/mlir_assign_opt.mlir create mode 100644 mlir-testing/output.ll create mode 100644 mlir-testing/run_compiled_mlir.py create mode 100644 mlir-testing/sample.py create mode 100644 mlir-testing/sample_2.py diff --git a/mlir-testing/assign_local_size.mlir b/mlir-testing/assign_local_size.mlir new file mode 100644 index 0000000000..0af0ddd9cd --- /dev/null +++ b/mlir-testing/assign_local_size.mlir @@ -0,0 +1,35 @@ +builtin.module { + func.func @pyop3_loop( + %dat_0: tensor, + %dat_1: tensor, + %idat_0: tensor, + %idat_1: tensor, + %idat_2: tensor, + %idat_3: tensor + ) -> tensor { + + %c0 = arith.constant 0 : index // iter var + %c1 = arith.constant 1 : index // iter var + %c17 = arith.constant 17 : index // + %c15 = arith.constant 15 : index + %c32 = arith.constant 32 : index + %c2f = arith.constant 2.0 : f64 + + scf.for %i_0 = %c0 to %c17 step %c1 iter_args() -> () {} + + %f2 = scf.for %i_2 = %c0 to %c15 step %c1 iter_args(%dat_it = %dat_0) -> (tensor) { + %e1 = tensor.extract %idat_2[%i_2] : tensor + %ii_2 = arith.index_cast %e1 : i32 to index + %e2 = tensor.extract %idat_0[%ii_2] : tensor + %iii_2 = arith.index_cast %e2 : i32 to index + %v1 = tensor.extract %dat_1[%iii_2] : tensor + %v2 = arith.mulf %c2f, %v1 : f64 + %res = tensor.insert %v2 into %dat_it[%iii_2] : tensor + scf.yield %res : tensor + } + + scf.for %i_3 = %c0 to %c32 step %c1 iter_args() -> () {} + + func.return %f2 : tensor + } +} diff --git a/mlir-testing/assign_local_size.txt b/mlir-testing/assign_local_size.txt new file mode 100644 index 0000000000..d5e7205a71 --- /dev/null +++ b/mlir-testing/assign_local_size.txt @@ -0,0 +1,29 @@ +******************************************************************************** +#include +#include +#include +#include + +void pyop3_loop(double *__restrict__ dat_0, double const *__restrict__ dat_1, int32_t const *__restrict__ idat_0, int32_t const *__restrict__ idat_1, int32_t const *__restrict__ idat_2, int32_t const *__restrict__ idat_3) +{ + for (int32_t i_0 = 0; i_0 <= 17; ++i_0) + { + } + for (int32_t i_2 = 0; i_2 <= 15; ++i_2) + dat_0[idat_0[idat_2[i_2]]] = 2.0 * dat_1[idat_0[idat_2[i_2]]]; + for (int32_t i_3 = 0; i_3 <= 32; ++i_3) + { + } + +} +******************************************************************************** +dat_0 (16) : [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] +dat_1 (16) : [10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10.] +idat_0 (67) : [ 0 0 0 0 0 1 2 3 3 3 3 4 4 4 4 5 5 5 5 6 6 6 6 7 + 7 7 7 8 8 8 8 9 9 9 9 9 9 10 10 10 10 11 11 11 11 12 12 12 + 12 13 13 13 13 13 13 13 13 14 14 14 14 15 15 15 15 15 15] +idat_1 (18) : [ 0 7 11 15 19 23 27 31 33 37 41 45 49 51 53 57 61 63] +idat_2 (16) : [ 4 5 6 10 14 18 22 26 30 36 40 44 48 56 60 66] +idat_3 (33) : [ 1 2 3 8 9 12 13 16 17 20 21 24 25 28 29 32 34 35 38 39 42 43 46 47 + 50 52 54 55 58 59 62 64 65] +******************************************************************************** diff --git a/mlir-testing/assign_mlir_demo.py b/mlir-testing/assign_mlir_demo.py new file mode 100644 index 0000000000..6fca2ae981 --- /dev/null +++ b/mlir-testing/assign_mlir_demo.py @@ -0,0 +1,17 @@ +from firedrake import * +import pyop3 as op3 +import numpy as np + +mesh = UnitSquareMesh(3,3) + +V = FunctionSpace(mesh, "CG", 1) +f = Function(V).assign(10) +g = Function(V) + +g.dat.assign( + 2 * f.dat, + eager=True, + eager_strategy="compile", + compiler_parameters={"codegen": "mlir"} +) +assert (g.dat.data_ro == 20).all() diff --git a/mlir-testing/assign_size.txt b/mlir-testing/assign_size.txt new file mode 100644 index 0000000000..21b8e370be --- /dev/null +++ b/mlir-testing/assign_size.txt @@ -0,0 +1,42 @@ +******************************************************************************** +#include +#include +#include +#include + +void pyop3_loop(int64_t const *__restrict__ dat_0, int64_t const *__restrict__ dat_1, double *__restrict__ dat_2, double const *__restrict__ dat_3, int64_t const *__restrict__ dat_4, int64_t const *__restrict__ dat_5, int64_t const *__restrict__ dat_6, int64_t const *__restrict__ dat_7, int32_t const *__restrict__ idat_0, int32_t const *__restrict__ idat_1, int32_t const *__restrict__ idat_2, int32_t const *__restrict__ idat_3) +{ + int32_t p_0; + int32_t p_1; + int32_t p_2; + + p_0 = (int32_t) (dat_0[0] + dat_1[0]); + for (int32_t i_0 = 0; i_0 <= -1 + p_0; ++i_0) + { + } + p_1 = (int32_t) (dat_4[0] + dat_5[0]); + for (int32_t i_2 = 0; i_2 <= -1 + p_1; ++i_2) + dat_2[idat_0[idat_2[i_2]]] = 2.0 * dat_3[idat_0[idat_2[i_2]]]; + p_2 = (int32_t) (dat_6[0] + dat_7[0]); + for (int32_t i_3 = 0; i_3 <= -1 + p_2; ++i_3) + { + } + +} +******************************************************************************** +dat_0 (1) : [18] +dat_1 (1) : [0] +dat_2 (16) : [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] +dat_3 (16) : [10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10.] +dat_4 (1) : [16] +dat_5 (1) : [0] +dat_6 (1) : [33] +dat_7 (1) : [0] +idat_0 (67) : [ 0 0 0 0 0 1 2 3 3 3 3 4 4 4 4 5 5 5 5 6 6 6 6 7 + 7 7 7 8 8 8 8 9 9 9 9 9 9 10 10 10 10 11 11 11 11 12 12 12 + 12 13 13 13 13 13 13 13 13 14 14 14 14 15 15 15 15 15 15] +idat_1 (18) : [ 0 7 11 15 19 23 27 31 33 37 41 45 49 51 53 57 61 63] +idat_2 (16) : [ 4 5 6 10 14 18 22 26 30 36 40 44 48 56 60 66] +idat_3 (33) : [ 1 2 3 8 9 12 13 16 17 20 21 24 25 28 29 32 34 35 38 39 42 43 46 47 + 50 52 54 55 58 59 62 64 65] +******************************************************************************** diff --git a/mlir-testing/fixed_output.mlir b/mlir-testing/fixed_output.mlir new file mode 100644 index 0000000000..07b3b0af03 --- /dev/null +++ b/mlir-testing/fixed_output.mlir @@ -0,0 +1,159 @@ +module { + llvm.func @memrefCopy(i64, !llvm.ptr, !llvm.ptr) + llvm.func @malloc(i64) -> !llvm.ptr + llvm.func @pyop3_loop(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: i64, %arg3: i64, %arg4: i64, %arg5: !llvm.ptr, %arg6: !llvm.ptr, %arg7: i64, %arg8: i64, %arg9: i64, %arg10: !llvm.ptr, %arg11: !llvm.ptr, %arg12: i64, %arg13: i64, %arg14: i64, %arg15: !llvm.ptr, %arg16: !llvm.ptr, %arg17: i64, %arg18: i64, %arg19: i64, %arg20: !llvm.ptr, %arg21: !llvm.ptr, %arg22: i64, %arg23: i64, %arg24: i64, %arg25: !llvm.ptr, %arg26: !llvm.ptr, %arg27: i64, %arg28: i64, %arg29: i64) -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> attributes {llvm.emit_c_interface} { + %0 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %1 = llvm.insertvalue %arg20, %0[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %2 = llvm.insertvalue %arg21, %1[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %3 = llvm.insertvalue %arg22, %2[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %4 = llvm.insertvalue %arg23, %3[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %5 = llvm.insertvalue %arg24, %4[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %6 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %7 = llvm.insertvalue %arg15, %6[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %8 = llvm.insertvalue %arg16, %7[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %9 = llvm.insertvalue %arg17, %8[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %10 = llvm.insertvalue %arg18, %9[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %11 = llvm.insertvalue %arg19, %10[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %12 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %13 = llvm.insertvalue %arg5, %12[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %14 = llvm.insertvalue %arg6, %13[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %15 = llvm.insertvalue %arg7, %14[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %16 = llvm.insertvalue %arg8, %15[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %17 = llvm.insertvalue %arg9, %16[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %18 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %19 = llvm.insertvalue %arg0, %18[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %20 = llvm.insertvalue %arg1, %19[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %21 = llvm.insertvalue %arg2, %20[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %22 = llvm.insertvalue %arg3, %21[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %23 = llvm.insertvalue %arg4, %22[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %24 = llvm.mlir.constant(16 : index) : i64 + %25 = llvm.mlir.constant(2.000000e+00 : f64) : f64 + %26 = llvm.mlir.constant(0 : index) : i64 + %27 = llvm.mlir.constant(1 : index) : i64 + llvm.br ^bb1(%26 : i64) + ^bb1(%28: i64): // 2 preds: ^bb0, ^bb2 + %29 = llvm.icmp "slt" %28, %24 : i64 + llvm.cond_br %29, ^bb2, ^bb3 + ^bb2: // pred: ^bb1 + %30 = llvm.extractvalue %5[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %31 = llvm.extractvalue %5[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %32 = llvm.getelementptr %30[%31] : (!llvm.ptr, i64) -> !llvm.ptr, i32 + %33 = llvm.extractvalue %5[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %34 = llvm.mul %28, %33 overflow : i64 + %35 = llvm.getelementptr inbounds|nuw %32[%34] : (!llvm.ptr, i64) -> !llvm.ptr, i32 + %36 = llvm.load %35 : !llvm.ptr -> i32 + %37 = llvm.sext %36 : i32 to i64 + %38 = llvm.extractvalue %17[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %39 = llvm.extractvalue %17[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %40 = llvm.getelementptr %38[%39] : (!llvm.ptr, i64) -> !llvm.ptr, i32 + %41 = llvm.extractvalue %17[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %42 = llvm.mul %37, %41 overflow : i64 + %43 = llvm.getelementptr inbounds|nuw %40[%42] : (!llvm.ptr, i64) -> !llvm.ptr, i32 + %44 = llvm.load %43 : !llvm.ptr -> i32 + %45 = llvm.sext %44 : i32 to i64 + %46 = llvm.extractvalue %11[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %47 = llvm.extractvalue %11[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %48 = llvm.getelementptr %46[%47] : (!llvm.ptr, i64) -> !llvm.ptr, f64 + %49 = llvm.extractvalue %11[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %50 = llvm.mul %45, %49 overflow : i64 + %51 = llvm.getelementptr inbounds|nuw %48[%50] : (!llvm.ptr, i64) -> !llvm.ptr, f64 + %52 = llvm.load %51 : !llvm.ptr -> f64 + %53 = llvm.fmul %52, %25 : f64 + %54 = llvm.extractvalue %23[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %55 = llvm.extractvalue %23[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %56 = llvm.getelementptr %54[%55] : (!llvm.ptr, i64) -> !llvm.ptr, f64 + %57 = llvm.extractvalue %23[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %58 = llvm.mul %45, %57 overflow : i64 + %59 = llvm.getelementptr inbounds|nuw %56[%58] : (!llvm.ptr, i64) -> !llvm.ptr, f64 + llvm.store %53, %59 : f64, !llvm.ptr + %60 = llvm.add %28, %27 : i64 + llvm.br ^bb1(%60 : i64) + ^bb3: // pred: ^bb1 + %61 = llvm.mlir.constant(1 : index) : i64 + %62 = llvm.extractvalue %23[3] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %63 = llvm.alloca %61 x !llvm.array<1 x i64> : (i64) -> !llvm.ptr + llvm.store %62, %63 : !llvm.array<1 x i64>, !llvm.ptr + %64 = llvm.getelementptr %63[0, %26] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<1 x i64> + %65 = llvm.load %64 : !llvm.ptr -> i64 + %66 = llvm.mlir.constant(1 : index) : i64 + %67 = llvm.mlir.zero : !llvm.ptr + %68 = llvm.getelementptr %67[%65] : (!llvm.ptr, i64) -> !llvm.ptr, f64 + %69 = llvm.ptrtoint %68 : !llvm.ptr to i64 + %70 = llvm.call @malloc(%69) : (i64) -> !llvm.ptr + %71 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %72 = llvm.insertvalue %70, %71[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %73 = llvm.insertvalue %70, %72[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %74 = llvm.mlir.constant(0 : index) : i64 + %75 = llvm.insertvalue %74, %73[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %76 = llvm.insertvalue %65, %75[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %77 = llvm.insertvalue %66, %76[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %78 = llvm.intr.stacksave : !llvm.ptr + %79 = llvm.mlir.constant(1 : i64) : i64 + %80 = llvm.mlir.constant(1 : index) : i64 + %81 = llvm.alloca %80 x !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> : (i64) -> !llvm.ptr + llvm.store %23, %81 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>, !llvm.ptr + %82 = llvm.mlir.poison : !llvm.struct<(i64, ptr)> + %83 = llvm.insertvalue %79, %82[0] : !llvm.struct<(i64, ptr)> + %84 = llvm.insertvalue %81, %83[1] : !llvm.struct<(i64, ptr)> + %85 = llvm.mlir.constant(1 : i64) : i64 + %86 = llvm.mlir.constant(1 : index) : i64 + %87 = llvm.alloca %86 x !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> : (i64) -> !llvm.ptr + llvm.store %77, %87 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>, !llvm.ptr + %88 = llvm.mlir.poison : !llvm.struct<(i64, ptr)> + %89 = llvm.insertvalue %85, %88[0] : !llvm.struct<(i64, ptr)> + %90 = llvm.insertvalue %87, %89[1] : !llvm.struct<(i64, ptr)> + %91 = llvm.mlir.constant(1 : index) : i64 + %92 = llvm.alloca %91 x !llvm.struct<(i64, ptr)> : (i64) -> !llvm.ptr + llvm.store %84, %92 : !llvm.struct<(i64, ptr)>, !llvm.ptr + %93 = llvm.alloca %91 x !llvm.struct<(i64, ptr)> : (i64) -> !llvm.ptr + llvm.store %90, %93 : !llvm.struct<(i64, ptr)>, !llvm.ptr + %94 = llvm.mlir.zero : !llvm.ptr + %95 = llvm.getelementptr %94[1] : (!llvm.ptr) -> !llvm.ptr, f64 + %96 = llvm.ptrtoint %95 : !llvm.ptr to i64 + llvm.call @memrefCopy(%96, %92, %93) : (i64, !llvm.ptr, !llvm.ptr) -> () + llvm.intr.stackrestore %78 : !llvm.ptr + llvm.return %77 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + } + llvm.func @_mlir_ciface_pyop3_loop(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr, %arg3: !llvm.ptr, %arg4: !llvm.ptr, %arg5: !llvm.ptr, %arg6: !llvm.ptr) attributes {llvm.emit_c_interface} { + %0 = llvm.load %arg1 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %1 = llvm.extractvalue %0[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %2 = llvm.extractvalue %0[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %3 = llvm.extractvalue %0[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %4 = llvm.extractvalue %0[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %5 = llvm.extractvalue %0[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %6 = llvm.load %arg2 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %7 = llvm.extractvalue %6[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %8 = llvm.extractvalue %6[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %9 = llvm.extractvalue %6[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %10 = llvm.extractvalue %6[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %11 = llvm.extractvalue %6[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %12 = llvm.load %arg3 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %13 = llvm.extractvalue %12[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %14 = llvm.extractvalue %12[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %15 = llvm.extractvalue %12[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %16 = llvm.extractvalue %12[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %17 = llvm.extractvalue %12[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %18 = llvm.load %arg4 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %19 = llvm.extractvalue %18[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %20 = llvm.extractvalue %18[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %21 = llvm.extractvalue %18[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %22 = llvm.extractvalue %18[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %23 = llvm.extractvalue %18[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %24 = llvm.load %arg5 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %25 = llvm.extractvalue %24[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %26 = llvm.extractvalue %24[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %27 = llvm.extractvalue %24[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %28 = llvm.extractvalue %24[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %29 = llvm.extractvalue %24[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %30 = llvm.load %arg6 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %31 = llvm.extractvalue %30[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %32 = llvm.extractvalue %30[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %33 = llvm.extractvalue %30[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %34 = llvm.extractvalue %30[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %35 = llvm.extractvalue %30[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + %36 = llvm.call @pyop3_loop(%1, %2, %3, %4, %5, %7, %8, %9, %10, %11, %13, %14, %15, %16, %17, %19, %20, %21, %22, %23, %25, %26, %27, %28, %29, %31, %32, %33, %34, %35) : (!llvm.ptr, !llvm.ptr, i64, i64, i64, !llvm.ptr, !llvm.ptr, i64, i64, i64, !llvm.ptr, !llvm.ptr, i64, i64, i64, !llvm.ptr, !llvm.ptr, i64, i64, i64, !llvm.ptr, !llvm.ptr, i64, i64, i64, !llvm.ptr, !llvm.ptr, i64, i64, i64) -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> + llvm.store %36, %arg0 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>, !llvm.ptr + llvm.return + } +} + diff --git a/mlir-testing/integral_loop.txt b/mlir-testing/integral_loop.txt new file mode 100644 index 0000000000..feac230fa3 --- /dev/null +++ b/mlir-testing/integral_loop.txt @@ -0,0 +1,68 @@ +******************************************************************************** +#include +#include +#include +#include +#include + +static void form_cell_integral(double *__restrict__ A, double const *__restrict__ coords_0); +static void form_cell_integral(double *__restrict__ A, double const *__restrict__ coords_0) +{ + double t0; + double t1; + double t2; + double t3[3] = { 0.33333333333333337, 0.33333333333333326, 0.33333333333333326 }; + + t0 = -1.0 * coords_0[0]; + t1 = -1.0 * coords_0[1]; + t2 = 0.5 * fabs((t0 + coords_0[2]) * (t1 + coords_0[5]) + -1.0 * (t0 + coords_0[4]) * (t1 + coords_0[3])); + for (int32_t j = 0; j <= 2; ++j) + A[j] = A[j] + t3[j] * t2; + +} + +void pyop3_loop(double const *__restrict__ dat_0, double *__restrict__ dat_1, int32_t const *__restrict__ idat_0, int32_t const *__restrict__ idat_1) +{ + int32_t j_0; + int32_t j_1; + int32_t j_2; + double t_0[3l]; + double t_1[6l]; + + for (int32_t i_0 = 0; i_0 <= 17; ++i_0) + { + for (int32_t i_1 = 0; i_1 <= 2; ++i_1) + { + j_0 = 0 + 1 * (i_1 + 0); + t_0[j_0] = (double) (0.0); + } + for (int32_t i_2 = 0; i_2 <= 2; ++i_2) + for (int32_t i_3 = 0; i_3 <= 1; ++i_3) + { + j_1 = 0 + 1 * (i_2 * 2 + 0 + i_3); + t_1[j_1] = dat_0[idat_0[3 * i_0 + i_2] + i_3]; + } + form_cell_integral(&(t_0[0]), &(t_1[0])); + for (int32_t i_6 = 0; i_6 <= 2; ++i_6) + { + j_2 = 0 + 1 * (i_6 + 0); + dat_1[idat_1[3 * i_0 + i_6]] = dat_1[idat_1[3 * i_0 + i_6]] + t_0[j_2]; + } + } + +} +******************************************************************************** +dat_0 (32) : [0.33333333 0. 0. 0.33333333 0. 0. + 0.33333333 0.33333333 0. 0.66666667 0.66666667 0. + 0.33333333 0.66666667 0.66666667 0.33333333 0. 1. + 1. 0. 0.33333333 1. 0.66666667 0.66666667 + 1. 0.33333333 0.66666667 1. 1. 0.66666667 + 1. 1. ] +dat_1 (16) : [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] +idat_0 (54) : [ 0 2 4 0 2 6 2 6 8 0 6 10 6 8 12 6 10 14 8 12 16 6 12 14 + 10 14 18 12 16 20 12 14 22 14 18 24 12 20 22 14 22 24 20 22 26 22 24 28 + 22 26 28 26 28 30] +idat_1 (54) : [ 0 1 2 0 1 3 1 3 4 0 3 5 3 4 6 3 5 7 4 6 8 3 6 7 + 5 7 9 6 8 10 6 7 11 7 9 12 6 10 11 7 11 12 10 11 13 11 12 14 + 11 13 14 13 14 15] +******************************************************************************** diff --git a/mlir-testing/loopy_assign.txt b/mlir-testing/loopy_assign.txt new file mode 100644 index 0000000000..d14de3083d --- /dev/null +++ b/mlir-testing/loopy_assign.txt @@ -0,0 +1,38 @@ +--------------------------------------------------------------------------- +KERNEL: pyop3_loop +--------------------------------------------------------------------------- +ARGUMENTS: +dat_0: ArrayArg, type: np:dtype('float64'), shape: unknown in/out aspace: global +dat_1: ArrayArg, type: np:dtype('float64'), shape: unknown in aspace: global +idat_0: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global +idat_1: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global +idat_2: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global +idat_3: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global +--------------------------------------------------------------------------- +DOMAINS: +{ [i_0] : 0 <= i_0 <= 17 } +{ [i_1] : 1 = 0 } +{ [i_2] : 0 <= i_2 <= 15 } +{ [i_3] : 0 <= i_3 <= 32 } +{ [i_4] : 1 = 0 } +--------------------------------------------------------------------------- +INAME TAGS: +i_0: None +i_1: None +i_2: None +i_3: None +i_4: None +--------------------------------------------------------------------------- +INSTRUCTIONS: + for i_1, i_0 +↱ dat_0[0 + 1*(idat_0[0 + 1*idat_1[0 + 1*i_0]] + i_1)] = 2*dat_1[0 + 1*(idat_0[0 + 1*idat_1[0 + 1*i_0]] + i_1)] {id=insn_0} +│ end i_1, i_0 +│ for i_2 +└↱ dat_0[0 + 1*(idat_0[0 + 1*idat_2[0 + 1*i_2]] + 0)] = 2*dat_1[0 + 1*(idat_0[0 + 1*idat_2[0 + 1*i_2]] + 0) {id=insn_1} + │ end i_2 + │ for i_4, i_3 +↱└ dat_0[0 + 1*(idat_0[0 + 1*idat_3[0 + 1*i_3]] + i_4) = 2*dat_1[0 + 1*(idat_0[0 + 1*idat_3[0 + 1*i_3]] + i_4) {id=insn_2} +│ end i_4, i_3 +└ CODE(idat_2, dat_0, idat_0, idat_1, dat_1, idat_3|) {id=insn} + +--------------------------------------------------------------------------- diff --git a/mlir-testing/mlir_assign.mlir b/mlir-testing/mlir_assign.mlir new file mode 100644 index 0000000000..4eab0e02b9 --- /dev/null +++ b/mlir-testing/mlir_assign.mlir @@ -0,0 +1,172 @@ +builtin.module { + func.func @pyop3_loop(%0: tensor, %1: tensor, %2: tensor, %3: tensor, %4: tensor, %5: tensor) { + %6 = arith.constant 0 : i32 + %7 = arith.index_cast %6 : i32 to index + %8 = arith.constant 18 : i32 + %9 = arith.index_cast %8 : i32 to index + %10 = arith.constant 1 : i32 + %11 = arith.index_cast %10 : i32 to index + %12 = scf.for %13 = %7 to %9 step %11 iter_args(%14 = %0) -> (tensor) { + %15 = arith.constant 0 : i32 + %16 = arith.index_cast %15 : i32 to index + %17 = arith.constant 0 : i32 + %18 = arith.index_cast %17 : i32 to index + %19 = arith.constant 1 : i32 + %20 = arith.index_cast %19 : i32 to index + %21 = scf.for %22 = %16 to %18 step %20 iter_args(%23 = %14) -> (tensor) { + %24 = arith.constant 2 : i32 + %25 = arith.constant 0 : i32 + %26 = arith.constant 1 : i32 + %27 = arith.constant 0 : i32 + %28 = arith.constant 1 : i32 + %29 = arith.constant 0 : i32 + %30 = arith.constant 1 : i32 + %31 = arith.muli %30, %13 : i32 + %32 = arith.addi %29, %31 : i32 + %33 = arith.index_cast %32 : i32 to index + %34 = tensor.extract %2[%33] : tensor + %35 = arith.muli %28, %34 : i32 + %36 = arith.addi %27, %35 : i32 + %37 = arith.index_cast %36 : i32 to index + %38 = tensor.extract %1[%37] : tensor + %39 = arith.addi %38, %22 : i32 + %40 = arith.muli %26, %39 : i32 + %41 = arith.addi %25, %40 : i32 + %42 = arith.index_cast %41 : i32 to index + %43 = tensor.extract %3[%42] : tensor + %44 = arith.muli %24, %43 : i32 + %45 = arith.constant 0 : i32 + %46 = arith.constant 1 : i32 + %47 = arith.constant 0 : i32 + %48 = arith.constant 1 : i32 + %49 = arith.constant 0 : i32 + %50 = arith.constant 1 : i32 + %51 = arith.muli %50, %13 : i32 + %52 = arith.addi %49, %51 : i32 + %53 = arith.index_cast %52 : i32 to index + %54 = tensor.extract %2[%53] : tensor + %55 = arith.muli %48, %54 : i32 + %56 = arith.addi %47, %55 : i32 + %57 = arith.index_cast %56 : i32 to index + %58 = tensor.extract %1[%57] : tensor + %59 = arith.addi %58, %22 : i32 + %60 = arith.muli %46, %59 : i32 + %61 = arith.addi %45, %60 : i32 + %62 = arith.index_cast %61 : i32 to index + %63 = tensor.insert %44 into %23[%62] : tensor + scf.yield %63 : tensor + } + scf.yield %21 : tensor + } + %64 = arith.constant 0 : i32 + %65 = arith.index_cast %64 : i32 to index + %66 = arith.constant 16 : i32 + %67 = arith.index_cast %66 : i32 to index + %68 = arith.constant 1 : i32 + %69 = arith.index_cast %68 : i32 to index + %70 = scf.for %71 = %65 to %67 step %69 iter_args(%72 = %12) -> (tensor) { + %73 = arith.constant 2 : i32 + %74 = arith.constant 0 : i32 + %75 = arith.constant 1 : i32 + %76 = arith.constant 0 : i32 + %77 = arith.constant 1 : i32 + %78 = arith.constant 0 : i32 + %79 = arith.constant 1 : i32 + %80 = arith.muli %79, %71 : i32 + %81 = arith.addi %78, %80 : i32 + %82 = arith.index_cast %81 : i32 to index + %83 = tensor.extract %4[%82] : tensor + %84 = arith.muli %77, %83 : i32 + %85 = arith.addi %76, %84 : i32 + %86 = arith.index_cast %85 : i32 to index + %87 = tensor.extract %1[%86] : tensor + %88 = arith.constant 0 : i32 + %89 = arith.addi %87, %88 : i32 + %90 = arith.muli %75, %89 : i32 + %91 = arith.addi %74, %90 : i32 + %92 = arith.index_cast %91 : i32 to index + %93 = tensor.extract %3[%92] : tensor + %94 = arith.muli %73, %93 : i32 + %95 = arith.constant 0 : i32 + %96 = arith.constant 1 : i32 + %97 = arith.constant 0 : i32 + %98 = arith.constant 1 : i32 + %99 = arith.constant 0 : i32 + %100 = arith.constant 1 : i32 + %101 = arith.muli %100, %71 : i32 + %102 = arith.addi %99, %101 : i32 + %103 = arith.index_cast %102 : i32 to index + %104 = tensor.extract %4[%103] : tensor + %105 = arith.muli %98, %104 : i32 + %106 = arith.addi %97, %105 : i32 + %107 = arith.index_cast %106 : i32 to index + %108 = tensor.extract %1[%107] : tensor + %109 = arith.constant 0 : i32 + %110 = arith.addi %108, %109 : i32 + %111 = arith.muli %96, %110 : i32 + %112 = arith.addi %95, %111 : i32 + %113 = arith.index_cast %112 : i32 to index + %114 = tensor.insert %94 into %72[%113] : tensor + scf.yield %114 : tensor + } + %115 = arith.constant 0 : i32 + %116 = arith.index_cast %115 : i32 to index + %117 = arith.constant 33 : i32 + %118 = arith.index_cast %117 : i32 to index + %119 = arith.constant 1 : i32 + %120 = arith.index_cast %119 : i32 to index + %121 = scf.for %122 = %116 to %118 step %120 iter_args(%123 = %70) -> (tensor) { + %124 = arith.constant 0 : i32 + %125 = arith.index_cast %124 : i32 to index + %126 = arith.constant 0 : i32 + %127 = arith.index_cast %126 : i32 to index + %128 = arith.constant 1 : i32 + %129 = arith.index_cast %128 : i32 to index + %130 = scf.for %131 = %125 to %127 step %129 iter_args(%132 = %123) -> (tensor) { + %133 = arith.constant 2 : i32 + %134 = arith.constant 0 : i32 + %135 = arith.constant 1 : i32 + %136 = arith.constant 0 : i32 + %137 = arith.constant 1 : i32 + %138 = arith.constant 0 : i32 + %139 = arith.constant 1 : i32 + %140 = arith.muli %139, %122 : i32 + %141 = arith.addi %138, %140 : i32 + %142 = arith.index_cast %141 : i32 to index + %143 = tensor.extract %5[%142] : tensor + %144 = arith.muli %137, %143 : i32 + %145 = arith.addi %136, %144 : i32 + %146 = arith.index_cast %145 : i32 to index + %147 = tensor.extract %1[%146] : tensor + %148 = arith.addi %147, %131 : i32 + %149 = arith.muli %135, %148 : i32 + %150 = arith.addi %134, %149 : i32 + %151 = arith.index_cast %150 : i32 to index + %152 = tensor.extract %3[%151] : tensor + %153 = arith.muli %133, %152 : i32 + %154 = arith.constant 0 : i32 + %155 = arith.constant 1 : i32 + %156 = arith.constant 0 : i32 + %157 = arith.constant 1 : i32 + %158 = arith.constant 0 : i32 + %159 = arith.constant 1 : i32 + %160 = arith.muli %159, %122 : i32 + %161 = arith.addi %158, %160 : i32 + %162 = arith.index_cast %161 : i32 to index + %163 = tensor.extract %5[%162] : tensor + %164 = arith.muli %157, %163 : i32 + %165 = arith.addi %156, %164 : i32 + %166 = arith.index_cast %165 : i32 to index + %167 = tensor.extract %1[%166] : tensor + %168 = arith.addi %167, %131 : i32 + %169 = arith.muli %155, %168 : i32 + %170 = arith.addi %154, %169 : i32 + %171 = arith.index_cast %170 : i32 to index + %172 = tensor.insert %153 into %132[%171] : tensor + scf.yield %172 : tensor + } + scf.yield %130 : tensor + } + func.return + } +} diff --git a/mlir-testing/mlir_assign_fixed.mlir b/mlir-testing/mlir_assign_fixed.mlir new file mode 100644 index 0000000000..20fdd5bccc --- /dev/null +++ b/mlir-testing/mlir_assign_fixed.mlir @@ -0,0 +1,152 @@ +// 62 +builtin.module { + func.func @pyop3_loop(%0: tensor, %1: tensor, %2: tensor, %3: tensor, %4: tensor, %5: tensor) -> tensor attributes {llvm.emit_c_interface} { + %7 = arith.constant 0 : index + %9 = arith.constant 18 : index + %11 = arith.constant 1 : index + %12 = scf.for %13 = %7 to %9 step %11 iter_args(%14 = %0) -> (tensor) { + %16 = arith.constant 0 : index + %18 = arith.constant 0 : index + %20 = arith.constant 1 : index + %21 = scf.for %22 = %16 to %18 step %20 iter_args(%23 = %14) -> (tensor) { + %24 = arith.constant 2. : f64 + %25 = arith.constant 0 : index + %26 = arith.constant 1 : index + %27 = arith.constant 0 : index + %28 = arith.constant 1 : index + %29 = arith.constant 0 : index + %30 = arith.constant 1 : index + %31 = arith.muli %30, %13 : index + %32 = arith.addi %29, %31 : index + %34 = tensor.extract %2[%32] : tensor + %15 = arith.index_cast %34 : i32 to index + %35 = arith.muli %28, %15 : index + %36 = arith.addi %27, %35 : index + %38 = tensor.extract %1[%36] : tensor + %17 = arith.index_cast %38 : i32 to index + %39 = arith.addi %17, %22 : index + %40 = arith.muli %26, %39 : index + %41 = arith.addi %25, %40 : index + %43 = tensor.extract %3[%41] : tensor + %44 = arith.mulf %24, %43 : f64 + %45 = arith.constant 0 : index + %46 = arith.constant 1 : index + %47 = arith.constant 0 : index + %48 = arith.constant 1 : index + %49 = arith.constant 0 : index + %50 = arith.constant 1 : index + %51 = arith.muli %50, %13 : index + %52 = arith.addi %49, %51 : index + %54 = tensor.extract %2[%52] : tensor + %19 = arith.index_cast %54 : i32 to index + %55 = arith.muli %48, %19 : index + %56 = arith.addi %47, %55 : index + %58 = tensor.extract %1[%56] : tensor + %57 = arith.index_cast %58 : i32 to index + %59 = arith.addi %57, %22 : index + %60 = arith.muli %46, %59 : index + %61 = arith.addi %45, %60 : index + %63 = tensor.insert %44 into %23[%61] : tensor + scf.yield %63 : tensor + } + scf.yield %21 : tensor + } + %65 = arith.constant 0 : index + %67 = arith.constant 16 : index + %69 = arith.constant 1 : index + %70 = scf.for %71 = %65 to %67 step %69 iter_args(%72 = %12) -> (tensor) { + %73 = arith.constant 2. : f64 + %74 = arith.constant 0 : index + %75 = arith.constant 1 : index + %76 = arith.constant 0 : index + %77 = arith.constant 1 : index + %78 = arith.constant 0 : index + %79 = arith.constant 1 : index + %80 = arith.muli %79, %71 : index + %81 = arith.addi %78, %80 : index + %83 = tensor.extract %4[%81] : tensor + %82 = arith.index_cast %83 : i32 to index + %84 = arith.muli %77, %82 : index + %85 = arith.addi %76, %84 : index + %87 = tensor.extract %1[%85] : tensor + %86 = arith.index_cast %87 : i32 to index + %88 = arith.constant 0 : index + %89 = arith.addi %86, %88 : index + %90 = arith.muli %75, %89 : index + %91 = arith.addi %74, %90 : index + %93 = tensor.extract %3[%91] : tensor + %94 = arith.mulf %73, %93 : f64 + %95 = arith.constant 0 : index + %96 = arith.constant 1 : index + %97 = arith.constant 0 : index + %98 = arith.constant 1 : index + %99 = arith.constant 0 : index + %100 = arith.constant 1 : index + %101 = arith.muli %100, %71 : index + %102 = arith.addi %99, %101 : index + %104 = tensor.extract %4[%102] : tensor + %92 = arith.index_cast %104 : i32 to index + %105 = arith.muli %98, %92 : index + %106 = arith.addi %97, %105 : index + %108 = tensor.extract %1[%106] : tensor + %107 = arith.index_cast %108 : i32 to index + %109 = arith.constant 0 : index + %110 = arith.addi %107, %109 : index + %111 = arith.muli %96, %110 : index + %112 = arith.addi %95, %111 : index + %114 = tensor.insert %94 into %72[%112] : tensor + scf.yield %114 : tensor + } + %116 = arith.constant 0 : index + %118 = arith.constant 33 : index + %120 = arith.constant 1 : index + %121 = scf.for %122 = %116 to %118 step %120 iter_args(%123 = %70) -> (tensor) { + %125 = arith.constant 0 : index + %127 = arith.constant 0 : index + %129 = arith.constant 1 : index + %130 = scf.for %131 = %125 to %127 step %129 iter_args(%132 = %123) -> (tensor) { + %133 = arith.constant 2. : f64 + %134 = arith.constant 0 : index + %135 = arith.constant 1 : index + %136 = arith.constant 0 : index + %137 = arith.constant 1 : index + %138 = arith.constant 0 : index + %139 = arith.constant 1 : index + %140 = arith.muli %139, %122 : index + %141 = arith.addi %138, %140 : index + %143 = tensor.extract %5[%141] : tensor + %142 = arith.index_cast %143 : i32 to index + %144 = arith.muli %137, %142 : index + %145 = arith.addi %136, %144 : index + %147 = tensor.extract %1[%145] : tensor + %146 = arith.index_cast %147 : i32 to index + %148 = arith.addi %146, %131 : index + %149 = arith.muli %135, %148 : index + %150 = arith.addi %134, %149 : index + %152 = tensor.extract %3[%150] : tensor + %153 = arith.mulf %133, %152 : f64 + %154 = arith.constant 0 : index + %155 = arith.constant 1 : index + %156 = arith.constant 0 : index + %157 = arith.constant 1 : index + %158 = arith.constant 0 : index + %159 = arith.constant 1 : index + %160 = arith.muli %159, %122 : index + %161 = arith.addi %158, %160 : index + %163 = tensor.extract %5[%161] : tensor + %162 = arith.index_cast %163 : i32 to index + %164 = arith.muli %157, %162 : index + %165 = arith.addi %156, %164 : index + %167 = tensor.extract %1[%165] : tensor + %166 = arith.index_cast %167 : i32 to index + %168 = arith.addi %166, %131 : index + %169 = arith.muli %155, %168 : index + %170 = arith.addi %154, %169 : index + %172 = tensor.insert %153 into %132[%170] : tensor + scf.yield %172 : tensor + } + scf.yield %130 : tensor + } + func.return %121 : tensor + } +} diff --git a/mlir-testing/mlir_assign_opt.mlir b/mlir-testing/mlir_assign_opt.mlir new file mode 100644 index 0000000000..7ca17bd78f --- /dev/null +++ b/mlir-testing/mlir_assign_opt.mlir @@ -0,0 +1,20 @@ +module { + func.func @pyop3_loop(%arg0: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor, %arg5: tensor) -> tensor attributes {llvm.emit_c_interface} { + %c16 = arith.constant 16 : index + %cst = arith.constant 2.000000e+00 : f64 + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %0 = scf.for %arg6 = %c0 to %c16 step %c1 iter_args(%arg7 = %arg0) -> (tensor) { + %extracted = tensor.extract %arg4[%arg6] : tensor + %1 = arith.index_cast %extracted : i32 to index + %extracted_0 = tensor.extract %arg1[%1] : tensor + %2 = arith.index_cast %extracted_0 : i32 to index + %extracted_1 = tensor.extract %arg3[%2] : tensor + %3 = arith.mulf %extracted_1, %cst : f64 + %inserted = tensor.insert %3 into %arg7[%2] : tensor + scf.yield %inserted : tensor + } + return %0 : tensor + } +} + diff --git a/mlir-testing/output.ll b/mlir-testing/output.ll new file mode 100644 index 0000000000..9a0a3411cd --- /dev/null +++ b/mlir-testing/output.ll @@ -0,0 +1,153 @@ +; ModuleID = 'LLVMDialectModule' +source_filename = "LLVMDialectModule" + +declare void @memrefCopy(i64, ptr, ptr) + +declare ptr @malloc(i64) + +define { ptr, ptr, i64, [1 x i64], [1 x i64] } @pyop3_loop(ptr %0, ptr %1, i64 %2, i64 %3, i64 %4, ptr %5, ptr %6, i64 %7, i64 %8, i64 %9, ptr %10, ptr %11, i64 %12, i64 %13, i64 %14, ptr %15, ptr %16, i64 %17, i64 %18, i64 %19, ptr %20, ptr %21, i64 %22, i64 %23, i64 %24, ptr %25, ptr %26, i64 %27, i64 %28, i64 %29) { + %31 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } poison, ptr %20, 0 + %32 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %31, ptr %21, 1 + %33 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %32, i64 %22, 2 + %34 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %33, i64 %23, 3, 0 + %35 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %34, i64 %24, 4, 0 + %36 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } poison, ptr %15, 0 + %37 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %36, ptr %16, 1 + %38 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %37, i64 %17, 2 + %39 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %38, i64 %18, 3, 0 + %40 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %39, i64 %19, 4, 0 + %41 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } poison, ptr %5, 0 + %42 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %41, ptr %6, 1 + %43 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %42, i64 %7, 2 + %44 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %43, i64 %8, 3, 0 + %45 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %44, i64 %9, 4, 0 + %46 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } poison, ptr %0, 0 + %47 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %46, ptr %1, 1 + %48 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %47, i64 %2, 2 + %49 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %48, i64 %3, 3, 0 + %50 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %49, i64 %4, 4, 0 + br label %51 + +51: ; preds = %54, %30 + %52 = phi i64 [ %85, %54 ], [ 0, %30 ] + %53 = icmp slt i64 %52, 16 + br i1 %53, label %54, label %86 + +54: ; preds = %51 + %55 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %35, 1 + %56 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %35, 2 + %57 = getelementptr i32, ptr %55, i64 %56 + %58 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %35, 4, 0 + %59 = mul nuw nsw i64 %52, %58 + %60 = getelementptr inbounds nuw i32, ptr %57, i64 %59 + %61 = load i32, ptr %60, align 4 + %62 = sext i32 %61 to i64 + %63 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %45, 1 + %64 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %45, 2 + %65 = getelementptr i32, ptr %63, i64 %64 + %66 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %45, 4, 0 + %67 = mul nuw nsw i64 %62, %66 + %68 = getelementptr inbounds nuw i32, ptr %65, i64 %67 + %69 = load i32, ptr %68, align 4 + %70 = sext i32 %69 to i64 + %71 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %40, 1 + %72 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %40, 2 + %73 = getelementptr double, ptr %71, i64 %72 + %74 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %40, 4, 0 + %75 = mul nuw nsw i64 %70, %74 + %76 = getelementptr inbounds nuw double, ptr %73, i64 %75 + %77 = load double, ptr %76, align 8 + %78 = fmul double %77, 2.000000e+00 + %79 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %50, 1 + %80 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %50, 2 + %81 = getelementptr double, ptr %79, i64 %80 + %82 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %50, 4, 0 + %83 = mul nuw nsw i64 %70, %82 + %84 = getelementptr inbounds nuw double, ptr %81, i64 %83 + store double %78, ptr %84, align 8 + %85 = add i64 %52, 1 + br label %51 + +86: ; preds = %51 + %87 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %50, 3 + %88 = alloca [1 x i64], i64 1, align 8 + store [1 x i64] %87, ptr %88, align 4 + %89 = getelementptr [1 x i64], ptr %88, i32 0, i64 0 + %90 = load i64, ptr %89, align 4 + %91 = getelementptr double, ptr null, i64 %90 + %92 = ptrtoint ptr %91 to i64 + %93 = call ptr @malloc(i64 %92) + %94 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } poison, ptr %93, 0 + %95 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %94, ptr %93, 1 + %96 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %95, i64 0, 2 + %97 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %96, i64 %90, 3, 0 + %98 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %97, i64 1, 4, 0 + %99 = call ptr @llvm.stacksave.p0() + %100 = alloca { ptr, ptr, i64, [1 x i64], [1 x i64] }, i64 1, align 8 + store { ptr, ptr, i64, [1 x i64], [1 x i64] } %50, ptr %100, align 8 + %101 = insertvalue { i64, ptr } { i64 1, ptr poison }, ptr %100, 1 + %102 = alloca { ptr, ptr, i64, [1 x i64], [1 x i64] }, i64 1, align 8 + store { ptr, ptr, i64, [1 x i64], [1 x i64] } %98, ptr %102, align 8 + %103 = insertvalue { i64, ptr } { i64 1, ptr poison }, ptr %102, 1 + %104 = alloca { i64, ptr }, i64 1, align 8 + store { i64, ptr } %101, ptr %104, align 8 + %105 = alloca { i64, ptr }, i64 1, align 8 + store { i64, ptr } %103, ptr %105, align 8 + call void @memrefCopy(i64 8, ptr %104, ptr %105) + call void @llvm.stackrestore.p0(ptr %99) + ret { ptr, ptr, i64, [1 x i64], [1 x i64] } %98 +} + +define void @_mlir_ciface_pyop3_loop(ptr %0, ptr %1, ptr %2, ptr %3, ptr %4, ptr %5, ptr %6) { + %8 = load { ptr, ptr, i64, [1 x i64], [1 x i64] }, ptr %1, align 8 + %9 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %8, 0 + %10 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %8, 1 + %11 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %8, 2 + %12 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %8, 3, 0 + %13 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %8, 4, 0 + %14 = load { ptr, ptr, i64, [1 x i64], [1 x i64] }, ptr %2, align 8 + %15 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %14, 0 + %16 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %14, 1 + %17 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %14, 2 + %18 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %14, 3, 0 + %19 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %14, 4, 0 + %20 = load { ptr, ptr, i64, [1 x i64], [1 x i64] }, ptr %3, align 8 + %21 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %20, 0 + %22 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %20, 1 + %23 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %20, 2 + %24 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %20, 3, 0 + %25 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %20, 4, 0 + %26 = load { ptr, ptr, i64, [1 x i64], [1 x i64] }, ptr %4, align 8 + %27 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %26, 0 + %28 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %26, 1 + %29 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %26, 2 + %30 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %26, 3, 0 + %31 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %26, 4, 0 + %32 = load { ptr, ptr, i64, [1 x i64], [1 x i64] }, ptr %5, align 8 + %33 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %32, 0 + %34 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %32, 1 + %35 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %32, 2 + %36 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %32, 3, 0 + %37 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %32, 4, 0 + %38 = load { ptr, ptr, i64, [1 x i64], [1 x i64] }, ptr %6, align 8 + %39 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %38, 0 + %40 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %38, 1 + %41 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %38, 2 + %42 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %38, 3, 0 + %43 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %38, 4, 0 + %44 = call { ptr, ptr, i64, [1 x i64], [1 x i64] } @pyop3_loop(ptr %9, ptr %10, i64 %11, i64 %12, i64 %13, ptr %15, ptr %16, i64 %17, i64 %18, i64 %19, ptr %21, ptr %22, i64 %23, i64 %24, i64 %25, ptr %27, ptr %28, i64 %29, i64 %30, i64 %31, ptr %33, ptr %34, i64 %35, i64 %36, i64 %37, ptr %39, ptr %40, i64 %41, i64 %42, i64 %43) + store { ptr, ptr, i64, [1 x i64], [1 x i64] } %44, ptr %0, align 8 + ret void +} + +; Function Attrs: nocallback nofree nosync nounwind willreturn +declare ptr @llvm.stacksave.p0() #0 + +; Function Attrs: nocallback nofree nosync nounwind willreturn +declare void @llvm.stackrestore.p0(ptr) #0 + +attributes #0 = { nocallback nofree nosync nounwind willreturn } + +!llvm.module.flags = !{!0} + +!0 = !{i32 2, !"Debug Info Version", i32 3} diff --git a/mlir-testing/run_compiled_mlir.py b/mlir-testing/run_compiled_mlir.py new file mode 100644 index 0000000000..6905b37c82 --- /dev/null +++ b/mlir-testing/run_compiled_mlir.py @@ -0,0 +1,52 @@ +import ctypes +import numpy as np +import cupy as cp + +from ctypes import c_void_p, c_longlong, Structure + +class MemRefDescriptor(Structure): + _fields_ = [ + ("allocated", c_void_p), + ("aligned", c_void_p), + ("offset", c_longlong), + ("shape", c_longlong * 1), + ("stride", c_longlong * 1), + ] + +def numpy_to_memref(arr): + if not arr.flags["C_CONTIGUOUS"]: + arr = np.ascontiguousarray(arr) + + desc = MemRefDescriptor() + desc.allocated = arr.ctypes.data_as(c_void_p) + desc.aligned = desc.allocated + desc.offset = 0 + desc.shape[0] = arr.shape[0] + desc.stride[0] = 1 + + return desc + + +if __name__ == "__main__": + lib = ctypes.CDLL("./liboutput.dylib") + + array_add = lib._mlir_ciface_add + array_add.argtypes = [ + ctypes.POINTER(MemRefDescriptor) + ] * 3 + + size = 8 + a = np.ones(size, dtype=np.float64) + b = np.ones(size, dtype=np.float64) * 2 + c = np.zeros(size, dtype=np.float64) + + a_desc = numpy_to_memref(a) + b_desc = numpy_to_memref(b) + c_desc = numpy_to_memref(c) + + array_add(ctypes.byref(a_desc), ctypes.byref(b_desc), ctypes.byref(c_desc)) + + expected = a + b + np.testing.assert_array_almost_equal(c, expected) + print("Array addition successful!") + print(f"First few elements: {c[:5]}") diff --git a/mlir-testing/sample.py b/mlir-testing/sample.py new file mode 100644 index 0000000000..1e8c018eda --- /dev/null +++ b/mlir-testing/sample.py @@ -0,0 +1,98 @@ +# Basic functionality +# Make MLIR in xDSL that adds two unranked tensor arrays +# Take this MLIR and lower to LLVM in xDSL +# Compile JIT with llvmlite? + +import sys +from xdsl.dialects import arith, func, memref, scf, tensor, linalg +from xdsl.dialects.arith import ConstantOp, AddfOp +from xdsl.dialects.tensor import DimOp, EmptyOp +from xdsl.dialects.builtin import ( + DYNAMIC_INDEX, + ModuleOp, + IndexType, + IntegerAttr, + f64, + TensorType, + ArrayAttr, + AffineMap, + AffineMapAttr, + AffineDimExpr, + UnitAttr +) + + +from xdsl.dialects.linalg.attrs import IteratorTypeAttr +from xdsl.dialects.linalg.ops import YieldOp, GenericOp + +from xdsl.ir import Block, Region +from xdsl.context import Context +from xdsl.printer import Printer + +def build_array_add(n: int) -> ModuleOp: + tensor_type = TensorType(f64, [DYNAMIC_INDEX]) + + identity_1d = AffineMap(num_dims=1, num_symbols=0, results=(AffineDimExpr(0),)) + identity_attr = AffineMapAttr(identity_1d) + + parallel = IteratorTypeAttr.parallel() + + func_block = Block(arg_types=[tensor_type, tensor_type, tensor_type]) + a, b, out = func_block.args + + c0 = ConstantOp(IntegerAttr(0, IndexType())) + func_block.add_op(c0) + + body_block = Block(arg_types=[f64, f64, f64]) + x, y, _z = body_block.args + + add = AddfOp(x, y) + body_block.add_op(add) + + body_block.add_op(YieldOp(add.result)) + + generic = GenericOp( + inputs=[a, b], + outputs=[out], + body=Region([body_block]), + indexing_maps=[identity_attr, identity_attr, identity_attr], + iterator_types=[parallel], + result_types=[tensor_type], + ) + func_block.add_op(generic) + + func_block.add_op(func.ReturnOp()) + + func_region = Region([func_block]) + func_op = func.FuncOp( + "add", + ([tensor_type, tensor_type, tensor_type], []), + func_region, + ) + + func_op.attributes["llvm.emit_c_interface"] = UnitAttr() + + return ModuleOp([func_op]) + +def emit_mlir(module: ModuleOp) -> str: + """Return the MLIR text representation of a module.""" + import io + buf = io.StringIO() + Printer(stream=buf).print_op(module) + return buf.getvalue() + +if __name__ == "__main__": + n = int(sys.argv[1]) if len(sys.argv) > 1 else 8 + + # Register dialects so xDSL can verify the IR + ctx = Context() + ctx.load_dialect(func.Func) + ctx.load_dialect(arith.Arith) + ctx.load_dialect(linalg.Linalg) + ctx.load_dialect(tensor.Tensor) + + module = build_array_add(n) + mlir = emit_mlir(module) + print(mlir) + + diff --git a/mlir-testing/sample_2.py b/mlir-testing/sample_2.py new file mode 100644 index 0000000000..241d30a229 --- /dev/null +++ b/mlir-testing/sample_2.py @@ -0,0 +1,65 @@ +from xdsl.builder import ImplicitBuilder +from xdsl.dialects import arith, func, scf, tensor +from xdsl.dialects.builtin import ( + FloatAttr, + IndexType, + ModuleOp, + TensorType, + f32, + i32, +) +from xdsl.ir import Block, Region + +index = IndexType() + +N = 128 +t_f = TensorType(f32, [N]) # dat_2, dat_3 (float data) +t_i = TensorType(i32, [N]) # idat_0, idat_2 (index data) + +fn_block = Block(arg_types=[t_f, t_f, t_i, t_i]) + +with ImplicitBuilder(fn_block) as (dat_2, dat_3, idat_0, idat_2): + c0 = arith.ConstantOp.from_int_and_width(0, index) + c1 = arith.ConstantOp.from_int_and_width(1, index) + n = arith.ConstantOp.from_int_and_width(N, index) + two = arith.ConstantOp(FloatAttr(2.0, f32)) + + body = Block(arg_types=[index, t_f]) + with ImplicitBuilder(body) as (i2, acc): + # k = idat_2[i_2] + k = tensor.ExtractOp(idat_2, [i2], i32) + k_idx = arith.IndexCastOp(k.result, index) + + # j = idat_0[k] + j = tensor.ExtractOp(idat_0, [k_idx.result], i32) + j_idx = arith.IndexCastOp(j.result, index) + + # v = dat_3[j] + v = tensor.ExtractOp(dat_3, [j_idx.result], f32) + + # r = 2.0 * v + r = arith.MulfOp(two.result, v.result) + + # dat_2[j] = r (value semantics -> produces new tensor) + new = tensor.InsertOp(r.result, acc, [j_idx.result]) + + scf.YieldOp(new.result) + + loop = scf.ForOp( + lb=c0.result, + ub=n.result, + step=c1.result, + iter_args=[dat_2], + body=Region(body), + ) + + func.ReturnOp(loop.results[0]) + +fn = func.FuncOp( + "indirect_scale", + ((t_f, t_f, t_i, t_i), (t_f,)), + Region(fn_block), +) + +module = ModuleOp([fn]) +print(module) From 134d0cb8258144f2b954424481e4fa39aaf72b1a Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Tue, 25 Aug 2026 15:57:13 +0100 Subject: [PATCH 12/30] bug fixed, delayed dtype integration Still need to integrate dtype support for all op3 expressions. Removed abstractmethod while testing. --- pyop3/expr/base.py | 2 +- pyop3/lower/codegen.py | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/pyop3/expr/base.py b/pyop3/expr/base.py index f64c4f7fe7..894e18b512 100644 --- a/pyop3/expr/base.py +++ b/pyop3/expr/base.py @@ -28,7 +28,7 @@ def local_min(self) -> numbers.Number: raise NotImplementedError @property - @abc.abstractmethod + # @abc.abstractmethod def dtype(self) -> np.dtype: pass diff --git a/pyop3/lower/codegen.py b/pyop3/lower/codegen.py index fc2ed036ab..e19db131a3 100644 --- a/pyop3/lower/codegen.py +++ b/pyop3/lower/codegen.py @@ -150,15 +150,15 @@ def _(loop): @_collect_temporary_shapes.register(AbstractAssignment) @_collect_temporary_shapes.register(NullInstruction) @_collect_temporary_shapes.register(Exscan) -def _(assignment): +def _(assignment: AbstractAssignment, /) -> idict: return idict() -@_collect_temporary_shapes.register(StandaloneCalledFunction) -def _(call): +@_collect_temporary_shapes.register +def _(call: StandaloneCalledFunction): import loopy as lp # TODO: Remove once StandaloneCalledFunction/similar integrated with MLIR return idict( { - (arg.buffer.name, arg.buffer.nest_indices): lp_arg.shape + arg.buffer: lp_arg.shape for lp_arg, arg in zip( call.function.code.default_entrypoint.args, call.arguments, strict=True ) @@ -328,7 +328,7 @@ def compile_array_assignment( if component.local_size != 1: iname = codegen_context.unique_name("i") ext = codegen_context.register_extent( - component.local_size, + component.size, iname_replace_maps[-1], loop_indices ) From 8f18475d0a5b2afbc40c75ef9abdbf8e90252f93 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Tue, 25 Aug 2026 16:12:56 +0100 Subject: [PATCH 13/30] removing debugging python/text files --- assign_local_size.mlir | 35 --------- assign_local_size.txt | 29 ------- assign_size.txt | 42 ---------- integral_loop.txt | 68 ---------------- loopy_assign.txt | 38 --------- mlir_assign.mlir | 172 ----------------------------------------- run_compiled_mlir.py | 52 ------------- sample.py | 98 ----------------------- sample_2.py | 65 ---------------- 9 files changed, 599 deletions(-) delete mode 100644 assign_local_size.mlir delete mode 100644 assign_local_size.txt delete mode 100644 assign_size.txt delete mode 100644 integral_loop.txt delete mode 100644 loopy_assign.txt delete mode 100644 mlir_assign.mlir delete mode 100644 run_compiled_mlir.py delete mode 100644 sample.py delete mode 100644 sample_2.py diff --git a/assign_local_size.mlir b/assign_local_size.mlir deleted file mode 100644 index 0af0ddd9cd..0000000000 --- a/assign_local_size.mlir +++ /dev/null @@ -1,35 +0,0 @@ -builtin.module { - func.func @pyop3_loop( - %dat_0: tensor, - %dat_1: tensor, - %idat_0: tensor, - %idat_1: tensor, - %idat_2: tensor, - %idat_3: tensor - ) -> tensor { - - %c0 = arith.constant 0 : index // iter var - %c1 = arith.constant 1 : index // iter var - %c17 = arith.constant 17 : index // - %c15 = arith.constant 15 : index - %c32 = arith.constant 32 : index - %c2f = arith.constant 2.0 : f64 - - scf.for %i_0 = %c0 to %c17 step %c1 iter_args() -> () {} - - %f2 = scf.for %i_2 = %c0 to %c15 step %c1 iter_args(%dat_it = %dat_0) -> (tensor) { - %e1 = tensor.extract %idat_2[%i_2] : tensor - %ii_2 = arith.index_cast %e1 : i32 to index - %e2 = tensor.extract %idat_0[%ii_2] : tensor - %iii_2 = arith.index_cast %e2 : i32 to index - %v1 = tensor.extract %dat_1[%iii_2] : tensor - %v2 = arith.mulf %c2f, %v1 : f64 - %res = tensor.insert %v2 into %dat_it[%iii_2] : tensor - scf.yield %res : tensor - } - - scf.for %i_3 = %c0 to %c32 step %c1 iter_args() -> () {} - - func.return %f2 : tensor - } -} diff --git a/assign_local_size.txt b/assign_local_size.txt deleted file mode 100644 index d5e7205a71..0000000000 --- a/assign_local_size.txt +++ /dev/null @@ -1,29 +0,0 @@ -******************************************************************************** -#include -#include -#include -#include - -void pyop3_loop(double *__restrict__ dat_0, double const *__restrict__ dat_1, int32_t const *__restrict__ idat_0, int32_t const *__restrict__ idat_1, int32_t const *__restrict__ idat_2, int32_t const *__restrict__ idat_3) -{ - for (int32_t i_0 = 0; i_0 <= 17; ++i_0) - { - } - for (int32_t i_2 = 0; i_2 <= 15; ++i_2) - dat_0[idat_0[idat_2[i_2]]] = 2.0 * dat_1[idat_0[idat_2[i_2]]]; - for (int32_t i_3 = 0; i_3 <= 32; ++i_3) - { - } - -} -******************************************************************************** -dat_0 (16) : [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] -dat_1 (16) : [10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10.] -idat_0 (67) : [ 0 0 0 0 0 1 2 3 3 3 3 4 4 4 4 5 5 5 5 6 6 6 6 7 - 7 7 7 8 8 8 8 9 9 9 9 9 9 10 10 10 10 11 11 11 11 12 12 12 - 12 13 13 13 13 13 13 13 13 14 14 14 14 15 15 15 15 15 15] -idat_1 (18) : [ 0 7 11 15 19 23 27 31 33 37 41 45 49 51 53 57 61 63] -idat_2 (16) : [ 4 5 6 10 14 18 22 26 30 36 40 44 48 56 60 66] -idat_3 (33) : [ 1 2 3 8 9 12 13 16 17 20 21 24 25 28 29 32 34 35 38 39 42 43 46 47 - 50 52 54 55 58 59 62 64 65] -******************************************************************************** diff --git a/assign_size.txt b/assign_size.txt deleted file mode 100644 index 21b8e370be..0000000000 --- a/assign_size.txt +++ /dev/null @@ -1,42 +0,0 @@ -******************************************************************************** -#include -#include -#include -#include - -void pyop3_loop(int64_t const *__restrict__ dat_0, int64_t const *__restrict__ dat_1, double *__restrict__ dat_2, double const *__restrict__ dat_3, int64_t const *__restrict__ dat_4, int64_t const *__restrict__ dat_5, int64_t const *__restrict__ dat_6, int64_t const *__restrict__ dat_7, int32_t const *__restrict__ idat_0, int32_t const *__restrict__ idat_1, int32_t const *__restrict__ idat_2, int32_t const *__restrict__ idat_3) -{ - int32_t p_0; - int32_t p_1; - int32_t p_2; - - p_0 = (int32_t) (dat_0[0] + dat_1[0]); - for (int32_t i_0 = 0; i_0 <= -1 + p_0; ++i_0) - { - } - p_1 = (int32_t) (dat_4[0] + dat_5[0]); - for (int32_t i_2 = 0; i_2 <= -1 + p_1; ++i_2) - dat_2[idat_0[idat_2[i_2]]] = 2.0 * dat_3[idat_0[idat_2[i_2]]]; - p_2 = (int32_t) (dat_6[0] + dat_7[0]); - for (int32_t i_3 = 0; i_3 <= -1 + p_2; ++i_3) - { - } - -} -******************************************************************************** -dat_0 (1) : [18] -dat_1 (1) : [0] -dat_2 (16) : [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] -dat_3 (16) : [10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10.] -dat_4 (1) : [16] -dat_5 (1) : [0] -dat_6 (1) : [33] -dat_7 (1) : [0] -idat_0 (67) : [ 0 0 0 0 0 1 2 3 3 3 3 4 4 4 4 5 5 5 5 6 6 6 6 7 - 7 7 7 8 8 8 8 9 9 9 9 9 9 10 10 10 10 11 11 11 11 12 12 12 - 12 13 13 13 13 13 13 13 13 14 14 14 14 15 15 15 15 15 15] -idat_1 (18) : [ 0 7 11 15 19 23 27 31 33 37 41 45 49 51 53 57 61 63] -idat_2 (16) : [ 4 5 6 10 14 18 22 26 30 36 40 44 48 56 60 66] -idat_3 (33) : [ 1 2 3 8 9 12 13 16 17 20 21 24 25 28 29 32 34 35 38 39 42 43 46 47 - 50 52 54 55 58 59 62 64 65] -******************************************************************************** diff --git a/integral_loop.txt b/integral_loop.txt deleted file mode 100644 index feac230fa3..0000000000 --- a/integral_loop.txt +++ /dev/null @@ -1,68 +0,0 @@ -******************************************************************************** -#include -#include -#include -#include -#include - -static void form_cell_integral(double *__restrict__ A, double const *__restrict__ coords_0); -static void form_cell_integral(double *__restrict__ A, double const *__restrict__ coords_0) -{ - double t0; - double t1; - double t2; - double t3[3] = { 0.33333333333333337, 0.33333333333333326, 0.33333333333333326 }; - - t0 = -1.0 * coords_0[0]; - t1 = -1.0 * coords_0[1]; - t2 = 0.5 * fabs((t0 + coords_0[2]) * (t1 + coords_0[5]) + -1.0 * (t0 + coords_0[4]) * (t1 + coords_0[3])); - for (int32_t j = 0; j <= 2; ++j) - A[j] = A[j] + t3[j] * t2; - -} - -void pyop3_loop(double const *__restrict__ dat_0, double *__restrict__ dat_1, int32_t const *__restrict__ idat_0, int32_t const *__restrict__ idat_1) -{ - int32_t j_0; - int32_t j_1; - int32_t j_2; - double t_0[3l]; - double t_1[6l]; - - for (int32_t i_0 = 0; i_0 <= 17; ++i_0) - { - for (int32_t i_1 = 0; i_1 <= 2; ++i_1) - { - j_0 = 0 + 1 * (i_1 + 0); - t_0[j_0] = (double) (0.0); - } - for (int32_t i_2 = 0; i_2 <= 2; ++i_2) - for (int32_t i_3 = 0; i_3 <= 1; ++i_3) - { - j_1 = 0 + 1 * (i_2 * 2 + 0 + i_3); - t_1[j_1] = dat_0[idat_0[3 * i_0 + i_2] + i_3]; - } - form_cell_integral(&(t_0[0]), &(t_1[0])); - for (int32_t i_6 = 0; i_6 <= 2; ++i_6) - { - j_2 = 0 + 1 * (i_6 + 0); - dat_1[idat_1[3 * i_0 + i_6]] = dat_1[idat_1[3 * i_0 + i_6]] + t_0[j_2]; - } - } - -} -******************************************************************************** -dat_0 (32) : [0.33333333 0. 0. 0.33333333 0. 0. - 0.33333333 0.33333333 0. 0.66666667 0.66666667 0. - 0.33333333 0.66666667 0.66666667 0.33333333 0. 1. - 1. 0. 0.33333333 1. 0.66666667 0.66666667 - 1. 0.33333333 0.66666667 1. 1. 0.66666667 - 1. 1. ] -dat_1 (16) : [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] -idat_0 (54) : [ 0 2 4 0 2 6 2 6 8 0 6 10 6 8 12 6 10 14 8 12 16 6 12 14 - 10 14 18 12 16 20 12 14 22 14 18 24 12 20 22 14 22 24 20 22 26 22 24 28 - 22 26 28 26 28 30] -idat_1 (54) : [ 0 1 2 0 1 3 1 3 4 0 3 5 3 4 6 3 5 7 4 6 8 3 6 7 - 5 7 9 6 8 10 6 7 11 7 9 12 6 10 11 7 11 12 10 11 13 11 12 14 - 11 13 14 13 14 15] -******************************************************************************** diff --git a/loopy_assign.txt b/loopy_assign.txt deleted file mode 100644 index d14de3083d..0000000000 --- a/loopy_assign.txt +++ /dev/null @@ -1,38 +0,0 @@ ---------------------------------------------------------------------------- -KERNEL: pyop3_loop ---------------------------------------------------------------------------- -ARGUMENTS: -dat_0: ArrayArg, type: np:dtype('float64'), shape: unknown in/out aspace: global -dat_1: ArrayArg, type: np:dtype('float64'), shape: unknown in aspace: global -idat_0: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global -idat_1: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global -idat_2: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global -idat_3: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global ---------------------------------------------------------------------------- -DOMAINS: -{ [i_0] : 0 <= i_0 <= 17 } -{ [i_1] : 1 = 0 } -{ [i_2] : 0 <= i_2 <= 15 } -{ [i_3] : 0 <= i_3 <= 32 } -{ [i_4] : 1 = 0 } ---------------------------------------------------------------------------- -INAME TAGS: -i_0: None -i_1: None -i_2: None -i_3: None -i_4: None ---------------------------------------------------------------------------- -INSTRUCTIONS: - for i_1, i_0 -↱ dat_0[0 + 1*(idat_0[0 + 1*idat_1[0 + 1*i_0]] + i_1)] = 2*dat_1[0 + 1*(idat_0[0 + 1*idat_1[0 + 1*i_0]] + i_1)] {id=insn_0} -│ end i_1, i_0 -│ for i_2 -└↱ dat_0[0 + 1*(idat_0[0 + 1*idat_2[0 + 1*i_2]] + 0)] = 2*dat_1[0 + 1*(idat_0[0 + 1*idat_2[0 + 1*i_2]] + 0) {id=insn_1} - │ end i_2 - │ for i_4, i_3 -↱└ dat_0[0 + 1*(idat_0[0 + 1*idat_3[0 + 1*i_3]] + i_4) = 2*dat_1[0 + 1*(idat_0[0 + 1*idat_3[0 + 1*i_3]] + i_4) {id=insn_2} -│ end i_4, i_3 -└ CODE(idat_2, dat_0, idat_0, idat_1, dat_1, idat_3|) {id=insn} - ---------------------------------------------------------------------------- diff --git a/mlir_assign.mlir b/mlir_assign.mlir deleted file mode 100644 index 4eab0e02b9..0000000000 --- a/mlir_assign.mlir +++ /dev/null @@ -1,172 +0,0 @@ -builtin.module { - func.func @pyop3_loop(%0: tensor, %1: tensor, %2: tensor, %3: tensor, %4: tensor, %5: tensor) { - %6 = arith.constant 0 : i32 - %7 = arith.index_cast %6 : i32 to index - %8 = arith.constant 18 : i32 - %9 = arith.index_cast %8 : i32 to index - %10 = arith.constant 1 : i32 - %11 = arith.index_cast %10 : i32 to index - %12 = scf.for %13 = %7 to %9 step %11 iter_args(%14 = %0) -> (tensor) { - %15 = arith.constant 0 : i32 - %16 = arith.index_cast %15 : i32 to index - %17 = arith.constant 0 : i32 - %18 = arith.index_cast %17 : i32 to index - %19 = arith.constant 1 : i32 - %20 = arith.index_cast %19 : i32 to index - %21 = scf.for %22 = %16 to %18 step %20 iter_args(%23 = %14) -> (tensor) { - %24 = arith.constant 2 : i32 - %25 = arith.constant 0 : i32 - %26 = arith.constant 1 : i32 - %27 = arith.constant 0 : i32 - %28 = arith.constant 1 : i32 - %29 = arith.constant 0 : i32 - %30 = arith.constant 1 : i32 - %31 = arith.muli %30, %13 : i32 - %32 = arith.addi %29, %31 : i32 - %33 = arith.index_cast %32 : i32 to index - %34 = tensor.extract %2[%33] : tensor - %35 = arith.muli %28, %34 : i32 - %36 = arith.addi %27, %35 : i32 - %37 = arith.index_cast %36 : i32 to index - %38 = tensor.extract %1[%37] : tensor - %39 = arith.addi %38, %22 : i32 - %40 = arith.muli %26, %39 : i32 - %41 = arith.addi %25, %40 : i32 - %42 = arith.index_cast %41 : i32 to index - %43 = tensor.extract %3[%42] : tensor - %44 = arith.muli %24, %43 : i32 - %45 = arith.constant 0 : i32 - %46 = arith.constant 1 : i32 - %47 = arith.constant 0 : i32 - %48 = arith.constant 1 : i32 - %49 = arith.constant 0 : i32 - %50 = arith.constant 1 : i32 - %51 = arith.muli %50, %13 : i32 - %52 = arith.addi %49, %51 : i32 - %53 = arith.index_cast %52 : i32 to index - %54 = tensor.extract %2[%53] : tensor - %55 = arith.muli %48, %54 : i32 - %56 = arith.addi %47, %55 : i32 - %57 = arith.index_cast %56 : i32 to index - %58 = tensor.extract %1[%57] : tensor - %59 = arith.addi %58, %22 : i32 - %60 = arith.muli %46, %59 : i32 - %61 = arith.addi %45, %60 : i32 - %62 = arith.index_cast %61 : i32 to index - %63 = tensor.insert %44 into %23[%62] : tensor - scf.yield %63 : tensor - } - scf.yield %21 : tensor - } - %64 = arith.constant 0 : i32 - %65 = arith.index_cast %64 : i32 to index - %66 = arith.constant 16 : i32 - %67 = arith.index_cast %66 : i32 to index - %68 = arith.constant 1 : i32 - %69 = arith.index_cast %68 : i32 to index - %70 = scf.for %71 = %65 to %67 step %69 iter_args(%72 = %12) -> (tensor) { - %73 = arith.constant 2 : i32 - %74 = arith.constant 0 : i32 - %75 = arith.constant 1 : i32 - %76 = arith.constant 0 : i32 - %77 = arith.constant 1 : i32 - %78 = arith.constant 0 : i32 - %79 = arith.constant 1 : i32 - %80 = arith.muli %79, %71 : i32 - %81 = arith.addi %78, %80 : i32 - %82 = arith.index_cast %81 : i32 to index - %83 = tensor.extract %4[%82] : tensor - %84 = arith.muli %77, %83 : i32 - %85 = arith.addi %76, %84 : i32 - %86 = arith.index_cast %85 : i32 to index - %87 = tensor.extract %1[%86] : tensor - %88 = arith.constant 0 : i32 - %89 = arith.addi %87, %88 : i32 - %90 = arith.muli %75, %89 : i32 - %91 = arith.addi %74, %90 : i32 - %92 = arith.index_cast %91 : i32 to index - %93 = tensor.extract %3[%92] : tensor - %94 = arith.muli %73, %93 : i32 - %95 = arith.constant 0 : i32 - %96 = arith.constant 1 : i32 - %97 = arith.constant 0 : i32 - %98 = arith.constant 1 : i32 - %99 = arith.constant 0 : i32 - %100 = arith.constant 1 : i32 - %101 = arith.muli %100, %71 : i32 - %102 = arith.addi %99, %101 : i32 - %103 = arith.index_cast %102 : i32 to index - %104 = tensor.extract %4[%103] : tensor - %105 = arith.muli %98, %104 : i32 - %106 = arith.addi %97, %105 : i32 - %107 = arith.index_cast %106 : i32 to index - %108 = tensor.extract %1[%107] : tensor - %109 = arith.constant 0 : i32 - %110 = arith.addi %108, %109 : i32 - %111 = arith.muli %96, %110 : i32 - %112 = arith.addi %95, %111 : i32 - %113 = arith.index_cast %112 : i32 to index - %114 = tensor.insert %94 into %72[%113] : tensor - scf.yield %114 : tensor - } - %115 = arith.constant 0 : i32 - %116 = arith.index_cast %115 : i32 to index - %117 = arith.constant 33 : i32 - %118 = arith.index_cast %117 : i32 to index - %119 = arith.constant 1 : i32 - %120 = arith.index_cast %119 : i32 to index - %121 = scf.for %122 = %116 to %118 step %120 iter_args(%123 = %70) -> (tensor) { - %124 = arith.constant 0 : i32 - %125 = arith.index_cast %124 : i32 to index - %126 = arith.constant 0 : i32 - %127 = arith.index_cast %126 : i32 to index - %128 = arith.constant 1 : i32 - %129 = arith.index_cast %128 : i32 to index - %130 = scf.for %131 = %125 to %127 step %129 iter_args(%132 = %123) -> (tensor) { - %133 = arith.constant 2 : i32 - %134 = arith.constant 0 : i32 - %135 = arith.constant 1 : i32 - %136 = arith.constant 0 : i32 - %137 = arith.constant 1 : i32 - %138 = arith.constant 0 : i32 - %139 = arith.constant 1 : i32 - %140 = arith.muli %139, %122 : i32 - %141 = arith.addi %138, %140 : i32 - %142 = arith.index_cast %141 : i32 to index - %143 = tensor.extract %5[%142] : tensor - %144 = arith.muli %137, %143 : i32 - %145 = arith.addi %136, %144 : i32 - %146 = arith.index_cast %145 : i32 to index - %147 = tensor.extract %1[%146] : tensor - %148 = arith.addi %147, %131 : i32 - %149 = arith.muli %135, %148 : i32 - %150 = arith.addi %134, %149 : i32 - %151 = arith.index_cast %150 : i32 to index - %152 = tensor.extract %3[%151] : tensor - %153 = arith.muli %133, %152 : i32 - %154 = arith.constant 0 : i32 - %155 = arith.constant 1 : i32 - %156 = arith.constant 0 : i32 - %157 = arith.constant 1 : i32 - %158 = arith.constant 0 : i32 - %159 = arith.constant 1 : i32 - %160 = arith.muli %159, %122 : i32 - %161 = arith.addi %158, %160 : i32 - %162 = arith.index_cast %161 : i32 to index - %163 = tensor.extract %5[%162] : tensor - %164 = arith.muli %157, %163 : i32 - %165 = arith.addi %156, %164 : i32 - %166 = arith.index_cast %165 : i32 to index - %167 = tensor.extract %1[%166] : tensor - %168 = arith.addi %167, %131 : i32 - %169 = arith.muli %155, %168 : i32 - %170 = arith.addi %154, %169 : i32 - %171 = arith.index_cast %170 : i32 to index - %172 = tensor.insert %153 into %132[%171] : tensor - scf.yield %172 : tensor - } - scf.yield %130 : tensor - } - func.return - } -} diff --git a/run_compiled_mlir.py b/run_compiled_mlir.py deleted file mode 100644 index 6905b37c82..0000000000 --- a/run_compiled_mlir.py +++ /dev/null @@ -1,52 +0,0 @@ -import ctypes -import numpy as np -import cupy as cp - -from ctypes import c_void_p, c_longlong, Structure - -class MemRefDescriptor(Structure): - _fields_ = [ - ("allocated", c_void_p), - ("aligned", c_void_p), - ("offset", c_longlong), - ("shape", c_longlong * 1), - ("stride", c_longlong * 1), - ] - -def numpy_to_memref(arr): - if not arr.flags["C_CONTIGUOUS"]: - arr = np.ascontiguousarray(arr) - - desc = MemRefDescriptor() - desc.allocated = arr.ctypes.data_as(c_void_p) - desc.aligned = desc.allocated - desc.offset = 0 - desc.shape[0] = arr.shape[0] - desc.stride[0] = 1 - - return desc - - -if __name__ == "__main__": - lib = ctypes.CDLL("./liboutput.dylib") - - array_add = lib._mlir_ciface_add - array_add.argtypes = [ - ctypes.POINTER(MemRefDescriptor) - ] * 3 - - size = 8 - a = np.ones(size, dtype=np.float64) - b = np.ones(size, dtype=np.float64) * 2 - c = np.zeros(size, dtype=np.float64) - - a_desc = numpy_to_memref(a) - b_desc = numpy_to_memref(b) - c_desc = numpy_to_memref(c) - - array_add(ctypes.byref(a_desc), ctypes.byref(b_desc), ctypes.byref(c_desc)) - - expected = a + b - np.testing.assert_array_almost_equal(c, expected) - print("Array addition successful!") - print(f"First few elements: {c[:5]}") diff --git a/sample.py b/sample.py deleted file mode 100644 index 1e8c018eda..0000000000 --- a/sample.py +++ /dev/null @@ -1,98 +0,0 @@ -# Basic functionality -# Make MLIR in xDSL that adds two unranked tensor arrays -# Take this MLIR and lower to LLVM in xDSL -# Compile JIT with llvmlite? - -import sys -from xdsl.dialects import arith, func, memref, scf, tensor, linalg -from xdsl.dialects.arith import ConstantOp, AddfOp -from xdsl.dialects.tensor import DimOp, EmptyOp -from xdsl.dialects.builtin import ( - DYNAMIC_INDEX, - ModuleOp, - IndexType, - IntegerAttr, - f64, - TensorType, - ArrayAttr, - AffineMap, - AffineMapAttr, - AffineDimExpr, - UnitAttr -) - - -from xdsl.dialects.linalg.attrs import IteratorTypeAttr -from xdsl.dialects.linalg.ops import YieldOp, GenericOp - -from xdsl.ir import Block, Region -from xdsl.context import Context -from xdsl.printer import Printer - -def build_array_add(n: int) -> ModuleOp: - tensor_type = TensorType(f64, [DYNAMIC_INDEX]) - - identity_1d = AffineMap(num_dims=1, num_symbols=0, results=(AffineDimExpr(0),)) - identity_attr = AffineMapAttr(identity_1d) - - parallel = IteratorTypeAttr.parallel() - - func_block = Block(arg_types=[tensor_type, tensor_type, tensor_type]) - a, b, out = func_block.args - - c0 = ConstantOp(IntegerAttr(0, IndexType())) - func_block.add_op(c0) - - body_block = Block(arg_types=[f64, f64, f64]) - x, y, _z = body_block.args - - add = AddfOp(x, y) - body_block.add_op(add) - - body_block.add_op(YieldOp(add.result)) - - generic = GenericOp( - inputs=[a, b], - outputs=[out], - body=Region([body_block]), - indexing_maps=[identity_attr, identity_attr, identity_attr], - iterator_types=[parallel], - result_types=[tensor_type], - ) - func_block.add_op(generic) - - func_block.add_op(func.ReturnOp()) - - func_region = Region([func_block]) - func_op = func.FuncOp( - "add", - ([tensor_type, tensor_type, tensor_type], []), - func_region, - ) - - func_op.attributes["llvm.emit_c_interface"] = UnitAttr() - - return ModuleOp([func_op]) - -def emit_mlir(module: ModuleOp) -> str: - """Return the MLIR text representation of a module.""" - import io - buf = io.StringIO() - Printer(stream=buf).print_op(module) - return buf.getvalue() - -if __name__ == "__main__": - n = int(sys.argv[1]) if len(sys.argv) > 1 else 8 - - # Register dialects so xDSL can verify the IR - ctx = Context() - ctx.load_dialect(func.Func) - ctx.load_dialect(arith.Arith) - ctx.load_dialect(linalg.Linalg) - ctx.load_dialect(tensor.Tensor) - - module = build_array_add(n) - mlir = emit_mlir(module) - print(mlir) - - diff --git a/sample_2.py b/sample_2.py deleted file mode 100644 index 241d30a229..0000000000 --- a/sample_2.py +++ /dev/null @@ -1,65 +0,0 @@ -from xdsl.builder import ImplicitBuilder -from xdsl.dialects import arith, func, scf, tensor -from xdsl.dialects.builtin import ( - FloatAttr, - IndexType, - ModuleOp, - TensorType, - f32, - i32, -) -from xdsl.ir import Block, Region - -index = IndexType() - -N = 128 -t_f = TensorType(f32, [N]) # dat_2, dat_3 (float data) -t_i = TensorType(i32, [N]) # idat_0, idat_2 (index data) - -fn_block = Block(arg_types=[t_f, t_f, t_i, t_i]) - -with ImplicitBuilder(fn_block) as (dat_2, dat_3, idat_0, idat_2): - c0 = arith.ConstantOp.from_int_and_width(0, index) - c1 = arith.ConstantOp.from_int_and_width(1, index) - n = arith.ConstantOp.from_int_and_width(N, index) - two = arith.ConstantOp(FloatAttr(2.0, f32)) - - body = Block(arg_types=[index, t_f]) - with ImplicitBuilder(body) as (i2, acc): - # k = idat_2[i_2] - k = tensor.ExtractOp(idat_2, [i2], i32) - k_idx = arith.IndexCastOp(k.result, index) - - # j = idat_0[k] - j = tensor.ExtractOp(idat_0, [k_idx.result], i32) - j_idx = arith.IndexCastOp(j.result, index) - - # v = dat_3[j] - v = tensor.ExtractOp(dat_3, [j_idx.result], f32) - - # r = 2.0 * v - r = arith.MulfOp(two.result, v.result) - - # dat_2[j] = r (value semantics -> produces new tensor) - new = tensor.InsertOp(r.result, acc, [j_idx.result]) - - scf.YieldOp(new.result) - - loop = scf.ForOp( - lb=c0.result, - ub=n.result, - step=c1.result, - iter_args=[dat_2], - body=Region(body), - ) - - func.ReturnOp(loop.results[0]) - -fn = func.FuncOp( - "indirect_scale", - ((t_f, t_f, t_i, t_i), (t_f,)), - Region(fn_block), -) - -module = ModuleOp([fn]) -print(module) From c569f303d52577fca6b9c5c4a4365b6e9a8cc774 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Wed, 26 Aug 2026 12:32:44 +0100 Subject: [PATCH 14/30] cleaning up refactor, to add docstrings --- pyop3/lower/codegen.py | 248 +--------------------------------- pyop3/lower/context.py | 296 ++++++++++++++++++++++++++++++++++++++++- pyop3/lower/loopy.py | 6 +- 3 files changed, 299 insertions(+), 251 deletions(-) diff --git a/pyop3/lower/codegen.py b/pyop3/lower/codegen.py index e19db131a3..06d913cb67 100644 --- a/pyop3/lower/codegen.py +++ b/pyop3/lower/codegen.py @@ -45,6 +45,7 @@ ) from pyop3.lower.loopy import LoopyCodegenContext +from pyop3.lower.context import _collect_temporary_shapes, _compile # TODO: import other way around? from pyop3.lower.transform import ( @@ -126,250 +127,3 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed return translation_unit, kernel_name_to_global_buffer_info, global_buffer_intents -@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): - import loopy as lp # TODO: Remove once StandaloneCalledFunction/similar integrated with MLIR - 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: 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: - 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 = codegen_context.register_extent( - component.size, - iname_map, - loop_indices - ) - 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(StandaloneCalledFunction) -def _(call, loop_indices, codegen_context): - codegen_context.compile_standalone_function(call, loop_indices) - -@_compile.register(NonEmptyArrayAssignment) -def parse_assignment(assignment: NonEmptyArrayAssignment, loop_indices, context: CodegenContext): - if any(isinstance(arg, pyop3.expr.MatPetscMatBufferExpression) for arg in assignment.arguments): - context.compile_petsc_mat(assignment, loop_indices) - else: - compile_array_assignment( - assignment, - loop_indices, - context, - assignment.axis_trees, - ) - -# NOTE: Move this? Weird to have this one here and rest in context classes. -# Probably move to context.py if I can remove pym references. -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("Refactor needed") - - codegen_context.add_leaf_assignment( - assignment, - 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: pym.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, - codegen_context, - axis_trees, - iname_replace_maps=new_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_maps, - axis_tree=None, - paths=new_paths - ) - else: - codegen_context.add_leaf_assignment( - assignment, - new_paths, - new_maps, - loop_indices - ) - -@_compile.register(Exscan) -def _(exscan, loop_indices, codegen_context): - codegen_context.compile_exscan(exscan, loop_indices) diff --git a/pyop3/lower/context.py b/pyop3/lower/context.py index f05302f915..2c93b247e7 100644 --- a/pyop3/lower/context.py +++ b/pyop3/lower/context.py @@ -1,12 +1,41 @@ from abc import ABC, abstractmethod from typing import Any, List, Dict, Tuple import numbers +import functools -from pyop3 import utils +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): """ Abstract base class for code generation contexts. @@ -61,6 +90,15 @@ def _add_instruction(self, insn): # {{{ 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 @@ -120,18 +158,182 @@ def register_extent(self, obj: Any, inames, loop_indices): @abstractmethod def compile_standalone_function(self, call, loop_indices): + """ + Compiling standalone functions i.e. LACallable for target IR representations + """ pass @abstractmethod def compile_petsc_mat(self, assignment, loop_indices): + """ + """ pass @abstractmethod def compile_exscan(self, call, loop_indices): + """ + """ pass # }}} + + # {{{ general implementations + + + def compile_array_assignment( + self, + assignment, + loop_indices, + 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("Refactor needed") + + self.add_leaf_assignment( + assignment, + 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 = self.unique_name("i") + ext = self.register_extent( + component.size, + iname_replace_maps[-1], + loop_indices + ) + self.add_domain(iname, ext) + new_maps = iname_replace_maps.copy() + new_maps[-1] = iname_replace_maps[-1] | {axis.label: self.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 self.within_inames(within_inames): + if axis_tree.node_map[new_paths[-1]]: + self.compile_array_assignment( + assignment, + loop_indices, + axis_trees, + iname_replace_maps=new_maps, + axis_tree=axis_tree, + paths=new_paths + ) + elif axis_trees: + self.compile_array_assignment( + assignment, + loop_indices, + axis_trees, + iname_replace_maps=new_maps, + axis_tree=None, + paths=new_paths + ) + else: + self.add_leaf_assignment( + assignment, + new_paths, + new_maps, + loop_indices + ) + + def parse_loop_properly_this_time( + self, + loop, + axis_tree, + loop_indices, + *, + 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, + self, + ) + 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 = self.unique_name("i") + domain_var = self.register_extent( + component.size, + iname_map, + loop_indices + ) + self.add_domain(iname, domain_var) + iname_replace_map_ = iname_map | {axis.label: self.var(iname)} + within_inames = frozenset({iname}) + else: + iname_replace_map_ = iname_map | {axis.label: 0} + within_inames = set() + + with self.within_inames(within_inames): + if subaxis := axis_tree.node_map[path_]: + self.parse_loop_properly_this_time( + loop, + axis_tree, + loop_indices, + 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, + self, + ) + # }}} + def add_subkernel(self, subkernel): self._subkernels.append(subkernel) @@ -144,3 +346,95 @@ def __str__(self) -> str: ctx += f"Arguments: {str(self.arguments)}\n\n" ctx += f"Subkernels: {str(self.subkernels)}\n\n" return ctx + + +@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: + codegen_context.parse_loop_properly_this_time( + loop, + loop.index.iterset, + loop_indices, + ) + +@_compile.register(StandaloneCalledFunction) +def _(call, loop_indices, codegen_context): + codegen_context.compile_standalone_function(call, loop_indices) + +@_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): + codegen_context.compile_petsc_mat(assignment, loop_indices) + else: + codegen_context.compile_array_assignment( + assignment, + loop_indices, + assignment.axis_trees, + ) + +@_compile.register(Exscan) +def _(exscan, loop_indices, codegen_context): + codegen_context.compile_exscan(exscan, loop_indices) + + +# 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): + import loopy as lp # TODO: Remove once StandaloneCalledFunction/similar integrated with MLIR + 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/pyop3/lower/loopy.py b/pyop3/lower/loopy.py index cf6669ca12..37ec6ffa9c 100644 --- a/pyop3/lower/loopy.py +++ b/pyop3/lower/loopy.py @@ -86,6 +86,9 @@ def subkernels(self) -> tuple: 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) @@ -265,9 +268,6 @@ 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 add_leaf_assignment(self, assignment, paths, iname_maps, loop_indices): - pass def lower_buffer_access( self, From 28ea6c97ee4ff7657e29357e65dd76d2ca69fd97 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Wed, 26 Aug 2026 14:10:19 +0100 Subject: [PATCH 15/30] removing testing files from git --- mlir-testing/assign_local_size.mlir | 35 ------ mlir-testing/assign_local_size.txt | 29 ----- mlir-testing/assign_mlir_demo.py | 17 --- mlir-testing/assign_size.txt | 42 ------- mlir-testing/fixed_output.mlir | 159 ------------------------- mlir-testing/integral_loop.txt | 68 ----------- mlir-testing/loopy_assign.txt | 38 ------ mlir-testing/mlir_assign.mlir | 172 ---------------------------- mlir-testing/mlir_assign_fixed.mlir | 152 ------------------------ mlir-testing/mlir_assign_opt.mlir | 20 ---- mlir-testing/output.ll | 153 ------------------------- mlir-testing/run_compiled_mlir.py | 52 --------- mlir-testing/sample.py | 98 ---------------- mlir-testing/sample_2.py | 65 ----------- 14 files changed, 1100 deletions(-) delete mode 100644 mlir-testing/assign_local_size.mlir delete mode 100644 mlir-testing/assign_local_size.txt delete mode 100644 mlir-testing/assign_mlir_demo.py delete mode 100644 mlir-testing/assign_size.txt delete mode 100644 mlir-testing/fixed_output.mlir delete mode 100644 mlir-testing/integral_loop.txt delete mode 100644 mlir-testing/loopy_assign.txt delete mode 100644 mlir-testing/mlir_assign.mlir delete mode 100644 mlir-testing/mlir_assign_fixed.mlir delete mode 100644 mlir-testing/mlir_assign_opt.mlir delete mode 100644 mlir-testing/output.ll delete mode 100644 mlir-testing/run_compiled_mlir.py delete mode 100644 mlir-testing/sample.py delete mode 100644 mlir-testing/sample_2.py diff --git a/mlir-testing/assign_local_size.mlir b/mlir-testing/assign_local_size.mlir deleted file mode 100644 index 0af0ddd9cd..0000000000 --- a/mlir-testing/assign_local_size.mlir +++ /dev/null @@ -1,35 +0,0 @@ -builtin.module { - func.func @pyop3_loop( - %dat_0: tensor, - %dat_1: tensor, - %idat_0: tensor, - %idat_1: tensor, - %idat_2: tensor, - %idat_3: tensor - ) -> tensor { - - %c0 = arith.constant 0 : index // iter var - %c1 = arith.constant 1 : index // iter var - %c17 = arith.constant 17 : index // - %c15 = arith.constant 15 : index - %c32 = arith.constant 32 : index - %c2f = arith.constant 2.0 : f64 - - scf.for %i_0 = %c0 to %c17 step %c1 iter_args() -> () {} - - %f2 = scf.for %i_2 = %c0 to %c15 step %c1 iter_args(%dat_it = %dat_0) -> (tensor) { - %e1 = tensor.extract %idat_2[%i_2] : tensor - %ii_2 = arith.index_cast %e1 : i32 to index - %e2 = tensor.extract %idat_0[%ii_2] : tensor - %iii_2 = arith.index_cast %e2 : i32 to index - %v1 = tensor.extract %dat_1[%iii_2] : tensor - %v2 = arith.mulf %c2f, %v1 : f64 - %res = tensor.insert %v2 into %dat_it[%iii_2] : tensor - scf.yield %res : tensor - } - - scf.for %i_3 = %c0 to %c32 step %c1 iter_args() -> () {} - - func.return %f2 : tensor - } -} diff --git a/mlir-testing/assign_local_size.txt b/mlir-testing/assign_local_size.txt deleted file mode 100644 index d5e7205a71..0000000000 --- a/mlir-testing/assign_local_size.txt +++ /dev/null @@ -1,29 +0,0 @@ -******************************************************************************** -#include -#include -#include -#include - -void pyop3_loop(double *__restrict__ dat_0, double const *__restrict__ dat_1, int32_t const *__restrict__ idat_0, int32_t const *__restrict__ idat_1, int32_t const *__restrict__ idat_2, int32_t const *__restrict__ idat_3) -{ - for (int32_t i_0 = 0; i_0 <= 17; ++i_0) - { - } - for (int32_t i_2 = 0; i_2 <= 15; ++i_2) - dat_0[idat_0[idat_2[i_2]]] = 2.0 * dat_1[idat_0[idat_2[i_2]]]; - for (int32_t i_3 = 0; i_3 <= 32; ++i_3) - { - } - -} -******************************************************************************** -dat_0 (16) : [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] -dat_1 (16) : [10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10.] -idat_0 (67) : [ 0 0 0 0 0 1 2 3 3 3 3 4 4 4 4 5 5 5 5 6 6 6 6 7 - 7 7 7 8 8 8 8 9 9 9 9 9 9 10 10 10 10 11 11 11 11 12 12 12 - 12 13 13 13 13 13 13 13 13 14 14 14 14 15 15 15 15 15 15] -idat_1 (18) : [ 0 7 11 15 19 23 27 31 33 37 41 45 49 51 53 57 61 63] -idat_2 (16) : [ 4 5 6 10 14 18 22 26 30 36 40 44 48 56 60 66] -idat_3 (33) : [ 1 2 3 8 9 12 13 16 17 20 21 24 25 28 29 32 34 35 38 39 42 43 46 47 - 50 52 54 55 58 59 62 64 65] -******************************************************************************** diff --git a/mlir-testing/assign_mlir_demo.py b/mlir-testing/assign_mlir_demo.py deleted file mode 100644 index 6fca2ae981..0000000000 --- a/mlir-testing/assign_mlir_demo.py +++ /dev/null @@ -1,17 +0,0 @@ -from firedrake import * -import pyop3 as op3 -import numpy as np - -mesh = UnitSquareMesh(3,3) - -V = FunctionSpace(mesh, "CG", 1) -f = Function(V).assign(10) -g = Function(V) - -g.dat.assign( - 2 * f.dat, - eager=True, - eager_strategy="compile", - compiler_parameters={"codegen": "mlir"} -) -assert (g.dat.data_ro == 20).all() diff --git a/mlir-testing/assign_size.txt b/mlir-testing/assign_size.txt deleted file mode 100644 index 21b8e370be..0000000000 --- a/mlir-testing/assign_size.txt +++ /dev/null @@ -1,42 +0,0 @@ -******************************************************************************** -#include -#include -#include -#include - -void pyop3_loop(int64_t const *__restrict__ dat_0, int64_t const *__restrict__ dat_1, double *__restrict__ dat_2, double const *__restrict__ dat_3, int64_t const *__restrict__ dat_4, int64_t const *__restrict__ dat_5, int64_t const *__restrict__ dat_6, int64_t const *__restrict__ dat_7, int32_t const *__restrict__ idat_0, int32_t const *__restrict__ idat_1, int32_t const *__restrict__ idat_2, int32_t const *__restrict__ idat_3) -{ - int32_t p_0; - int32_t p_1; - int32_t p_2; - - p_0 = (int32_t) (dat_0[0] + dat_1[0]); - for (int32_t i_0 = 0; i_0 <= -1 + p_0; ++i_0) - { - } - p_1 = (int32_t) (dat_4[0] + dat_5[0]); - for (int32_t i_2 = 0; i_2 <= -1 + p_1; ++i_2) - dat_2[idat_0[idat_2[i_2]]] = 2.0 * dat_3[idat_0[idat_2[i_2]]]; - p_2 = (int32_t) (dat_6[0] + dat_7[0]); - for (int32_t i_3 = 0; i_3 <= -1 + p_2; ++i_3) - { - } - -} -******************************************************************************** -dat_0 (1) : [18] -dat_1 (1) : [0] -dat_2 (16) : [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] -dat_3 (16) : [10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10. 10.] -dat_4 (1) : [16] -dat_5 (1) : [0] -dat_6 (1) : [33] -dat_7 (1) : [0] -idat_0 (67) : [ 0 0 0 0 0 1 2 3 3 3 3 4 4 4 4 5 5 5 5 6 6 6 6 7 - 7 7 7 8 8 8 8 9 9 9 9 9 9 10 10 10 10 11 11 11 11 12 12 12 - 12 13 13 13 13 13 13 13 13 14 14 14 14 15 15 15 15 15 15] -idat_1 (18) : [ 0 7 11 15 19 23 27 31 33 37 41 45 49 51 53 57 61 63] -idat_2 (16) : [ 4 5 6 10 14 18 22 26 30 36 40 44 48 56 60 66] -idat_3 (33) : [ 1 2 3 8 9 12 13 16 17 20 21 24 25 28 29 32 34 35 38 39 42 43 46 47 - 50 52 54 55 58 59 62 64 65] -******************************************************************************** diff --git a/mlir-testing/fixed_output.mlir b/mlir-testing/fixed_output.mlir deleted file mode 100644 index 07b3b0af03..0000000000 --- a/mlir-testing/fixed_output.mlir +++ /dev/null @@ -1,159 +0,0 @@ -module { - llvm.func @memrefCopy(i64, !llvm.ptr, !llvm.ptr) - llvm.func @malloc(i64) -> !llvm.ptr - llvm.func @pyop3_loop(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: i64, %arg3: i64, %arg4: i64, %arg5: !llvm.ptr, %arg6: !llvm.ptr, %arg7: i64, %arg8: i64, %arg9: i64, %arg10: !llvm.ptr, %arg11: !llvm.ptr, %arg12: i64, %arg13: i64, %arg14: i64, %arg15: !llvm.ptr, %arg16: !llvm.ptr, %arg17: i64, %arg18: i64, %arg19: i64, %arg20: !llvm.ptr, %arg21: !llvm.ptr, %arg22: i64, %arg23: i64, %arg24: i64, %arg25: !llvm.ptr, %arg26: !llvm.ptr, %arg27: i64, %arg28: i64, %arg29: i64) -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> attributes {llvm.emit_c_interface} { - %0 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %1 = llvm.insertvalue %arg20, %0[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %2 = llvm.insertvalue %arg21, %1[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %3 = llvm.insertvalue %arg22, %2[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %4 = llvm.insertvalue %arg23, %3[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %5 = llvm.insertvalue %arg24, %4[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %6 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %7 = llvm.insertvalue %arg15, %6[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %8 = llvm.insertvalue %arg16, %7[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %9 = llvm.insertvalue %arg17, %8[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %10 = llvm.insertvalue %arg18, %9[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %11 = llvm.insertvalue %arg19, %10[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %12 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %13 = llvm.insertvalue %arg5, %12[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %14 = llvm.insertvalue %arg6, %13[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %15 = llvm.insertvalue %arg7, %14[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %16 = llvm.insertvalue %arg8, %15[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %17 = llvm.insertvalue %arg9, %16[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %18 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %19 = llvm.insertvalue %arg0, %18[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %20 = llvm.insertvalue %arg1, %19[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %21 = llvm.insertvalue %arg2, %20[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %22 = llvm.insertvalue %arg3, %21[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %23 = llvm.insertvalue %arg4, %22[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %24 = llvm.mlir.constant(16 : index) : i64 - %25 = llvm.mlir.constant(2.000000e+00 : f64) : f64 - %26 = llvm.mlir.constant(0 : index) : i64 - %27 = llvm.mlir.constant(1 : index) : i64 - llvm.br ^bb1(%26 : i64) - ^bb1(%28: i64): // 2 preds: ^bb0, ^bb2 - %29 = llvm.icmp "slt" %28, %24 : i64 - llvm.cond_br %29, ^bb2, ^bb3 - ^bb2: // pred: ^bb1 - %30 = llvm.extractvalue %5[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %31 = llvm.extractvalue %5[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %32 = llvm.getelementptr %30[%31] : (!llvm.ptr, i64) -> !llvm.ptr, i32 - %33 = llvm.extractvalue %5[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %34 = llvm.mul %28, %33 overflow : i64 - %35 = llvm.getelementptr inbounds|nuw %32[%34] : (!llvm.ptr, i64) -> !llvm.ptr, i32 - %36 = llvm.load %35 : !llvm.ptr -> i32 - %37 = llvm.sext %36 : i32 to i64 - %38 = llvm.extractvalue %17[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %39 = llvm.extractvalue %17[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %40 = llvm.getelementptr %38[%39] : (!llvm.ptr, i64) -> !llvm.ptr, i32 - %41 = llvm.extractvalue %17[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %42 = llvm.mul %37, %41 overflow : i64 - %43 = llvm.getelementptr inbounds|nuw %40[%42] : (!llvm.ptr, i64) -> !llvm.ptr, i32 - %44 = llvm.load %43 : !llvm.ptr -> i32 - %45 = llvm.sext %44 : i32 to i64 - %46 = llvm.extractvalue %11[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %47 = llvm.extractvalue %11[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %48 = llvm.getelementptr %46[%47] : (!llvm.ptr, i64) -> !llvm.ptr, f64 - %49 = llvm.extractvalue %11[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %50 = llvm.mul %45, %49 overflow : i64 - %51 = llvm.getelementptr inbounds|nuw %48[%50] : (!llvm.ptr, i64) -> !llvm.ptr, f64 - %52 = llvm.load %51 : !llvm.ptr -> f64 - %53 = llvm.fmul %52, %25 : f64 - %54 = llvm.extractvalue %23[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %55 = llvm.extractvalue %23[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %56 = llvm.getelementptr %54[%55] : (!llvm.ptr, i64) -> !llvm.ptr, f64 - %57 = llvm.extractvalue %23[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %58 = llvm.mul %45, %57 overflow : i64 - %59 = llvm.getelementptr inbounds|nuw %56[%58] : (!llvm.ptr, i64) -> !llvm.ptr, f64 - llvm.store %53, %59 : f64, !llvm.ptr - %60 = llvm.add %28, %27 : i64 - llvm.br ^bb1(%60 : i64) - ^bb3: // pred: ^bb1 - %61 = llvm.mlir.constant(1 : index) : i64 - %62 = llvm.extractvalue %23[3] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %63 = llvm.alloca %61 x !llvm.array<1 x i64> : (i64) -> !llvm.ptr - llvm.store %62, %63 : !llvm.array<1 x i64>, !llvm.ptr - %64 = llvm.getelementptr %63[0, %26] : (!llvm.ptr, i64) -> !llvm.ptr, !llvm.array<1 x i64> - %65 = llvm.load %64 : !llvm.ptr -> i64 - %66 = llvm.mlir.constant(1 : index) : i64 - %67 = llvm.mlir.zero : !llvm.ptr - %68 = llvm.getelementptr %67[%65] : (!llvm.ptr, i64) -> !llvm.ptr, f64 - %69 = llvm.ptrtoint %68 : !llvm.ptr to i64 - %70 = llvm.call @malloc(%69) : (i64) -> !llvm.ptr - %71 = llvm.mlir.poison : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %72 = llvm.insertvalue %70, %71[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %73 = llvm.insertvalue %70, %72[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %74 = llvm.mlir.constant(0 : index) : i64 - %75 = llvm.insertvalue %74, %73[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %76 = llvm.insertvalue %65, %75[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %77 = llvm.insertvalue %66, %76[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %78 = llvm.intr.stacksave : !llvm.ptr - %79 = llvm.mlir.constant(1 : i64) : i64 - %80 = llvm.mlir.constant(1 : index) : i64 - %81 = llvm.alloca %80 x !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> : (i64) -> !llvm.ptr - llvm.store %23, %81 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>, !llvm.ptr - %82 = llvm.mlir.poison : !llvm.struct<(i64, ptr)> - %83 = llvm.insertvalue %79, %82[0] : !llvm.struct<(i64, ptr)> - %84 = llvm.insertvalue %81, %83[1] : !llvm.struct<(i64, ptr)> - %85 = llvm.mlir.constant(1 : i64) : i64 - %86 = llvm.mlir.constant(1 : index) : i64 - %87 = llvm.alloca %86 x !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> : (i64) -> !llvm.ptr - llvm.store %77, %87 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>, !llvm.ptr - %88 = llvm.mlir.poison : !llvm.struct<(i64, ptr)> - %89 = llvm.insertvalue %85, %88[0] : !llvm.struct<(i64, ptr)> - %90 = llvm.insertvalue %87, %89[1] : !llvm.struct<(i64, ptr)> - %91 = llvm.mlir.constant(1 : index) : i64 - %92 = llvm.alloca %91 x !llvm.struct<(i64, ptr)> : (i64) -> !llvm.ptr - llvm.store %84, %92 : !llvm.struct<(i64, ptr)>, !llvm.ptr - %93 = llvm.alloca %91 x !llvm.struct<(i64, ptr)> : (i64) -> !llvm.ptr - llvm.store %90, %93 : !llvm.struct<(i64, ptr)>, !llvm.ptr - %94 = llvm.mlir.zero : !llvm.ptr - %95 = llvm.getelementptr %94[1] : (!llvm.ptr) -> !llvm.ptr, f64 - %96 = llvm.ptrtoint %95 : !llvm.ptr to i64 - llvm.call @memrefCopy(%96, %92, %93) : (i64, !llvm.ptr, !llvm.ptr) -> () - llvm.intr.stackrestore %78 : !llvm.ptr - llvm.return %77 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - } - llvm.func @_mlir_ciface_pyop3_loop(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: !llvm.ptr, %arg3: !llvm.ptr, %arg4: !llvm.ptr, %arg5: !llvm.ptr, %arg6: !llvm.ptr) attributes {llvm.emit_c_interface} { - %0 = llvm.load %arg1 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %1 = llvm.extractvalue %0[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %2 = llvm.extractvalue %0[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %3 = llvm.extractvalue %0[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %4 = llvm.extractvalue %0[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %5 = llvm.extractvalue %0[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %6 = llvm.load %arg2 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %7 = llvm.extractvalue %6[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %8 = llvm.extractvalue %6[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %9 = llvm.extractvalue %6[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %10 = llvm.extractvalue %6[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %11 = llvm.extractvalue %6[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %12 = llvm.load %arg3 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %13 = llvm.extractvalue %12[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %14 = llvm.extractvalue %12[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %15 = llvm.extractvalue %12[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %16 = llvm.extractvalue %12[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %17 = llvm.extractvalue %12[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %18 = llvm.load %arg4 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %19 = llvm.extractvalue %18[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %20 = llvm.extractvalue %18[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %21 = llvm.extractvalue %18[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %22 = llvm.extractvalue %18[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %23 = llvm.extractvalue %18[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %24 = llvm.load %arg5 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %25 = llvm.extractvalue %24[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %26 = llvm.extractvalue %24[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %27 = llvm.extractvalue %24[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %28 = llvm.extractvalue %24[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %29 = llvm.extractvalue %24[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %30 = llvm.load %arg6 : !llvm.ptr -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %31 = llvm.extractvalue %30[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %32 = llvm.extractvalue %30[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %33 = llvm.extractvalue %30[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %34 = llvm.extractvalue %30[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %35 = llvm.extractvalue %30[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - %36 = llvm.call @pyop3_loop(%1, %2, %3, %4, %5, %7, %8, %9, %10, %11, %13, %14, %15, %16, %17, %19, %20, %21, %22, %23, %25, %26, %27, %28, %29, %31, %32, %33, %34, %35) : (!llvm.ptr, !llvm.ptr, i64, i64, i64, !llvm.ptr, !llvm.ptr, i64, i64, i64, !llvm.ptr, !llvm.ptr, i64, i64, i64, !llvm.ptr, !llvm.ptr, i64, i64, i64, !llvm.ptr, !llvm.ptr, i64, i64, i64, !llvm.ptr, !llvm.ptr, i64, i64, i64) -> !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)> - llvm.store %36, %arg0 : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>, !llvm.ptr - llvm.return - } -} - diff --git a/mlir-testing/integral_loop.txt b/mlir-testing/integral_loop.txt deleted file mode 100644 index feac230fa3..0000000000 --- a/mlir-testing/integral_loop.txt +++ /dev/null @@ -1,68 +0,0 @@ -******************************************************************************** -#include -#include -#include -#include -#include - -static void form_cell_integral(double *__restrict__ A, double const *__restrict__ coords_0); -static void form_cell_integral(double *__restrict__ A, double const *__restrict__ coords_0) -{ - double t0; - double t1; - double t2; - double t3[3] = { 0.33333333333333337, 0.33333333333333326, 0.33333333333333326 }; - - t0 = -1.0 * coords_0[0]; - t1 = -1.0 * coords_0[1]; - t2 = 0.5 * fabs((t0 + coords_0[2]) * (t1 + coords_0[5]) + -1.0 * (t0 + coords_0[4]) * (t1 + coords_0[3])); - for (int32_t j = 0; j <= 2; ++j) - A[j] = A[j] + t3[j] * t2; - -} - -void pyop3_loop(double const *__restrict__ dat_0, double *__restrict__ dat_1, int32_t const *__restrict__ idat_0, int32_t const *__restrict__ idat_1) -{ - int32_t j_0; - int32_t j_1; - int32_t j_2; - double t_0[3l]; - double t_1[6l]; - - for (int32_t i_0 = 0; i_0 <= 17; ++i_0) - { - for (int32_t i_1 = 0; i_1 <= 2; ++i_1) - { - j_0 = 0 + 1 * (i_1 + 0); - t_0[j_0] = (double) (0.0); - } - for (int32_t i_2 = 0; i_2 <= 2; ++i_2) - for (int32_t i_3 = 0; i_3 <= 1; ++i_3) - { - j_1 = 0 + 1 * (i_2 * 2 + 0 + i_3); - t_1[j_1] = dat_0[idat_0[3 * i_0 + i_2] + i_3]; - } - form_cell_integral(&(t_0[0]), &(t_1[0])); - for (int32_t i_6 = 0; i_6 <= 2; ++i_6) - { - j_2 = 0 + 1 * (i_6 + 0); - dat_1[idat_1[3 * i_0 + i_6]] = dat_1[idat_1[3 * i_0 + i_6]] + t_0[j_2]; - } - } - -} -******************************************************************************** -dat_0 (32) : [0.33333333 0. 0. 0.33333333 0. 0. - 0.33333333 0.33333333 0. 0.66666667 0.66666667 0. - 0.33333333 0.66666667 0.66666667 0.33333333 0. 1. - 1. 0. 0.33333333 1. 0.66666667 0.66666667 - 1. 0.33333333 0.66666667 1. 1. 0.66666667 - 1. 1. ] -dat_1 (16) : [0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0. 0.] -idat_0 (54) : [ 0 2 4 0 2 6 2 6 8 0 6 10 6 8 12 6 10 14 8 12 16 6 12 14 - 10 14 18 12 16 20 12 14 22 14 18 24 12 20 22 14 22 24 20 22 26 22 24 28 - 22 26 28 26 28 30] -idat_1 (54) : [ 0 1 2 0 1 3 1 3 4 0 3 5 3 4 6 3 5 7 4 6 8 3 6 7 - 5 7 9 6 8 10 6 7 11 7 9 12 6 10 11 7 11 12 10 11 13 11 12 14 - 11 13 14 13 14 15] -******************************************************************************** diff --git a/mlir-testing/loopy_assign.txt b/mlir-testing/loopy_assign.txt deleted file mode 100644 index d14de3083d..0000000000 --- a/mlir-testing/loopy_assign.txt +++ /dev/null @@ -1,38 +0,0 @@ ---------------------------------------------------------------------------- -KERNEL: pyop3_loop ---------------------------------------------------------------------------- -ARGUMENTS: -dat_0: ArrayArg, type: np:dtype('float64'), shape: unknown in/out aspace: global -dat_1: ArrayArg, type: np:dtype('float64'), shape: unknown in aspace: global -idat_0: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global -idat_1: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global -idat_2: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global -idat_3: ArrayArg, type: np:dtype('int32'), shape: unknown in aspace: global ---------------------------------------------------------------------------- -DOMAINS: -{ [i_0] : 0 <= i_0 <= 17 } -{ [i_1] : 1 = 0 } -{ [i_2] : 0 <= i_2 <= 15 } -{ [i_3] : 0 <= i_3 <= 32 } -{ [i_4] : 1 = 0 } ---------------------------------------------------------------------------- -INAME TAGS: -i_0: None -i_1: None -i_2: None -i_3: None -i_4: None ---------------------------------------------------------------------------- -INSTRUCTIONS: - for i_1, i_0 -↱ dat_0[0 + 1*(idat_0[0 + 1*idat_1[0 + 1*i_0]] + i_1)] = 2*dat_1[0 + 1*(idat_0[0 + 1*idat_1[0 + 1*i_0]] + i_1)] {id=insn_0} -│ end i_1, i_0 -│ for i_2 -└↱ dat_0[0 + 1*(idat_0[0 + 1*idat_2[0 + 1*i_2]] + 0)] = 2*dat_1[0 + 1*(idat_0[0 + 1*idat_2[0 + 1*i_2]] + 0) {id=insn_1} - │ end i_2 - │ for i_4, i_3 -↱└ dat_0[0 + 1*(idat_0[0 + 1*idat_3[0 + 1*i_3]] + i_4) = 2*dat_1[0 + 1*(idat_0[0 + 1*idat_3[0 + 1*i_3]] + i_4) {id=insn_2} -│ end i_4, i_3 -└ CODE(idat_2, dat_0, idat_0, idat_1, dat_1, idat_3|) {id=insn} - ---------------------------------------------------------------------------- diff --git a/mlir-testing/mlir_assign.mlir b/mlir-testing/mlir_assign.mlir deleted file mode 100644 index 4eab0e02b9..0000000000 --- a/mlir-testing/mlir_assign.mlir +++ /dev/null @@ -1,172 +0,0 @@ -builtin.module { - func.func @pyop3_loop(%0: tensor, %1: tensor, %2: tensor, %3: tensor, %4: tensor, %5: tensor) { - %6 = arith.constant 0 : i32 - %7 = arith.index_cast %6 : i32 to index - %8 = arith.constant 18 : i32 - %9 = arith.index_cast %8 : i32 to index - %10 = arith.constant 1 : i32 - %11 = arith.index_cast %10 : i32 to index - %12 = scf.for %13 = %7 to %9 step %11 iter_args(%14 = %0) -> (tensor) { - %15 = arith.constant 0 : i32 - %16 = arith.index_cast %15 : i32 to index - %17 = arith.constant 0 : i32 - %18 = arith.index_cast %17 : i32 to index - %19 = arith.constant 1 : i32 - %20 = arith.index_cast %19 : i32 to index - %21 = scf.for %22 = %16 to %18 step %20 iter_args(%23 = %14) -> (tensor) { - %24 = arith.constant 2 : i32 - %25 = arith.constant 0 : i32 - %26 = arith.constant 1 : i32 - %27 = arith.constant 0 : i32 - %28 = arith.constant 1 : i32 - %29 = arith.constant 0 : i32 - %30 = arith.constant 1 : i32 - %31 = arith.muli %30, %13 : i32 - %32 = arith.addi %29, %31 : i32 - %33 = arith.index_cast %32 : i32 to index - %34 = tensor.extract %2[%33] : tensor - %35 = arith.muli %28, %34 : i32 - %36 = arith.addi %27, %35 : i32 - %37 = arith.index_cast %36 : i32 to index - %38 = tensor.extract %1[%37] : tensor - %39 = arith.addi %38, %22 : i32 - %40 = arith.muli %26, %39 : i32 - %41 = arith.addi %25, %40 : i32 - %42 = arith.index_cast %41 : i32 to index - %43 = tensor.extract %3[%42] : tensor - %44 = arith.muli %24, %43 : i32 - %45 = arith.constant 0 : i32 - %46 = arith.constant 1 : i32 - %47 = arith.constant 0 : i32 - %48 = arith.constant 1 : i32 - %49 = arith.constant 0 : i32 - %50 = arith.constant 1 : i32 - %51 = arith.muli %50, %13 : i32 - %52 = arith.addi %49, %51 : i32 - %53 = arith.index_cast %52 : i32 to index - %54 = tensor.extract %2[%53] : tensor - %55 = arith.muli %48, %54 : i32 - %56 = arith.addi %47, %55 : i32 - %57 = arith.index_cast %56 : i32 to index - %58 = tensor.extract %1[%57] : tensor - %59 = arith.addi %58, %22 : i32 - %60 = arith.muli %46, %59 : i32 - %61 = arith.addi %45, %60 : i32 - %62 = arith.index_cast %61 : i32 to index - %63 = tensor.insert %44 into %23[%62] : tensor - scf.yield %63 : tensor - } - scf.yield %21 : tensor - } - %64 = arith.constant 0 : i32 - %65 = arith.index_cast %64 : i32 to index - %66 = arith.constant 16 : i32 - %67 = arith.index_cast %66 : i32 to index - %68 = arith.constant 1 : i32 - %69 = arith.index_cast %68 : i32 to index - %70 = scf.for %71 = %65 to %67 step %69 iter_args(%72 = %12) -> (tensor) { - %73 = arith.constant 2 : i32 - %74 = arith.constant 0 : i32 - %75 = arith.constant 1 : i32 - %76 = arith.constant 0 : i32 - %77 = arith.constant 1 : i32 - %78 = arith.constant 0 : i32 - %79 = arith.constant 1 : i32 - %80 = arith.muli %79, %71 : i32 - %81 = arith.addi %78, %80 : i32 - %82 = arith.index_cast %81 : i32 to index - %83 = tensor.extract %4[%82] : tensor - %84 = arith.muli %77, %83 : i32 - %85 = arith.addi %76, %84 : i32 - %86 = arith.index_cast %85 : i32 to index - %87 = tensor.extract %1[%86] : tensor - %88 = arith.constant 0 : i32 - %89 = arith.addi %87, %88 : i32 - %90 = arith.muli %75, %89 : i32 - %91 = arith.addi %74, %90 : i32 - %92 = arith.index_cast %91 : i32 to index - %93 = tensor.extract %3[%92] : tensor - %94 = arith.muli %73, %93 : i32 - %95 = arith.constant 0 : i32 - %96 = arith.constant 1 : i32 - %97 = arith.constant 0 : i32 - %98 = arith.constant 1 : i32 - %99 = arith.constant 0 : i32 - %100 = arith.constant 1 : i32 - %101 = arith.muli %100, %71 : i32 - %102 = arith.addi %99, %101 : i32 - %103 = arith.index_cast %102 : i32 to index - %104 = tensor.extract %4[%103] : tensor - %105 = arith.muli %98, %104 : i32 - %106 = arith.addi %97, %105 : i32 - %107 = arith.index_cast %106 : i32 to index - %108 = tensor.extract %1[%107] : tensor - %109 = arith.constant 0 : i32 - %110 = arith.addi %108, %109 : i32 - %111 = arith.muli %96, %110 : i32 - %112 = arith.addi %95, %111 : i32 - %113 = arith.index_cast %112 : i32 to index - %114 = tensor.insert %94 into %72[%113] : tensor - scf.yield %114 : tensor - } - %115 = arith.constant 0 : i32 - %116 = arith.index_cast %115 : i32 to index - %117 = arith.constant 33 : i32 - %118 = arith.index_cast %117 : i32 to index - %119 = arith.constant 1 : i32 - %120 = arith.index_cast %119 : i32 to index - %121 = scf.for %122 = %116 to %118 step %120 iter_args(%123 = %70) -> (tensor) { - %124 = arith.constant 0 : i32 - %125 = arith.index_cast %124 : i32 to index - %126 = arith.constant 0 : i32 - %127 = arith.index_cast %126 : i32 to index - %128 = arith.constant 1 : i32 - %129 = arith.index_cast %128 : i32 to index - %130 = scf.for %131 = %125 to %127 step %129 iter_args(%132 = %123) -> (tensor) { - %133 = arith.constant 2 : i32 - %134 = arith.constant 0 : i32 - %135 = arith.constant 1 : i32 - %136 = arith.constant 0 : i32 - %137 = arith.constant 1 : i32 - %138 = arith.constant 0 : i32 - %139 = arith.constant 1 : i32 - %140 = arith.muli %139, %122 : i32 - %141 = arith.addi %138, %140 : i32 - %142 = arith.index_cast %141 : i32 to index - %143 = tensor.extract %5[%142] : tensor - %144 = arith.muli %137, %143 : i32 - %145 = arith.addi %136, %144 : i32 - %146 = arith.index_cast %145 : i32 to index - %147 = tensor.extract %1[%146] : tensor - %148 = arith.addi %147, %131 : i32 - %149 = arith.muli %135, %148 : i32 - %150 = arith.addi %134, %149 : i32 - %151 = arith.index_cast %150 : i32 to index - %152 = tensor.extract %3[%151] : tensor - %153 = arith.muli %133, %152 : i32 - %154 = arith.constant 0 : i32 - %155 = arith.constant 1 : i32 - %156 = arith.constant 0 : i32 - %157 = arith.constant 1 : i32 - %158 = arith.constant 0 : i32 - %159 = arith.constant 1 : i32 - %160 = arith.muli %159, %122 : i32 - %161 = arith.addi %158, %160 : i32 - %162 = arith.index_cast %161 : i32 to index - %163 = tensor.extract %5[%162] : tensor - %164 = arith.muli %157, %163 : i32 - %165 = arith.addi %156, %164 : i32 - %166 = arith.index_cast %165 : i32 to index - %167 = tensor.extract %1[%166] : tensor - %168 = arith.addi %167, %131 : i32 - %169 = arith.muli %155, %168 : i32 - %170 = arith.addi %154, %169 : i32 - %171 = arith.index_cast %170 : i32 to index - %172 = tensor.insert %153 into %132[%171] : tensor - scf.yield %172 : tensor - } - scf.yield %130 : tensor - } - func.return - } -} diff --git a/mlir-testing/mlir_assign_fixed.mlir b/mlir-testing/mlir_assign_fixed.mlir deleted file mode 100644 index 20fdd5bccc..0000000000 --- a/mlir-testing/mlir_assign_fixed.mlir +++ /dev/null @@ -1,152 +0,0 @@ -// 62 -builtin.module { - func.func @pyop3_loop(%0: tensor, %1: tensor, %2: tensor, %3: tensor, %4: tensor, %5: tensor) -> tensor attributes {llvm.emit_c_interface} { - %7 = arith.constant 0 : index - %9 = arith.constant 18 : index - %11 = arith.constant 1 : index - %12 = scf.for %13 = %7 to %9 step %11 iter_args(%14 = %0) -> (tensor) { - %16 = arith.constant 0 : index - %18 = arith.constant 0 : index - %20 = arith.constant 1 : index - %21 = scf.for %22 = %16 to %18 step %20 iter_args(%23 = %14) -> (tensor) { - %24 = arith.constant 2. : f64 - %25 = arith.constant 0 : index - %26 = arith.constant 1 : index - %27 = arith.constant 0 : index - %28 = arith.constant 1 : index - %29 = arith.constant 0 : index - %30 = arith.constant 1 : index - %31 = arith.muli %30, %13 : index - %32 = arith.addi %29, %31 : index - %34 = tensor.extract %2[%32] : tensor - %15 = arith.index_cast %34 : i32 to index - %35 = arith.muli %28, %15 : index - %36 = arith.addi %27, %35 : index - %38 = tensor.extract %1[%36] : tensor - %17 = arith.index_cast %38 : i32 to index - %39 = arith.addi %17, %22 : index - %40 = arith.muli %26, %39 : index - %41 = arith.addi %25, %40 : index - %43 = tensor.extract %3[%41] : tensor - %44 = arith.mulf %24, %43 : f64 - %45 = arith.constant 0 : index - %46 = arith.constant 1 : index - %47 = arith.constant 0 : index - %48 = arith.constant 1 : index - %49 = arith.constant 0 : index - %50 = arith.constant 1 : index - %51 = arith.muli %50, %13 : index - %52 = arith.addi %49, %51 : index - %54 = tensor.extract %2[%52] : tensor - %19 = arith.index_cast %54 : i32 to index - %55 = arith.muli %48, %19 : index - %56 = arith.addi %47, %55 : index - %58 = tensor.extract %1[%56] : tensor - %57 = arith.index_cast %58 : i32 to index - %59 = arith.addi %57, %22 : index - %60 = arith.muli %46, %59 : index - %61 = arith.addi %45, %60 : index - %63 = tensor.insert %44 into %23[%61] : tensor - scf.yield %63 : tensor - } - scf.yield %21 : tensor - } - %65 = arith.constant 0 : index - %67 = arith.constant 16 : index - %69 = arith.constant 1 : index - %70 = scf.for %71 = %65 to %67 step %69 iter_args(%72 = %12) -> (tensor) { - %73 = arith.constant 2. : f64 - %74 = arith.constant 0 : index - %75 = arith.constant 1 : index - %76 = arith.constant 0 : index - %77 = arith.constant 1 : index - %78 = arith.constant 0 : index - %79 = arith.constant 1 : index - %80 = arith.muli %79, %71 : index - %81 = arith.addi %78, %80 : index - %83 = tensor.extract %4[%81] : tensor - %82 = arith.index_cast %83 : i32 to index - %84 = arith.muli %77, %82 : index - %85 = arith.addi %76, %84 : index - %87 = tensor.extract %1[%85] : tensor - %86 = arith.index_cast %87 : i32 to index - %88 = arith.constant 0 : index - %89 = arith.addi %86, %88 : index - %90 = arith.muli %75, %89 : index - %91 = arith.addi %74, %90 : index - %93 = tensor.extract %3[%91] : tensor - %94 = arith.mulf %73, %93 : f64 - %95 = arith.constant 0 : index - %96 = arith.constant 1 : index - %97 = arith.constant 0 : index - %98 = arith.constant 1 : index - %99 = arith.constant 0 : index - %100 = arith.constant 1 : index - %101 = arith.muli %100, %71 : index - %102 = arith.addi %99, %101 : index - %104 = tensor.extract %4[%102] : tensor - %92 = arith.index_cast %104 : i32 to index - %105 = arith.muli %98, %92 : index - %106 = arith.addi %97, %105 : index - %108 = tensor.extract %1[%106] : tensor - %107 = arith.index_cast %108 : i32 to index - %109 = arith.constant 0 : index - %110 = arith.addi %107, %109 : index - %111 = arith.muli %96, %110 : index - %112 = arith.addi %95, %111 : index - %114 = tensor.insert %94 into %72[%112] : tensor - scf.yield %114 : tensor - } - %116 = arith.constant 0 : index - %118 = arith.constant 33 : index - %120 = arith.constant 1 : index - %121 = scf.for %122 = %116 to %118 step %120 iter_args(%123 = %70) -> (tensor) { - %125 = arith.constant 0 : index - %127 = arith.constant 0 : index - %129 = arith.constant 1 : index - %130 = scf.for %131 = %125 to %127 step %129 iter_args(%132 = %123) -> (tensor) { - %133 = arith.constant 2. : f64 - %134 = arith.constant 0 : index - %135 = arith.constant 1 : index - %136 = arith.constant 0 : index - %137 = arith.constant 1 : index - %138 = arith.constant 0 : index - %139 = arith.constant 1 : index - %140 = arith.muli %139, %122 : index - %141 = arith.addi %138, %140 : index - %143 = tensor.extract %5[%141] : tensor - %142 = arith.index_cast %143 : i32 to index - %144 = arith.muli %137, %142 : index - %145 = arith.addi %136, %144 : index - %147 = tensor.extract %1[%145] : tensor - %146 = arith.index_cast %147 : i32 to index - %148 = arith.addi %146, %131 : index - %149 = arith.muli %135, %148 : index - %150 = arith.addi %134, %149 : index - %152 = tensor.extract %3[%150] : tensor - %153 = arith.mulf %133, %152 : f64 - %154 = arith.constant 0 : index - %155 = arith.constant 1 : index - %156 = arith.constant 0 : index - %157 = arith.constant 1 : index - %158 = arith.constant 0 : index - %159 = arith.constant 1 : index - %160 = arith.muli %159, %122 : index - %161 = arith.addi %158, %160 : index - %163 = tensor.extract %5[%161] : tensor - %162 = arith.index_cast %163 : i32 to index - %164 = arith.muli %157, %162 : index - %165 = arith.addi %156, %164 : index - %167 = tensor.extract %1[%165] : tensor - %166 = arith.index_cast %167 : i32 to index - %168 = arith.addi %166, %131 : index - %169 = arith.muli %155, %168 : index - %170 = arith.addi %154, %169 : index - %172 = tensor.insert %153 into %132[%170] : tensor - scf.yield %172 : tensor - } - scf.yield %130 : tensor - } - func.return %121 : tensor - } -} diff --git a/mlir-testing/mlir_assign_opt.mlir b/mlir-testing/mlir_assign_opt.mlir deleted file mode 100644 index 7ca17bd78f..0000000000 --- a/mlir-testing/mlir_assign_opt.mlir +++ /dev/null @@ -1,20 +0,0 @@ -module { - func.func @pyop3_loop(%arg0: tensor, %arg1: tensor, %arg2: tensor, %arg3: tensor, %arg4: tensor, %arg5: tensor) -> tensor attributes {llvm.emit_c_interface} { - %c16 = arith.constant 16 : index - %cst = arith.constant 2.000000e+00 : f64 - %c0 = arith.constant 0 : index - %c1 = arith.constant 1 : index - %0 = scf.for %arg6 = %c0 to %c16 step %c1 iter_args(%arg7 = %arg0) -> (tensor) { - %extracted = tensor.extract %arg4[%arg6] : tensor - %1 = arith.index_cast %extracted : i32 to index - %extracted_0 = tensor.extract %arg1[%1] : tensor - %2 = arith.index_cast %extracted_0 : i32 to index - %extracted_1 = tensor.extract %arg3[%2] : tensor - %3 = arith.mulf %extracted_1, %cst : f64 - %inserted = tensor.insert %3 into %arg7[%2] : tensor - scf.yield %inserted : tensor - } - return %0 : tensor - } -} - diff --git a/mlir-testing/output.ll b/mlir-testing/output.ll deleted file mode 100644 index 9a0a3411cd..0000000000 --- a/mlir-testing/output.ll +++ /dev/null @@ -1,153 +0,0 @@ -; ModuleID = 'LLVMDialectModule' -source_filename = "LLVMDialectModule" - -declare void @memrefCopy(i64, ptr, ptr) - -declare ptr @malloc(i64) - -define { ptr, ptr, i64, [1 x i64], [1 x i64] } @pyop3_loop(ptr %0, ptr %1, i64 %2, i64 %3, i64 %4, ptr %5, ptr %6, i64 %7, i64 %8, i64 %9, ptr %10, ptr %11, i64 %12, i64 %13, i64 %14, ptr %15, ptr %16, i64 %17, i64 %18, i64 %19, ptr %20, ptr %21, i64 %22, i64 %23, i64 %24, ptr %25, ptr %26, i64 %27, i64 %28, i64 %29) { - %31 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } poison, ptr %20, 0 - %32 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %31, ptr %21, 1 - %33 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %32, i64 %22, 2 - %34 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %33, i64 %23, 3, 0 - %35 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %34, i64 %24, 4, 0 - %36 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } poison, ptr %15, 0 - %37 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %36, ptr %16, 1 - %38 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %37, i64 %17, 2 - %39 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %38, i64 %18, 3, 0 - %40 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %39, i64 %19, 4, 0 - %41 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } poison, ptr %5, 0 - %42 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %41, ptr %6, 1 - %43 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %42, i64 %7, 2 - %44 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %43, i64 %8, 3, 0 - %45 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %44, i64 %9, 4, 0 - %46 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } poison, ptr %0, 0 - %47 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %46, ptr %1, 1 - %48 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %47, i64 %2, 2 - %49 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %48, i64 %3, 3, 0 - %50 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %49, i64 %4, 4, 0 - br label %51 - -51: ; preds = %54, %30 - %52 = phi i64 [ %85, %54 ], [ 0, %30 ] - %53 = icmp slt i64 %52, 16 - br i1 %53, label %54, label %86 - -54: ; preds = %51 - %55 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %35, 1 - %56 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %35, 2 - %57 = getelementptr i32, ptr %55, i64 %56 - %58 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %35, 4, 0 - %59 = mul nuw nsw i64 %52, %58 - %60 = getelementptr inbounds nuw i32, ptr %57, i64 %59 - %61 = load i32, ptr %60, align 4 - %62 = sext i32 %61 to i64 - %63 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %45, 1 - %64 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %45, 2 - %65 = getelementptr i32, ptr %63, i64 %64 - %66 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %45, 4, 0 - %67 = mul nuw nsw i64 %62, %66 - %68 = getelementptr inbounds nuw i32, ptr %65, i64 %67 - %69 = load i32, ptr %68, align 4 - %70 = sext i32 %69 to i64 - %71 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %40, 1 - %72 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %40, 2 - %73 = getelementptr double, ptr %71, i64 %72 - %74 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %40, 4, 0 - %75 = mul nuw nsw i64 %70, %74 - %76 = getelementptr inbounds nuw double, ptr %73, i64 %75 - %77 = load double, ptr %76, align 8 - %78 = fmul double %77, 2.000000e+00 - %79 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %50, 1 - %80 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %50, 2 - %81 = getelementptr double, ptr %79, i64 %80 - %82 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %50, 4, 0 - %83 = mul nuw nsw i64 %70, %82 - %84 = getelementptr inbounds nuw double, ptr %81, i64 %83 - store double %78, ptr %84, align 8 - %85 = add i64 %52, 1 - br label %51 - -86: ; preds = %51 - %87 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %50, 3 - %88 = alloca [1 x i64], i64 1, align 8 - store [1 x i64] %87, ptr %88, align 4 - %89 = getelementptr [1 x i64], ptr %88, i32 0, i64 0 - %90 = load i64, ptr %89, align 4 - %91 = getelementptr double, ptr null, i64 %90 - %92 = ptrtoint ptr %91 to i64 - %93 = call ptr @malloc(i64 %92) - %94 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } poison, ptr %93, 0 - %95 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %94, ptr %93, 1 - %96 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %95, i64 0, 2 - %97 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %96, i64 %90, 3, 0 - %98 = insertvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %97, i64 1, 4, 0 - %99 = call ptr @llvm.stacksave.p0() - %100 = alloca { ptr, ptr, i64, [1 x i64], [1 x i64] }, i64 1, align 8 - store { ptr, ptr, i64, [1 x i64], [1 x i64] } %50, ptr %100, align 8 - %101 = insertvalue { i64, ptr } { i64 1, ptr poison }, ptr %100, 1 - %102 = alloca { ptr, ptr, i64, [1 x i64], [1 x i64] }, i64 1, align 8 - store { ptr, ptr, i64, [1 x i64], [1 x i64] } %98, ptr %102, align 8 - %103 = insertvalue { i64, ptr } { i64 1, ptr poison }, ptr %102, 1 - %104 = alloca { i64, ptr }, i64 1, align 8 - store { i64, ptr } %101, ptr %104, align 8 - %105 = alloca { i64, ptr }, i64 1, align 8 - store { i64, ptr } %103, ptr %105, align 8 - call void @memrefCopy(i64 8, ptr %104, ptr %105) - call void @llvm.stackrestore.p0(ptr %99) - ret { ptr, ptr, i64, [1 x i64], [1 x i64] } %98 -} - -define void @_mlir_ciface_pyop3_loop(ptr %0, ptr %1, ptr %2, ptr %3, ptr %4, ptr %5, ptr %6) { - %8 = load { ptr, ptr, i64, [1 x i64], [1 x i64] }, ptr %1, align 8 - %9 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %8, 0 - %10 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %8, 1 - %11 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %8, 2 - %12 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %8, 3, 0 - %13 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %8, 4, 0 - %14 = load { ptr, ptr, i64, [1 x i64], [1 x i64] }, ptr %2, align 8 - %15 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %14, 0 - %16 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %14, 1 - %17 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %14, 2 - %18 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %14, 3, 0 - %19 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %14, 4, 0 - %20 = load { ptr, ptr, i64, [1 x i64], [1 x i64] }, ptr %3, align 8 - %21 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %20, 0 - %22 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %20, 1 - %23 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %20, 2 - %24 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %20, 3, 0 - %25 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %20, 4, 0 - %26 = load { ptr, ptr, i64, [1 x i64], [1 x i64] }, ptr %4, align 8 - %27 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %26, 0 - %28 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %26, 1 - %29 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %26, 2 - %30 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %26, 3, 0 - %31 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %26, 4, 0 - %32 = load { ptr, ptr, i64, [1 x i64], [1 x i64] }, ptr %5, align 8 - %33 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %32, 0 - %34 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %32, 1 - %35 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %32, 2 - %36 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %32, 3, 0 - %37 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %32, 4, 0 - %38 = load { ptr, ptr, i64, [1 x i64], [1 x i64] }, ptr %6, align 8 - %39 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %38, 0 - %40 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %38, 1 - %41 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %38, 2 - %42 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %38, 3, 0 - %43 = extractvalue { ptr, ptr, i64, [1 x i64], [1 x i64] } %38, 4, 0 - %44 = call { ptr, ptr, i64, [1 x i64], [1 x i64] } @pyop3_loop(ptr %9, ptr %10, i64 %11, i64 %12, i64 %13, ptr %15, ptr %16, i64 %17, i64 %18, i64 %19, ptr %21, ptr %22, i64 %23, i64 %24, i64 %25, ptr %27, ptr %28, i64 %29, i64 %30, i64 %31, ptr %33, ptr %34, i64 %35, i64 %36, i64 %37, ptr %39, ptr %40, i64 %41, i64 %42, i64 %43) - store { ptr, ptr, i64, [1 x i64], [1 x i64] } %44, ptr %0, align 8 - ret void -} - -; Function Attrs: nocallback nofree nosync nounwind willreturn -declare ptr @llvm.stacksave.p0() #0 - -; Function Attrs: nocallback nofree nosync nounwind willreturn -declare void @llvm.stackrestore.p0(ptr) #0 - -attributes #0 = { nocallback nofree nosync nounwind willreturn } - -!llvm.module.flags = !{!0} - -!0 = !{i32 2, !"Debug Info Version", i32 3} diff --git a/mlir-testing/run_compiled_mlir.py b/mlir-testing/run_compiled_mlir.py deleted file mode 100644 index 6905b37c82..0000000000 --- a/mlir-testing/run_compiled_mlir.py +++ /dev/null @@ -1,52 +0,0 @@ -import ctypes -import numpy as np -import cupy as cp - -from ctypes import c_void_p, c_longlong, Structure - -class MemRefDescriptor(Structure): - _fields_ = [ - ("allocated", c_void_p), - ("aligned", c_void_p), - ("offset", c_longlong), - ("shape", c_longlong * 1), - ("stride", c_longlong * 1), - ] - -def numpy_to_memref(arr): - if not arr.flags["C_CONTIGUOUS"]: - arr = np.ascontiguousarray(arr) - - desc = MemRefDescriptor() - desc.allocated = arr.ctypes.data_as(c_void_p) - desc.aligned = desc.allocated - desc.offset = 0 - desc.shape[0] = arr.shape[0] - desc.stride[0] = 1 - - return desc - - -if __name__ == "__main__": - lib = ctypes.CDLL("./liboutput.dylib") - - array_add = lib._mlir_ciface_add - array_add.argtypes = [ - ctypes.POINTER(MemRefDescriptor) - ] * 3 - - size = 8 - a = np.ones(size, dtype=np.float64) - b = np.ones(size, dtype=np.float64) * 2 - c = np.zeros(size, dtype=np.float64) - - a_desc = numpy_to_memref(a) - b_desc = numpy_to_memref(b) - c_desc = numpy_to_memref(c) - - array_add(ctypes.byref(a_desc), ctypes.byref(b_desc), ctypes.byref(c_desc)) - - expected = a + b - np.testing.assert_array_almost_equal(c, expected) - print("Array addition successful!") - print(f"First few elements: {c[:5]}") diff --git a/mlir-testing/sample.py b/mlir-testing/sample.py deleted file mode 100644 index 1e8c018eda..0000000000 --- a/mlir-testing/sample.py +++ /dev/null @@ -1,98 +0,0 @@ -# Basic functionality -# Make MLIR in xDSL that adds two unranked tensor arrays -# Take this MLIR and lower to LLVM in xDSL -# Compile JIT with llvmlite? - -import sys -from xdsl.dialects import arith, func, memref, scf, tensor, linalg -from xdsl.dialects.arith import ConstantOp, AddfOp -from xdsl.dialects.tensor import DimOp, EmptyOp -from xdsl.dialects.builtin import ( - DYNAMIC_INDEX, - ModuleOp, - IndexType, - IntegerAttr, - f64, - TensorType, - ArrayAttr, - AffineMap, - AffineMapAttr, - AffineDimExpr, - UnitAttr -) - - -from xdsl.dialects.linalg.attrs import IteratorTypeAttr -from xdsl.dialects.linalg.ops import YieldOp, GenericOp - -from xdsl.ir import Block, Region -from xdsl.context import Context -from xdsl.printer import Printer - -def build_array_add(n: int) -> ModuleOp: - tensor_type = TensorType(f64, [DYNAMIC_INDEX]) - - identity_1d = AffineMap(num_dims=1, num_symbols=0, results=(AffineDimExpr(0),)) - identity_attr = AffineMapAttr(identity_1d) - - parallel = IteratorTypeAttr.parallel() - - func_block = Block(arg_types=[tensor_type, tensor_type, tensor_type]) - a, b, out = func_block.args - - c0 = ConstantOp(IntegerAttr(0, IndexType())) - func_block.add_op(c0) - - body_block = Block(arg_types=[f64, f64, f64]) - x, y, _z = body_block.args - - add = AddfOp(x, y) - body_block.add_op(add) - - body_block.add_op(YieldOp(add.result)) - - generic = GenericOp( - inputs=[a, b], - outputs=[out], - body=Region([body_block]), - indexing_maps=[identity_attr, identity_attr, identity_attr], - iterator_types=[parallel], - result_types=[tensor_type], - ) - func_block.add_op(generic) - - func_block.add_op(func.ReturnOp()) - - func_region = Region([func_block]) - func_op = func.FuncOp( - "add", - ([tensor_type, tensor_type, tensor_type], []), - func_region, - ) - - func_op.attributes["llvm.emit_c_interface"] = UnitAttr() - - return ModuleOp([func_op]) - -def emit_mlir(module: ModuleOp) -> str: - """Return the MLIR text representation of a module.""" - import io - buf = io.StringIO() - Printer(stream=buf).print_op(module) - return buf.getvalue() - -if __name__ == "__main__": - n = int(sys.argv[1]) if len(sys.argv) > 1 else 8 - - # Register dialects so xDSL can verify the IR - ctx = Context() - ctx.load_dialect(func.Func) - ctx.load_dialect(arith.Arith) - ctx.load_dialect(linalg.Linalg) - ctx.load_dialect(tensor.Tensor) - - module = build_array_add(n) - mlir = emit_mlir(module) - print(mlir) - - diff --git a/mlir-testing/sample_2.py b/mlir-testing/sample_2.py deleted file mode 100644 index 241d30a229..0000000000 --- a/mlir-testing/sample_2.py +++ /dev/null @@ -1,65 +0,0 @@ -from xdsl.builder import ImplicitBuilder -from xdsl.dialects import arith, func, scf, tensor -from xdsl.dialects.builtin import ( - FloatAttr, - IndexType, - ModuleOp, - TensorType, - f32, - i32, -) -from xdsl.ir import Block, Region - -index = IndexType() - -N = 128 -t_f = TensorType(f32, [N]) # dat_2, dat_3 (float data) -t_i = TensorType(i32, [N]) # idat_0, idat_2 (index data) - -fn_block = Block(arg_types=[t_f, t_f, t_i, t_i]) - -with ImplicitBuilder(fn_block) as (dat_2, dat_3, idat_0, idat_2): - c0 = arith.ConstantOp.from_int_and_width(0, index) - c1 = arith.ConstantOp.from_int_and_width(1, index) - n = arith.ConstantOp.from_int_and_width(N, index) - two = arith.ConstantOp(FloatAttr(2.0, f32)) - - body = Block(arg_types=[index, t_f]) - with ImplicitBuilder(body) as (i2, acc): - # k = idat_2[i_2] - k = tensor.ExtractOp(idat_2, [i2], i32) - k_idx = arith.IndexCastOp(k.result, index) - - # j = idat_0[k] - j = tensor.ExtractOp(idat_0, [k_idx.result], i32) - j_idx = arith.IndexCastOp(j.result, index) - - # v = dat_3[j] - v = tensor.ExtractOp(dat_3, [j_idx.result], f32) - - # r = 2.0 * v - r = arith.MulfOp(two.result, v.result) - - # dat_2[j] = r (value semantics -> produces new tensor) - new = tensor.InsertOp(r.result, acc, [j_idx.result]) - - scf.YieldOp(new.result) - - loop = scf.ForOp( - lb=c0.result, - ub=n.result, - step=c1.result, - iter_args=[dat_2], - body=Region(body), - ) - - func.ReturnOp(loop.results[0]) - -fn = func.FuncOp( - "indirect_scale", - ((t_f, t_f, t_i, t_i), (t_f,)), - Region(fn_block), -) - -module = ModuleOp([fn]) -print(module) From ea0996fe6a17eda778a40b2bb9cf0d520ad6babf Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Wed, 26 Aug 2026 14:13:47 +0100 Subject: [PATCH 16/30] removing offloading demo files from remote --- assign_offloading_demo.py | 26 -------------------------- indirect_offloading_demo.py | 15 --------------- 2 files changed, 41 deletions(-) delete mode 100644 assign_offloading_demo.py delete mode 100644 indirect_offloading_demo.py diff --git a/assign_offloading_demo.py b/assign_offloading_demo.py deleted file mode 100644 index ead9195f58..0000000000 --- a/assign_offloading_demo.py +++ /dev/null @@ -1,26 +0,0 @@ -from firedrake import * -import pyop3 as op3 -import numpy as np, cupy as cp - -import pyop3.debug_flags - -mesh = UnitSquareMesh(3,3) - -V = FunctionSpace(mesh, "CG", 1) -f = Function(V).assign(10) -g = Function(V) - -gpu = op3.CUDAGPU() - -pyop3.debug_flags.hit_assign = False -g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile", compiler_parameters={"codegen": "loopy"}) -pyop3.debug_flags.hit_assign = False - -# with op3.offloading(gpu): -# g.dat.assign(2 * f.dat, eager=True, eager_strategy="compile") -# assert isinstance(g.dat.data_ro, cp.ndarray) # Device -# assert (g.dat.data_ro == 20).all() - -assert isinstance(g.dat.data_ro, np.ndarray) # Host -assert (g.dat.data_ro == 20).all() - diff --git a/indirect_offloading_demo.py b/indirect_offloading_demo.py deleted file mode 100644 index 8b4e8e6148..0000000000 --- a/indirect_offloading_demo.py +++ /dev/null @@ -1,15 +0,0 @@ -from firedrake import * -import pyop3 as op3 -import numpy as np, cupy as cp - -import pyop3.debug_flags - -mesh = UnitSquareMesh(3,3) - -V = FunctionSpace(mesh, "CG", 1) -v = TestFunction(V) - -pyop3.debug_flags.hit_assign = True -b = assemble(conj(v) * dx) -pyop3.debug_flags.hit_assign = False - From 950ff8efaf43a416fd5a332af43ed4fa46f5f407 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Wed, 26 Aug 2026 14:17:35 +0100 Subject: [PATCH 17/30] removing mlir class for merge --- pyop3/lower/mlir.py | 406 -------------------------------------------- 1 file changed, 406 deletions(-) delete mode 100644 pyop3/lower/mlir.py diff --git a/pyop3/lower/mlir.py b/pyop3/lower/mlir.py deleted file mode 100644 index 6e3403c43b..0000000000 --- a/pyop3/lower/mlir.py +++ /dev/null @@ -1,406 +0,0 @@ -''' - SIGNIFICANT REWRITE OF THIS CLASS DUE TO FOLLOWING CHANGE: - FROM PYM->MLIR - TO PYOP3->MLIR -''' - -import contextlib -import functools -import numbers -import dataclasses -import numpy as np - -import pymbolic as pym -import loopy as lp # NOTE: For typing, temporary until fully separated - -from xdsl.dialects import arith, func, tensor, scf - -from xdsl.dialects.builtin import ( - DYNAMIC_INDEX, - ModuleOp, - IntegerAttr, - IntegerType, - IndexType, - FunctionType, - TensorType, - i32, - i64, - f64 -) - -from xdsl.dialects.func import ( - FuncOp -) - -from xdsl.builder import Builder, InsertPoint -from xdsl.ir import SSAValue, Block, Region - -import pyop3 -from pyop3.buffer import AbstractBuffer, ConcreteBuffer, PetscMatBuffer, ArrayBuffer, NullBuffer -from pyop3.dtypes import IntType - -from pyop3.lower.context import CodegenContext - -from pyop3.insn.base import ( - Intent -) -NUMPY_TO_XDSL = { - np.dtype(np.float64): f64, - np.dtype(np.int32): i32, - np.dtype(np.int64): i64, -} - -@dataclasses.dataclass -class Assignment: - assignee: object - expression: object - within_inames: object - id: str - - def __str__(self): - return f"{self.assignee} = {self.expression}" - - -class Argument: - def __init__(self, name, dtype, shape): - self.name = name - self.dtype = dtype - self.shape = shape - - def __str__(self): - return self.name - - def __repr__(self): - return f"<{self.name}, dtype: {self.dtype}, shape: {self.shape if self.shape else '?'}>" - -class SymbolTable: - ''' - Symbol Table that acts as lookup for pymbol to MLIR SSAValue - - Context manager works around MLIR's region and block based system - ''' - - def __init__(self): - self._scopes: list[dict[str, SSAValue]] = [{}] - - @contextlib.contextmanager - def scope(self): - self._scopes.append({}) - try: - yield self - finally: - self._scopes.pop() - - def insert(self, name: str, value: SSAValue): - self._scopes[-1][name] = value - - def lookup(self, name: str) -> SSAValue: - for sc in reversed(self._scopes): - if name in sc: - return sc[name] - raise KeyError(f"Unknown variable: {name}") - - def __str__(self): - return str(self._scopes) - -class MLIRCodegenContext(CodegenContext): - - def __init__(self, *, check_negatives): - super().__init__(check_negatives=check_negatives) - - self.symbol_table = SymbolTable() - - # NOTE: Temporary & unused while I rewrite lower/ - self._within_inames = frozenset() - self._domains = dict() - - def add_domain(self, iname, *args): - nargs = len(args) - if nargs == 1: - start, stop = 0, args[0] - else: - assert nargs == 2 - start, stop = args[0], args[1] - self._domains[iname] = (start, stop) - - def add_assignment(self, assignee, expression, prefix="insn"): - # Assignee and expression come in pymbolic expression - insn = Assignment( - assignee=assignee, - expression=expression, - within_inames=self._within_inames, - id=self.unique_name(prefix) - ) - self._instructions.append(insn) - - def add_function_call(self, assignees, expression, prefix="insn"): - pass - - def add_buffer(self, buffer, intent: Intent | None = None) -> str: - # TODO: This only works for np.ndarrays for development atm - - buffer_key = (buffer.name, buffer.nest_indices) - if isinstance(buffer, NullBuffer): - assert not buffer.nest_indices - - if buffer_key in self._kernel_names: - return self._kernel_names[buffer_key] - shape = self._temporary_shapes.get(buffer_key, (buffer.size,)) - assert isinstance(shape, tuple) and all(isinstance(s, numbers.Integral) for s in shape) - name_in_kernel = self.add_temporary("t", buffer.dtype, shape=shape) - else: - if intent is None: - raise ValueError("Global data must declare intent") - - if buffer_key in self._kernel_names: - if intent != self.global_buffer_intents[buffer_key]: - # We are accessing a buffer with different intents so have to - # pessimally claim RW access - self.global_buffer_intents[buffer_key] = RW - return self._kernel_names[buffer_key] - - if isinstance(buffer.handle, np.ndarray): - if isinstance(buffer.dtype, np.dtypes.IntDType): - name_in_kernel = self.unique_name("idat") - else: - name_in_kernel = self.unique_name("dat") - - # If the buffer is being passed straight through to a function then we - # have to make sure that the shapes match - shape = self._temporary_shapes.get(buffer_key, None) - # TODO: An equivalent of lp.GlobalArg is required here - # GlobalArg represents array, dtype, shape, address space (local or global variable) - iter_arg = Argument(name_in_kernel, buffer.dtype, shape) - else: - assert isinstance(buffer, PetscMatBuffer) - assert buffer.mat_type not in {"nest", "python"} - - name_in_kernel = self.unique_name("mat") - iter_arg = Argument(name_in_kernel, pyop3.dtypes.OpaqueType("mat")) - - self.global_buffers[buffer_key] = buffer - self.global_buffer_intents[buffer_key] = intent - self._arguments.append(iter_arg) - - self._kernel_names[buffer_key] = name_in_kernel - return name_in_kernel - - def add_subkernel(self, subkernel): - pass - - def add_instruction(self, insn): - # TODO: Ignoring CInstruction for now because no MLIR equivalent built - if isinstance(insn, lp.CInstruction): - raise ValueError("Cannot deal with Loopy CInstructions") - - self._instructions.append(insn) - self._last_insn_id = insn.id - - - # NOTE: Here while I port code, to be removed - @contextlib.contextmanager - def within_inames(self, inames) -> None: - orig_within_inames = self._within_inames - self._within_inames |= inames - yield - self._within_inames = orig_within_inames - # NOTE: Temporary while we work with basic kernels - # Without petsc mats or standalone functions, this is just empty idict - def set_temporary_shapes(self, shapes): - self._temporary_shapes = shapes - - - @functools.singledispatchmethod - def translate_expr(self, expr) -> SSAValue: - if isinstance(expr, tuple): - breakpoint() - raise ValueError(f"{type(expr)} not implemented yet.") - - - @translate_expr.register(pyop3.expr.Scalar) - def _(self, scalar: pyop3.expr.Scalar): - buffer_ref = scalar.buffer - name_in_kernel = context.add_buffer(buffer_ref) - return buffer_ref - - - @translate_expr.register(pym.primitives.Subscript) - def _(self, expr: pym.primitives.Subscript) -> SSAValue: - array, index_ssa = self._resolve_subscript(expr) - extract = tensor.ExtractOp.build( - operands=[array, index_ssa], - result_types=[array.type.element_type], - ) - self.builder.insert(extract) - return extract.result - - # NOTE: Could maybe clean up Sum and Product as they are essentially same, just reduction ops. - # TODO: Need to figure out how to solve the dtype inference. Fine to assume i32 for indexing but not for compute generally - @translate_expr.register(pym.primitives.Sum) - def _(self, expr): - children = [self.translate_expr(c) for c in expr.children] - result = children[0] - for child in children[1:]: - result = self._arith_op(arith.AddiOp, arith.AddfOp, result, child) - return result - - @translate_expr.register(pym.primitives.Product) - def _(self, expr): - children = [self.translate_expr(c) for c in expr.children] - result = children[0] - for child in children[1:]: - result = self._arith_op(arith.MuliOp, arith.MulfOp, result, child) - return result - - @translate_expr.register(pym.primitives.Variable) - def _(self, expr: pym.primitives.Variable) -> SSAValue: - return self.symbol_table.lookup(expr.name) - - @translate_expr.register(numbers.Number) - def _(self, expr: numbers.Number) -> SSAValue: - if isinstance(expr, int): - attr = IntegerAttr.from_int_and_width(expr, 32) - else: - attr = FloatAttr(float(expr), f64) - - const = arith.ConstantOp(attr) - self.builder.insert(const) - return const.result - - def _translate_assignment(self, ins): - - match ins.assignee: - case pym.primitives.Variable(): - value = self.translate_expr(ins.expression) - self.symbol_table.insert(ins.assignee.name, value) - - case pym.primitives.Subscript(): - value = self.translate_expr(ins.expression) - array, index_ssa = self._resolve_subscript(ins.assignee) - insert = tensor.InsertOp.build( - operands=[value, array, index_ssa], - result_types=[array.type], - ) - self.builder.insert(insert) - self.symbol_table.insert(ins.assignee.aggregate.name, insert.result) - - - def make_kernel(self): - # TODO: Update when moving away from loopy-style-defined argument variables - arg_types = [TensorType(NUMPY_TO_XDSL[arg.dtype], [DYNAMIC_INDEX]) for arg in self._arguments] - - func_op = FuncOp("pyop3_loop", FunctionType.from_lists(arg_types, [])) - - entry = func_op.body.blocks[0] - for arg, ssa in zip(self._arguments, entry.args): - self.symbol_table.insert(arg.name, ssa) - - self.builder = Builder(InsertPoint.at_end(entry)) - - self._build_nest(self._instructions, frozenset()) - - self.builder.insert(func.ReturnOp()) - self.module = ModuleOp([func_op]) - return self.module - - - def _build_nest(self, instructions, entered): - ''' Building nesting order to deal with instructions within loops ''' - - for ins in (i for i in instructions if i.within_inames == entered): - self._translate_assignment(ins) - - deeper = [i for i in instructions if i.within_inames != entered] - if not deeper: - return - - def next_iname(ins): - needed = ins.within_inames - entered - for iname in self._domains: - if iname in needed: - return iname - raise RuntimeError("inconsistent iname state") - - groups: dict[str, list] = {} - for ins in deeper: - groups.setdefault(next_iname(ins), []).append(ins) - - for iname in self._domains: - if iname in groups: - self._build_loop(iname, groups[iname], entered) - - - def _build_loop(self, iname, instructions, entered): - ''' Build scf for loops, this is where change would happen if we want to switch from scf ''' - - start, stop = self._domains[iname] - - lb = self._to_index(self.translate_expr(start)) - ub = self._to_index(self.translate_expr(stop)) - step = self._to_index(self.translate_expr(1)) - - carried = self._written_arrays(instructions) - init_values = [self.symbol_table.lookup(name) for name in carried] - - block_arg_types = [IndexType()] + [v.type for v in init_values] - body = Block(arg_types=block_arg_types) - - for_op = scf.ForOp(lb, ub, step, init_values, Region(body)) - self.builder.insert(for_op) - - with self.symbol_table.scope(): - self.symbol_table.insert(iname, body.args[0]) - # bind carried names to this loop's block args, not the outer values - for name, block_arg in zip(carried, body.args[1:]): - self.symbol_table.insert(name, block_arg) - - old = self.builder - self.builder = Builder(InsertPoint.at_end(body)) - self._build_nest(instructions, entered | {iname}) - - yielded = [self.symbol_table.lookup(name) for name in carried] - self.builder.insert(scf.YieldOp(*yielded)) - self.builder = old - - # send result/yield back to outer scope - for name, result in zip(carried, for_op.results): - self.symbol_table.insert(name, result) - - def _to_index(self, value): - if isinstance(value.type, IndexType): - return value - cast = arith.IndexCastOp.build(operands=[value], result_types=[IndexType()]) - self.builder.insert(cast) - return cast.result - - def _written_arrays(self, instructions): - ''' Finding all arrays that are written to within a loop for iterator arguments ''' - written = [] - seen = set() - for ins in instructions: - if isinstance(ins.assignee, pym.primitives.Subscript): - name = ins.assignee.aggregate.name - if name not in seen: - seen.add(name) - written.append(name) - return written - - def _resolve_subscript(self, subscript): - ''' Helper function for indices in tuples ''' - array = self.symbol_table.lookup(subscript.aggregate.name) - indices = subscript.index if isinstance(subscript.index, tuple) else (subscript.index,) - index_ssa = [self._to_index(self.translate_expr(i)) for i in indices] - return array, index_ssa - - def _arith_op(self, int_op, float_op, lhs, rhs): - ''' - Helper function to resolve typing between int and float ops - This is both super ugly and makes the assumption that type is determined by one side - ''' - t = lhs.type - if isinstance(t, (IntegerType, IndexType)): - op = int_op(lhs, rhs) - else: # float type - op = float_op(lhs, rhs) - self.builder.insert(op) - return op.result From 6090c48cfc5d4fb52bc81f767ff766a12a01bdbc Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Wed, 26 Aug 2026 14:40:57 +0100 Subject: [PATCH 18/30] adding docstrings --- pyop3/lower/codegen.py | 5 +++-- pyop3/lower/context.py | 37 +++++++++++++++++++------------------ 2 files changed, 22 insertions(+), 20 deletions(-) diff --git a/pyop3/lower/codegen.py b/pyop3/lower/codegen.py index 06d913cb67..43b644d87f 100644 --- a/pyop3/lower/codegen.py +++ b/pyop3/lower/codegen.py @@ -45,7 +45,7 @@ ) from pyop3.lower.loopy import LoopyCodegenContext -from pyop3.lower.context import _collect_temporary_shapes, _compile +from pyop3.lower.context import _collect_temporary_shapes, _compile # NOTE: Maybe these functions could both go in transform? # TODO: import other way around? from pyop3.lower.transform import ( @@ -65,6 +65,7 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed """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 ------- @@ -83,7 +84,7 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed if compiler_parameters.codegen == "loopy": ContextClass = LoopyCodegenContext elif compiler_parameters.codegen == "mlir": - raise NotImplementedError("Still implementing this class") + raise NotImplementedError("Class is still being implemented.") context = ContextClass( propagate_negatives=compiler_parameters.propagate_negatives, diff --git a/pyop3/lower/context.py b/pyop3/lower/context.py index 2c93b247e7..fc0f55e1dd 100644 --- a/pyop3/lower/context.py +++ b/pyop3/lower/context.py @@ -38,9 +38,9 @@ class CodegenContext(ABC): """ - Abstract base class for code generation contexts. + Base class for code generation backends - Abstract methods required for auto-generating based on _compile_static in codegen.py + 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: @@ -123,7 +123,6 @@ def add_subkernel(self, subkernel) -> None: def set_temporary_shapes(self, shapes) -> None: pass - ''' Lowering passes for respective codegen context ''' @abstractmethod def lower_expr(self, expr, iname_maps, loop_indices, intent: Intent | None = None, paths = None): @@ -146,6 +145,9 @@ def lower_buffer_access( *, intent ): + """ + Determine indexing and lower buffer expression to respective IR + """ pass @abstractmethod @@ -158,30 +160,23 @@ def register_extent(self, obj: Any, inames, loop_indices): @abstractmethod def compile_standalone_function(self, call, loop_indices): - """ - Compiling standalone functions i.e. LACallable for target IR representations - """ pass @abstractmethod def compile_petsc_mat(self, assignment, loop_indices): - """ - """ pass @abstractmethod def compile_exscan(self, call, loop_indices): - """ - """ pass # }}} - # {{{ general implementations + # {{{ class methods - def compile_array_assignment( + def _compile_array_assignment( self, assignment, loop_indices, @@ -241,7 +236,7 @@ def compile_array_assignment( with self.within_inames(within_inames): if axis_tree.node_map[new_paths[-1]]: - self.compile_array_assignment( + self._compile_array_assignment( assignment, loop_indices, axis_trees, @@ -250,7 +245,7 @@ def compile_array_assignment( paths=new_paths ) elif axis_trees: - self.compile_array_assignment( + self._compile_array_assignment( assignment, loop_indices, axis_trees, @@ -266,7 +261,7 @@ def compile_array_assignment( loop_indices ) - def parse_loop_properly_this_time( + def _parse_loop_properly_this_time( self, loop, axis_tree, @@ -313,7 +308,7 @@ def parse_loop_properly_this_time( with self.within_inames(within_inames): if subaxis := axis_tree.node_map[path_]: - self.parse_loop_properly_this_time( + self._parse_loop_properly_this_time( loop, axis_tree, loop_indices, @@ -341,6 +336,9 @@ 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" @@ -348,6 +346,9 @@ def __str__(self) -> str: return ctx +# NOTE: Not a big fan of how compile sits in this file. +# No need to make a class method in CodegenContext. +# Could maybe be a class in transform? Similar to collect_temporary_shapes @functools.singledispatch def _compile(expr: Any, loop_indices: Dict, codegen_context: CodegenContext) -> None: raise TypeError(f"No handler defined for {type(expr).__name__}") @@ -371,7 +372,7 @@ def _( loop_indices, codegen_context ) -> None: - codegen_context.parse_loop_properly_this_time( + codegen_context._parse_loop_properly_this_time( loop, loop.index.iterset, loop_indices, @@ -386,7 +387,7 @@ def parse_assignment(assignment: NonEmptyArrayAssignment, loop_indices, codegen_ if any(isinstance(arg, pyop3.expr.MatPetscMatBufferExpression) for arg in assignment.arguments): codegen_context.compile_petsc_mat(assignment, loop_indices) else: - codegen_context.compile_array_assignment( + codegen_context._compile_array_assignment( assignment, loop_indices, assignment.axis_trees, From bf476367c1e8e6d2f313f0069bae8d1a9458ff13 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Wed, 26 Aug 2026 15:40:21 +0100 Subject: [PATCH 19/30] removing dtype method for future pr --- pyop3/debug_flags.py | 1 - pyop3/expr/base.py | 43 ------------------------------------------- 2 files changed, 44 deletions(-) delete mode 100644 pyop3/debug_flags.py diff --git a/pyop3/debug_flags.py b/pyop3/debug_flags.py deleted file mode 100644 index 63c21bc953..0000000000 --- a/pyop3/debug_flags.py +++ /dev/null @@ -1 +0,0 @@ -hit_assign = False diff --git a/pyop3/expr/base.py b/pyop3/expr/base.py index 894e18b512..1b4c88608c 100644 --- a/pyop3/expr/base.py +++ b/pyop3/expr/base.py @@ -27,11 +27,6 @@ def local_max(self) -> numbers.Number: def local_min(self) -> numbers.Number: raise NotImplementedError - @property - # @abc.abstractmethod - def dtype(self) -> np.dtype: - pass - @property @abc.abstractmethod def _full_str(self) -> str: @@ -296,32 +291,6 @@ def get_disk_cache_key(self, visitor): def operands(self) -> tuple[ExpressionT, ExpressionT]: return (self.a, self.b) - @property - def dtype(self) -> np.dtype: - - if isinstance(self.a, Expression): - a_dtype = np.dtype(self.a.dtype) - else: - a_dtype = np.dtype(type(self.a)) - - if isinstance(self.b, Expression): - b_dtype = np.dtype(self.b.dtype) - else: - b_dtype = np.dtype(type(self.b)) - - is_a_float = np.issubdtype(a_dtype, np.floating) - is_b_float = np.issubdtype(b_dtype, np.floating) - - if is_a_float or is_b_float: - # Keep only float dtypes and pick the highest precision one - float_dtypes = [ - dt for dt, is_float in [(a_dtype, is_a_float), (b_dtype, is_b_float)] - if is_float - ] - return max(float_dtypes, key=lambda dt: dt.itemsize) - - return max(a_dtype, b_dtype, key=lambda dt: dt.itemsize) - def with_operands(self, operands): a, b = operands return self.record_new(a=a, b=b) @@ -661,10 +630,6 @@ def local_max(self) -> numbers.Number: def local_min(self) -> numbers.Number: raise TypeError("not sure that this makes sense") - @property - def dtype(self) -> np.dtype: - raise TypeError("Not sure this makes sense") - @property def _full_str(self) -> str: return f"i_{{{self.axis.label}}}" @@ -699,10 +664,6 @@ def local_max(self) -> NoReturn: def local_min(self) -> NoReturn: raise TypeError - @property - def dtype(self) -> np.dtype: - return None - _full_str = "NaN" # }}} @@ -755,10 +716,6 @@ def local_max(self) -> numbers.Number: def local_min(self) -> numbers.Number: raise TypeError("not sure that this makes sense") - @property - def dtype(self) -> np.dtype: - return TypeError("not sure that this makes sense") # possibly int32/64? - @property def _full_str(self) -> str: return f"L_{{{self.loop_index.id}, {self.axis.label}}}" From 6a18af4d9d2f62b7830749aed6739f9a3b17aa8d Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Wed, 26 Aug 2026 15:40:36 +0100 Subject: [PATCH 20/30] removing debug flags --- pyop3/insn/exec.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/pyop3/insn/exec.py b/pyop3/insn/exec.py index ed3210f6ec..287808bb33 100644 --- a/pyop3/insn/exec.py +++ b/pyop3/insn/exec.py @@ -28,8 +28,6 @@ from pyop3.cache import cached_method, memory_cache from pyop3.constants import INC, MAX_RW, MAX_WRITE, MIN_RW, MIN_WRITE, READ, RW, WRITE -import pyop3.debug_flags - @dataclasses.dataclass(frozen=True, kw_only=True) class CompilerParameters: @@ -231,7 +229,6 @@ 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 From e6b37d5104599b9d8dc4ea1e423c670716ef4ce7 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Wed, 26 Aug 2026 15:41:36 +0100 Subject: [PATCH 21/30] removing unnecessary requirements-build update --- requirements-build.txt | 16 ---------------- 1 file changed, 16 deletions(-) delete mode 100644 requirements-build.txt diff --git a/requirements-build.txt b/requirements-build.txt deleted file mode 100644 index 2d2cdf2044..0000000000 --- a/requirements-build.txt +++ /dev/null @@ -1,16 +0,0 @@ -# Core build dependencies (adapted from pyproject.toml) -Cython>=3.0 -firedrake-rtree>=2026.2.0 -libsupermesh>=2026.0 -mpi4py>3; python_version >= '3.13' -mpi4py; python_version < '3.13' -numpy -pkgconfig -petsctools @ git+https://github.com/firedrakeproject/petsctools.git@main -pybind11 -setuptools>=77.0.3 - -# Transitive build dependencies -hatchling -meson-python -scikit_build_core \ No newline at end of file From 0fcd377c13a9b9b61df4fff9dc650af497b66e79 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Wed, 26 Aug 2026 15:44:49 +0100 Subject: [PATCH 22/30] returning requirements.txt --- requirements-build.txt | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 requirements-build.txt diff --git a/requirements-build.txt b/requirements-build.txt new file mode 100644 index 0000000000..2d2cdf2044 --- /dev/null +++ b/requirements-build.txt @@ -0,0 +1,16 @@ +# Core build dependencies (adapted from pyproject.toml) +Cython>=3.0 +firedrake-rtree>=2026.2.0 +libsupermesh>=2026.0 +mpi4py>3; python_version >= '3.13' +mpi4py; python_version < '3.13' +numpy +pkgconfig +petsctools @ git+https://github.com/firedrakeproject/petsctools.git@main +pybind11 +setuptools>=77.0.3 + +# Transitive build dependencies +hatchling +meson-python +scikit_build_core \ No newline at end of file From f82365bdc36e978c89240aeda996a0573a5acaa8 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Thu, 27 Aug 2026 10:17:56 +0100 Subject: [PATCH 23/30] resolving simple PR comments - rename compiler option `codegen` -> `backend` - documenting `backend` compiler option - functional variable renaming --- pyop3/insn/exec.py | 3 ++- pyop3/lower/codegen.py | 8 ++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/pyop3/insn/exec.py b/pyop3/insn/exec.py index 287808bb33..583580598b 100644 --- a/pyop3/insn/exec.py +++ b/pyop3/insn/exec.py @@ -56,7 +56,8 @@ class CompilerParameters: # TODO: handle these - need to build CompilerOptions - codegen: str = "loopy" + backend: str = "loopy" + """ Option to select 'loopy' or 'mlir' as code generation backends. """ # extra_cflags: tuple[str, ...] = () # extra_ldflags: tuple[str, ...] = () diff --git a/pyop3/lower/codegen.py b/pyop3/lower/codegen.py index 43b644d87f..893d1c8e7e 100644 --- a/pyop3/lower/codegen.py +++ b/pyop3/lower/codegen.py @@ -81,12 +81,12 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed else: cs_expr = (insn,) - if compiler_parameters.codegen == "loopy": - ContextClass = LoopyCodegenContext - elif compiler_parameters.codegen == "mlir": + if compiler_parameters.backend == "loopy": + make_context = LoopyCodegenContext + elif compiler_parameters.backend == "mlir": raise NotImplementedError("Class is still being implemented.") - context = ContextClass( + context = make_context( propagate_negatives=compiler_parameters.propagate_negatives, mask_array_accesses=compiler_parameters.mask_array_accesses, ) From b45bbf719171600bde353423c29b370803d820ae Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Thu, 27 Aug 2026 11:03:40 +0100 Subject: [PATCH 24/30] moving _compile to codegen - moving dispatch _compile function from context.py -> codegen.py - local import in _parse_loop to avoid circular import - moving dispatch _collect_temporary_shapes from context.py -> transform.py --- pyop3/lower/codegen.py | 53 +++++++++++++++++++--- pyop3/lower/context.py | 97 +--------------------------------------- pyop3/lower/transform.py | 52 +++++++++++++++++++++ 3 files changed, 102 insertions(+), 100 deletions(-) diff --git a/pyop3/lower/codegen.py b/pyop3/lower/codegen.py index 893d1c8e7e..3103809e10 100644 --- a/pyop3/lower/codegen.py +++ b/pyop3/lower/codegen.py @@ -10,7 +10,6 @@ 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 @@ -33,25 +32,22 @@ 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 ( - AbstractAssignment, - AssignmentType, Exscan, InstructionList, Loop, NonEmptyArrayAssignment, NullInstruction, StandaloneCalledFunction, - assignment_type_as_intent, ) from pyop3.lower.loopy import LoopyCodegenContext -from pyop3.lower.context import _collect_temporary_shapes, _compile # NOTE: Maybe these functions could both go in transform? # TODO: import other way around? from pyop3.lower.transform import ( with_attach_debugger, with_likwid_markers, with_petsc_event, + _collect_temporary_shapes ) def _compile_static_hashkey(op: PreprocessedOperation, compiler_parameters: ParsedCompilerParameters) -> Hashable: @@ -128,3 +124,50 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed 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: + codegen_context._parse_loop_properly_this_time( + loop, + loop.index.iterset, + loop_indices, + ) + +@_compile.register(StandaloneCalledFunction) +def _(call, loop_indices, codegen_context): + codegen_context.compile_standalone_function(call, loop_indices) + +@_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): + codegen_context.compile_petsc_mat(assignment, loop_indices) + else: + codegen_context._compile_array_assignment( + assignment, + loop_indices, + assignment.axis_trees, + ) + +@_compile.register(Exscan) +def _(exscan, loop_indices, codegen_context): + codegen_context.compile_exscan(exscan, loop_indices) diff --git a/pyop3/lower/context.py b/pyop3/lower/context.py index fc0f55e1dd..2b94fb3e63 100644 --- a/pyop3/lower/context.py +++ b/pyop3/lower/context.py @@ -35,7 +35,6 @@ assignment_type_as_intent, ) - class CodegenContext(ABC): """ Base class for code generation backends @@ -271,6 +270,8 @@ def _parse_loop_properly_this_time( path=None, iname_map=None, ) -> None: + from pyop3.lower.codegen import _compile + if axis_tree is UNIT_AXIS_TREE: # NOTE: might need an expression here sometimes for statement in loop.statements: @@ -345,97 +346,3 @@ def __str__(self) -> str: ctx += f"Subkernels: {str(self.subkernels)}\n\n" return ctx - -# NOTE: Not a big fan of how compile sits in this file. -# No need to make a class method in CodegenContext. -# Could maybe be a class in transform? Similar to collect_temporary_shapes -@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: - codegen_context._parse_loop_properly_this_time( - loop, - loop.index.iterset, - loop_indices, - ) - -@_compile.register(StandaloneCalledFunction) -def _(call, loop_indices, codegen_context): - codegen_context.compile_standalone_function(call, loop_indices) - -@_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): - codegen_context.compile_petsc_mat(assignment, loop_indices) - else: - codegen_context._compile_array_assignment( - assignment, - loop_indices, - assignment.axis_trees, - ) - -@_compile.register(Exscan) -def _(exscan, loop_indices, codegen_context): - codegen_context.compile_exscan(exscan, loop_indices) - - -# 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): - import loopy as lp # TODO: Remove once StandaloneCalledFunction/similar integrated with MLIR - 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/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) + } + ) From 4a2d5abace086a7a13cebf3d4329e240c090fb89 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Thu, 27 Aug 2026 15:08:48 +0100 Subject: [PATCH 25/30] abstracting backend from pyop3 language - codegen holds more pyop3 traversal. - backend-specific functions introduced where necessary notable flaw is bloated function call for add_petsc_mat. --- pyop3/lower/codegen.py | 246 +++++++++++++++++++++++++++++++++++++++-- pyop3/lower/context.py | 194 ++++---------------------------- pyop3/lower/loopy.py | 164 ++++++++++----------------- 3 files changed, 318 insertions(+), 286 deletions(-) diff --git a/pyop3/lower/codegen.py b/pyop3/lower/codegen.py index 3103809e10..8d93639b8a 100644 --- a/pyop3/lower/codegen.py +++ b/pyop3/lower/codegen.py @@ -10,6 +10,7 @@ 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 @@ -44,9 +45,6 @@ # TODO: import other way around? from pyop3.lower.transform import ( - with_attach_debugger, - with_likwid_markers, - with_petsc_event, _collect_temporary_shapes ) @@ -147,27 +145,259 @@ def _( loop_indices, codegen_context ) -> None: - codegen_context._parse_loop_properly_this_time( + _compile_loop( loop, loop.index.iterset, loop_indices, + codegen_context, ) @_compile.register(StandaloneCalledFunction) def _(call, loop_indices, codegen_context): - codegen_context.compile_standalone_function(call, loop_indices) + 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): - codegen_context.compile_petsc_mat(assignment, loop_indices) + _compile_petsc_mat(assignment, loop_indices, codegen_context) else: - codegen_context._compile_array_assignment( + _compile_array_assignment( assignment, loop_indices, assignment.axis_trees, + codegen_context, ) +def _compile_petsc_mat( + assignment, + loop_indices, + codegen_context, +): + # 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 + + assert isinstance(expr, pyop3.expr.BufferExpression) + + # 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 + + codegen_context.add_petsc_mat( + mat.buffer_view, + expr.buffer_view, + setting_mat_values, + assignment.assignment_type, + assignment.axis_trees, + layout_exprs, + loop_indices + ) + + +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, + 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, + 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, + ) + +# NOTE: This could become backend-agnostic if I implement an backend-alternative to pym.substitute @_compile.register(Exscan) def _(exscan, loop_indices, codegen_context): - codegen_context.compile_exscan(exscan, loop_indices) + 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) + + codegen_context.add_exscan( + lexpr, + rexpr, + iname, + iname_var + ) diff --git a/pyop3/lower/context.py b/pyop3/lower/context.py index 2b94fb3e63..0dbe32ec3a 100644 --- a/pyop3/lower/context.py +++ b/pyop3/lower/context.py @@ -118,6 +118,29 @@ def add_buffer(self, buffer_view: IndexedBuffer, intent: Intent | None = None) - def add_subkernel(self, subkernel) -> None: pass + @abstractmethod + def add_exscan( + self, + lexpr, + rexpr, + iname, + iname_var + ) -> None: + pass + + @abstractmethod + def add_petsc_mat( + self, + mat_view, + expr_view, + setting_mat_values, + assignment_type, + axis_trees, + layout_exprs, + loop_indices + ) -> None: + pass + @abstractmethod def set_temporary_shapes(self, shapes) -> None: pass @@ -157,179 +180,8 @@ def add_leaf_assignment(self, assignment, paths, iname_maps, loop_indices): def register_extent(self, obj: Any, inames, loop_indices): pass - @abstractmethod - def compile_standalone_function(self, call, loop_indices): - pass - - @abstractmethod - def compile_petsc_mat(self, assignment, loop_indices): - pass - - @abstractmethod - def compile_exscan(self, call, loop_indices): - pass - # }}} - - # {{{ class methods - - - def _compile_array_assignment( - self, - assignment, - loop_indices, - 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("Refactor needed") - - self.add_leaf_assignment( - assignment, - 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 = self.unique_name("i") - ext = self.register_extent( - component.size, - iname_replace_maps[-1], - loop_indices - ) - self.add_domain(iname, ext) - new_maps = iname_replace_maps.copy() - new_maps[-1] = iname_replace_maps[-1] | {axis.label: self.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 self.within_inames(within_inames): - if axis_tree.node_map[new_paths[-1]]: - self._compile_array_assignment( - assignment, - loop_indices, - axis_trees, - iname_replace_maps=new_maps, - axis_tree=axis_tree, - paths=new_paths - ) - elif axis_trees: - self._compile_array_assignment( - assignment, - loop_indices, - axis_trees, - iname_replace_maps=new_maps, - axis_tree=None, - paths=new_paths - ) - else: - self.add_leaf_assignment( - assignment, - new_paths, - new_maps, - loop_indices - ) - - def _parse_loop_properly_this_time( - self, - loop, - axis_tree, - loop_indices, - *, - axis=None, - path=None, - iname_map=None, - ) -> None: - from pyop3.lower.codegen import _compile - - 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, - self, - ) - 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 = self.unique_name("i") - domain_var = self.register_extent( - component.size, - iname_map, - loop_indices - ) - self.add_domain(iname, domain_var) - iname_replace_map_ = iname_map | {axis.label: self.var(iname)} - within_inames = frozenset({iname}) - else: - iname_replace_map_ = iname_map | {axis.label: 0} - within_inames = set() - - with self.within_inames(within_inames): - if subaxis := axis_tree.node_map[path_]: - self._parse_loop_properly_this_time( - loop, - axis_tree, - loop_indices, - 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, - self, - ) - # }}} - def add_subkernel(self, subkernel): self._subkernels.append(subkernel) diff --git a/pyop3/lower/loopy.py b/pyop3/lower/loopy.py index 37ec6ffa9c..dcf817aa0f 100644 --- a/pyop3/lower/loopy.py +++ b/pyop3/lower/loopy.py @@ -123,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, @@ -333,73 +366,22 @@ def maybe_multiindex(self, buffer, offset_expr): return indices - def compile_standalone_function( - self, - call: StandaloneCalledFunction, - loop_indices - ) -> 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 = 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 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} - ), - ) - - self.add_function_call(assignees, expression) - subkernel = call.function.code.with_entrypoints(frozenset()) - self.add_subkernel(subkernel) - - def compile_petsc_mat( + def add_petsc_mat( self, - assignment: ConcretizedNonEmptyArrayAssignment, + mat_view, + expr_view, + setting_mat_values, # bool flag for if updating or reading matrix + assignment_type, + axis_trees, + layout_exprs, loop_indices ) -> 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 + row_axis_tree, column_axis_tree = axis_trees + # Emit the right line of code, this should properly be a lp.ScalarCallable # https://petsc.org/release/manualpages/Mat/MatGetValuesLocal/ - mat_name = self.add_buffer(mat.buffer_view, assignment_type_as_intent(assignment.assignment_type)) - array_name = self.add_buffer(expr.buffer_view, READ) + mat_name = self.add_buffer(mat_view, assignment_type_as_intent(assignment_type)) + array_name = self.add_buffer(expr_view, READ) rsize = row_axis_tree.size csize = column_axis_tree.size @@ -417,34 +399,17 @@ def compile_petsc_mat( 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 = self.lower_expr(subst_layout, ((),), loop_indices) - layout_exprs.append(layout_expr) - irow, icol = layout_exprs - # FIXME: blocked = False + irow, icol = layout_exprs + # hacky myargs = [ - assignment, mat_name, array_name, rsize_var, csize_var, irow, icol, blocked + mat_name, array_name, rsize_var, csize_var, irow, icol, blocked ] if setting_mat_values: - match assignment.assignment_type: + match assignment_type: case AssignmentType.WRITE: call_str = _petsc_mat_store(*myargs) case AssignmentType.INC: @@ -500,28 +465,13 @@ def add_leaf_assignment( self.add_assignment(lexpr, rexpr) - def compile_exscan( + def add_exscan( self, - exscan: Exscan, - loop_indices + lexpr, + rexpr, + iname, + iname_var ) -> None: - assert isinstance(exscan, Exscan) - - if exscan.scan_type != "+": - raise NotImplementedError - domain_var = self.register_extent( - exscan.extent, - {}, - loop_indices, - ) - iname = self.unique_name("i") - self.add_domain(iname, domain_var) - - iname_var = pym.var(iname) - iname_map = {exscan.scan_axis.label: pym.var(iname)} - - lexpr = self.lower_expr(exscan.assignee, [iname_map], loop_indices, intent=RW) - rexpr = lexpr + self.lower_expr(exscan.expression, [iname_map], loop_indices) lexpr = pym.substitute(lexpr, {iname: iname_var+1}) self.add_assignment(lexpr, rexpr) @@ -850,21 +800,21 @@ def generate_preambles(self, target): assert isinstance(target, type(target)) yield ("solve", solve_preamble) -def _petsc_mat_load(assignment, mat_name, array_name, nrow, ncol, irow, icol, blocked): +def _petsc_mat_load(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): +def _petsc_mat_store(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): +def _petsc_mat_add(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: From af7d47ea9d67cbdaf3360a35bb378935583ad7bc Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Thu, 27 Aug 2026 15:39:27 +0100 Subject: [PATCH 26/30] moving petsc_mat and exscan to codegen - not implementing methods for MLIR, loopy-only. --- pyop3/lower/codegen.py | 100 ++++++++++++++++++++++++++++++++--------- pyop3/lower/context.py | 33 ++++---------- pyop3/lower/loopy.py | 99 ++++------------------------------------ 3 files changed, 97 insertions(+), 135 deletions(-) diff --git a/pyop3/lower/codegen.py b/pyop3/lower/codegen.py index 8d93639b8a..b0c3d93b75 100644 --- a/pyop3/lower/codegen.py +++ b/pyop3/lower/codegen.py @@ -33,12 +33,14 @@ 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 @@ -176,6 +178,8 @@ def _compile_petsc_mat( 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): @@ -186,9 +190,32 @@ def _compile_petsc_mat( 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: # @@ -208,15 +235,25 @@ def _compile_petsc_mat( layout_exprs.append(layout_expr) irow, icol = layout_exprs - codegen_context.add_petsc_mat( - mat.buffer_view, - expr.buffer_view, - setting_mat_values, - assignment.assignment_type, - assignment.axis_trees, - layout_exprs, - loop_indices - ) + # 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( @@ -246,7 +283,9 @@ def _compile_array_assignment( raise NotImplementedError("Refactor needed") codegen_context.add_leaf_assignment( - assignment, + assignment.assignee, + assignment.expression, + assignment.assignment_type, paths, iname_replace_maps, loop_indices @@ -300,7 +339,9 @@ def _compile_array_assignment( ) else: codegen_context.add_leaf_assignment( - assignment, + assignment.assignee, + assignment.expression, + assignment.assignment_type, new_paths, new_maps, loop_indices @@ -374,9 +415,11 @@ def _compile_loop( codegen_context, ) -# NOTE: This could become backend-agnostic if I implement an backend-alternative to pym.substitute @_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 @@ -395,9 +438,26 @@ def _(exscan, loop_indices, codegen_context): 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) - codegen_context.add_exscan( - lexpr, - rexpr, - iname, - iname_var - ) + 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 index 0dbe32ec3a..45b4b07504 100644 --- a/pyop3/lower/context.py +++ b/pyop3/lower/context.py @@ -118,29 +118,6 @@ def add_buffer(self, buffer_view: IndexedBuffer, intent: Intent | None = None) - def add_subkernel(self, subkernel) -> None: pass - @abstractmethod - def add_exscan( - self, - lexpr, - rexpr, - iname, - iname_var - ) -> None: - pass - - @abstractmethod - def add_petsc_mat( - self, - mat_view, - expr_view, - setting_mat_values, - assignment_type, - axis_trees, - layout_exprs, - loop_indices - ) -> None: - pass - @abstractmethod def set_temporary_shapes(self, shapes) -> None: pass @@ -173,7 +150,15 @@ def lower_buffer_access( pass @abstractmethod - def add_leaf_assignment(self, assignment, paths, iname_maps, loop_indices): + def add_leaf_assignment( + self, + assignee, + expression, + assignment_type, + paths, + iname_maps, + loop_indices + ): pass @abstractmethod diff --git a/pyop3/lower/loopy.py b/pyop3/lower/loopy.py index dcf817aa0f..bd74bde6bd 100644 --- a/pyop3/lower/loopy.py +++ b/pyop3/lower/loopy.py @@ -366,84 +366,32 @@ def maybe_multiindex(self, buffer, offset_expr): return indices - def add_petsc_mat( - self, - mat_view, - expr_view, - setting_mat_values, # bool flag for if updating or reading matrix - assignment_type, - axis_trees, - layout_exprs, - loop_indices - ) -> None: - - row_axis_tree, column_axis_tree = axis_trees - # Emit the right line of code, this should properly be a lp.ScalarCallable - # https://petsc.org/release/manualpages/Mat/MatGetValuesLocal/ - mat_name = self.add_buffer(mat_view, assignment_type_as_intent(assignment_type)) - array_name = self.add_buffer(expr_view, READ) - - rsize = row_axis_tree.size - csize = column_axis_tree.size - - # these sizes can be expressions that need evaluating - rsize_var = self.register_extent( - rsize, - {}, - loop_indices, - ) - - csize_var = self.register_extent( - csize, - {}, - loop_indices, - ) - - # FIXME: - blocked = False - - irow, icol = layout_exprs - - # hacky - myargs = [ - mat_name, array_name, rsize_var, csize_var, irow, icol, blocked - ] - if setting_mat_values: - match 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) - - self.add_cinstruction(call_str) - + # NOTE: This could probably be refactored def add_leaf_assignment( self, - assignment, + assignee, + expression, + assignment_type, paths, iname_replace_maps, loop_indices, ): - intent = assignment_type_as_intent(assignment.assignment_type) + intent = assignment_type_as_intent(assignment_type) lexpr = self.lower_expr( - assignment.assignee, + assignee, iname_replace_maps, loop_indices, intent=intent, paths=paths, ) rexpr = self.lower_expr( - assignment.expression, + expression, iname_replace_maps, loop_indices, paths=paths, ) - match assignment.assignment_type: + match assignment_type: case AssignmentType.WRITE: pass case AssignmentType.INC: @@ -465,16 +413,6 @@ def add_leaf_assignment( self.add_assignment(lexpr, rexpr) - def add_exscan( - self, - lexpr, - rexpr, - iname, - iname_var - ) -> None: - lexpr = pym.substitute(lexpr, {iname: iname_var+1}) - 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) @@ -532,7 +470,6 @@ def _(self, cond, /, *args, **kwargs) -> pym.Expression: 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] @@ -799,23 +736,3 @@ class SolveCallable(LACallable): def generate_preambles(self, target): assert isinstance(target, type(target)) yield ("solve", solve_preamble) - -def _petsc_mat_load(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(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(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);" From 5dba8f5b4db536d59b24b2dfbf555599831374a0 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Fri, 28 Aug 2026 10:31:37 +0100 Subject: [PATCH 27/30] fix typo and move transform as title --- pyop3/lower/codegen.py | 39 +++++++++++++++++++++++++++++++++++++++ pyop3/lower/loopy.py | 2 +- pyop3/lower/transform.py | 40 ---------------------------------------- 3 files changed, 40 insertions(+), 41 deletions(-) diff --git a/pyop3/lower/codegen.py b/pyop3/lower/codegen.py index b0c3d93b75..13ed3af867 100644 --- a/pyop3/lower/codegen.py +++ b/pyop3/lower/codegen.py @@ -441,6 +441,45 @@ def _(exscan, loop_indices, codegen_context): lexpr = pym.substitute(lexpr, {iname: iname_var+1}) codegen_context.add_assignment(lexpr, rexpr) +# NOTE: Make this overloaded function into class? +@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) + } + ) def _petsc_mat_load(assignment, mat_name, array_name, nrow, ncol, irow, icol, blocked): if blocked: diff --git a/pyop3/lower/loopy.py b/pyop3/lower/loopy.py index bd74bde6bd..b507a5d132 100644 --- a/pyop3/lower/loopy.py +++ b/pyop3/lower/loopy.py @@ -315,7 +315,7 @@ def lower_buffer_access( buffer = buffer_view.buffer if isinstance(buffer, PetscMatBuffer): - buffer = buffer_view.denested.getPythonself().buffer + 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 diff --git a/pyop3/lower/transform.py b/pyop3/lower/transform.py index ccfd832404..f533c55665 100644 --- a/pyop3/lower/transform.py +++ b/pyop3/lower/transform.py @@ -74,43 +74,3 @@ def with_attach_debugger(kernel): ) 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) - } - ) From ef3a195bd1ec25f056fcf78b95d3c86d0a0acc0c Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Fri, 28 Aug 2026 10:47:54 +0100 Subject: [PATCH 28/30] restructuring lower_expr - moving lower_expr to free function - going to revisit this once I have a working MLIR lowering implementation --- pyop3/lower/codegen.py | 6 +- pyop3/lower/loopy.py | 249 +++++++++++++++++++++-------------------- 2 files changed, 126 insertions(+), 129 deletions(-) diff --git a/pyop3/lower/codegen.py b/pyop3/lower/codegen.py index 13ed3af867..8a5f1dfc5f 100644 --- a/pyop3/lower/codegen.py +++ b/pyop3/lower/codegen.py @@ -33,6 +33,7 @@ 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 ( + AbstractAssignment, AssignmentType, Exscan, InstructionList, @@ -45,11 +46,6 @@ 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) diff --git a/pyop3/lower/loopy.py b/pyop3/lower/loopy.py index b507a5d132..27bab00a3a 100644 --- a/pyop3/lower/loopy.py +++ b/pyop3/lower/loopy.py @@ -414,130 +414,7 @@ def add_leaf_assignment( 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, - ) + return _lower_expr(expr, iname_maps, loop_indices, intent=intent, paths=paths, context=self) def finalize_kernel(self, function_name, compiler_parameters): preambles = [ @@ -595,6 +472,130 @@ def _(self, expr: pyop3.expr.Expression, inames, loop_indices): self.add_assignment(pym.var(extent_name), pym_expr) return extent_name + +# 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.Mod, /, *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, + *, + intent, + context, + **kwargs, +) -> pym.ExpressionNode: + return context.lower_buffer_access(expr.buffer_view, [0], iname_maps, loop_indices, intent=intent) + + +@_lower_expr.register(pyop3.expr.LinearDatBufferExpression) +def _(expr: pyop3.expr.LinearDatBufferExpression, /, iname_maps, loop_indices, *, intent, context, **kwargs) -> pym.Expression: + return context.lower_buffer_access(expr.buffer_view, [expr.layout], iname_maps, loop_indices, intent=intent) + + +@_lower_expr.register(pyop3.expr.NonlinearDatBufferExpression) +def _(expr: pyop3.expr.NonlinearDatBufferExpression, /, iname_maps, loop_indices, *, intent, paths, context, **kwargs) -> pym.Expression: + path = utils.just_one(paths) + return context.lower_buffer_access( + expr.buffer_view, + [expr.layouts[path]], + iname_maps, + loop_indices, + intent=intent, + ) + + +@_lower_expr.register(pyop3.expr.MatPetscMatBufferExpression) +def _(mat_expr: pyop3.expr.MatPetscMatBufferExpression, /, iname_maps, loop_indices, *, intent, paths, context) -> pym.Expression: + row_path, column_path = paths + layouts = ( + mat_expr.row_layout.linearize(row_path), + mat_expr.column_layout.linearize(column_path), + ) + return context.lower_buffer_access( + mat_expr.buffer_view, + layouts, + iname_maps, + loop_indices, + intent=intent, + ) + + +@_lower_expr.register(pyop3.expr.MatArrayBufferExpression) +def _(expr: pyop3.expr.MatArrayBufferExpression, /, iname_maps, loop_indices, *, intent, paths, context) -> pym.Expression: + row_path, column_path = paths + layouts = (expr.row_layouts[row_path], expr.column_layouts[column_path]) + return context.lower_buffer_access( + expr.buffer_view, + layouts, + iname_maps, + loop_indices, + intent=intent, + ) + class _MinSubscriptOffsetMapper(pym.mapper.IdentityMapper): def __init__(self): From 694d7a1518da4ad468b8403683e9e607e408a14f Mon Sep 17 00:00:00 2001 From: SamSJackson <86316114+SamSJackson@users.noreply.github.com> Date: Fri, 28 Aug 2026 11:40:16 +0100 Subject: [PATCH 29/30] Removing unnecessary imports Co-authored-by: Connor Ward --- pyop3/lower/transform.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/pyop3/lower/transform.py b/pyop3/lower/transform.py index f533c55665..5842a45d06 100644 --- a/pyop3/lower/transform.py +++ b/pyop3/lower/transform.py @@ -1,18 +1,6 @@ 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): """ From 982632b448560dbcd5eeeada33d7953f34474a86 Mon Sep 17 00:00:00 2001 From: SamSJackson Date: Fri, 28 Aug 2026 15:07:05 +0100 Subject: [PATCH 30/30] cleaning and fixed bug with component.local size --- pyop3/lower/codegen.py | 4 ++-- pyop3/lower/context.py | 2 -- pyop3/lower/loopy.py | 12 +++++++----- 3 files changed, 9 insertions(+), 9 deletions(-) diff --git a/pyop3/lower/codegen.py b/pyop3/lower/codegen.py index 8a5f1dfc5f..bdae968f8c 100644 --- a/pyop3/lower/codegen.py +++ b/pyop3/lower/codegen.py @@ -76,7 +76,7 @@ def _compile_static(op: InstructionExecutionContext, compiler_parameters: Parsed if compiler_parameters.backend == "loopy": make_context = LoopyCodegenContext elif compiler_parameters.backend == "mlir": - raise NotImplementedError("Class is still being implemented.") + raise NotImplementedError("MLIR code generation is still being implemented.") context = make_context( propagate_negatives=compiler_parameters.propagate_negatives, @@ -296,7 +296,7 @@ def _compile_array_assignment( if axis_tree.linearize(new_paths[-1], partial=True).size == 0: continue - if component.local_size != 1: + elif component.size != 1: iname = codegen_context.unique_name("i") ext = codegen_context.register_extent( component.size, diff --git a/pyop3/lower/context.py b/pyop3/lower/context.py index 45b4b07504..800eec7b3b 100644 --- a/pyop3/lower/context.py +++ b/pyop3/lower/context.py @@ -89,7 +89,6 @@ def _add_instruction(self, insn): # {{{ abstract methods - @abstractmethod def var(self, iname: str, *args) -> str | pym.primitives.Variable: """ @@ -97,7 +96,6 @@ def var(self, iname: str, *args) -> str | pym.primitives.Variable: """ pass - @abstractmethod def add_domain(self, iname: str, *args) -> None: pass diff --git a/pyop3/lower/loopy.py b/pyop3/lower/loopy.py index 27bab00a3a..5d326c805b 100644 --- a/pyop3/lower/loopy.py +++ b/pyop3/lower/loopy.py @@ -478,12 +478,10 @@ def _(self, expr: pyop3.expr.Expression, inames, loop_indices): 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) @@ -511,12 +509,12 @@ def _(or_: pyop3.expr.Or, /, *args, **kwargs) -> pym.Expression: @_lower_expr.register(pyop3.expr.Neg) def _(neg: pyop3.expr.Neg, /, *args, **kwargs) -> pym.Expression: - return _lower_expr(neg.a, *args, **kwargs) + 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) +def _(fdiv: pyop3.expr.FloorDiv, /, *args, **kwargs) -> pym.Expression: + return _lower_expr(fdiv.a, *args, **kwargs) // _lower_expr(fdiv.b, *args, **kwargs) @_lower_expr.register(pyop3.expr.Comparison) @@ -527,6 +525,10 @@ def _(cond, /, *args, **kwargs) -> pym.Expression: _lower_expr(cond.b, *args, **kwargs), ) +@_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)) + @_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]