b2b PR

devito/ir/ietxdsl/cluster_to_ssa.py

@@ -298,4 +298,169 @@ def is_int(val: SSAValue):
     return isinstance(val.typ, builtin.IntegerType)
 
 def is_float(val: SSAValue):
-    return val.typ in (builtin.f32, builtin.f64)
+    return val.typ in (builtin.f32, builtin.f64)
+
+
+
+#### -------------------------------------------------------- ####
+####                                                          ####
+####           devito.stencil  ---> stencil dialect           ####
+####                                                          ####
+#### -------------------------------------------------------- ####
+
+from xdsl.pattern_rewriter import RewritePattern, PatternRewriter, GreedyRewritePatternApplier, op_type_rewrite_pattern, PatternRewriteWalker
+
+from devito.ir.ietxdsl.lowering import recalc_func_type
+
+class _DevitoStencilToStencilStencil(RewritePattern):
+    """
+    This converts the `devito.stencil` op into the following:
+
+    ```
+        // we do magic here to accomodate triple buffering
+        // this will be tackled in the lowering
+        %t0_field = iet.get_stencil(%t0)
+        %t1_field = iet.get_stencil(%t1)
+        %t2_field = iet.get_stencil(%t2)
+
+        // do cast and load ops
+        %t0_w_size = stencil.cast(%t0_field) [-4, -4, -4] to [68, 68, 68]
+        %t1_w_size = stencil.cast(%t1_field) [-4, -4, -4] to [68, 68, 68]
+        %t2_w_size = stencil.cast(%t1_field) [-4, -4, -4] to [68, 68, 68]
+        // put ub - lb here (make the temp ?x?x?)
+        %t0_temp = stencil.load(%t0_w_size) : (!stencil.field<72x72x72xf64>) -> !stencil.temp<?x?x?xf64>
+        %t1_temp = stencil.load(%t1_w_size) : (!stencil.field<72x72x72xf64>) -> !stencil.temp<?x?x?xf64>
+
+        // apply with t0 and t1 data to generate t2 data
+        %t2_output = stencil.apply(%t0_temp, %t1_temp) ({
+        ^1(%t0_buff : !stencil.temp<?x?x?xf64>, %t1_buff : !stencil.temp<?x?x?xf64>):
+            ...
+            // stencil.body is placed here
+        }) : (!stencil.temp<?x?x?xf64>, !stencil.temp<?x?x?xf64>) -> !stencil.temp<?x?x?xf64>
+
+        // write generated temp back to field
+        // what are the bound notations here?
+        stencil.store(%t2_output, %t2_w_size) [0, 0, 0] : [64, 64, 64]
+    ```
+    """
+    @op_type_rewrite_pattern
+    def match_and_rewrite(self, op: iet_ssa.Stencil, rewriter: PatternRewriter, /):
+        field_t = stencil.FieldType.from_shape(
+            op.shape, op.field_type
+        )
+        rank = len(op.shape.data)
+
+        data = iet_ssa.LoadSymbolic.get('data', memref.MemRefType.from_element_type_and_shape(
+            stencil.FieldType.from_shape([-1] * rank, op.field_type), [op.time_buffers.value.data]
+        ))
+        lb = stencil.IndexAttr.get(*list(-halo_elm.data[0].value.data for halo_elm in op.halo.data))
+        ub = stencil.IndexAttr.get(*list(shape_elm.value.data + halo_elm.data[1].value.data for shape_elm, halo_elm in zip(op.shape.data, op.halo.data)))
+        fields = list(
+            iet_ssa.GetField.get(data, t_idx, field_t, lb, ub)
+            for t_idx in (*op.input_indices, op.output)
+        )
+        # name the resulting fields
+        for field in fields:
+            field.field.name = f"{field.t_index.name}_w_size"
+
+        loads = list(
+            stencil.LoadOp.get(field)
+            for field in fields[:-1]
+        )
+        rewriter.replace_matched_op([
+            iet_ssa.Statement.get("// get data obj"),
+            data,
+            iet_ssa.Statement.get("// get individual fields"),
+            *fields,
+            iet_ssa.Statement.get("// stencil loads"),
+            *loads,
+            out := stencil.ApplyOp.get(
+                loads, op.body.detach_block(0)
+            ),
+            stencil.StoreOp.get(out, fields[-1])
+        ])
+
+class _LowerGetField(RewritePattern):
+    @op_type_rewrite_pattern
+    def match_and_rewrite(self, op: iet_ssa.GetField, rewriter: PatternRewriter, /):
+
+        rewriter.replace_matched_op([
+            idx := arith.IndexCastOp.get(op.t_index, builtin.IndexType()),
+            field := memref.Load.get(op.data, [idx]),
+            field_w_size := stencil.CastOp.get(field, op.lb, op.ub, op.field.typ)
+        ], [field_w_size.result])
+
+
+def get_containing_func(op: Operation) -> func.FuncOp | None:
+    while op is not None and not isinstance(op, func.FuncOp):
+        op = op.parent_op()
+    if op is None:
+        return None
+    return op
+
+class _LowerLoadSymbolidToFuncArgs(RewritePattern):
+    func_to_args: dict[func.FuncOp, dict[str, SSAValue]]
+
+    def __init__(self):
+        from collections import defaultdict
+        self.func_to_args = defaultdict(dict)
+
+    @op_type_rewrite_pattern
+    def match_and_rewrite(self, op: iet_ssa.LoadSymbolic, rewriter: PatternRewriter, /):
+        parent: func.FuncOp | None = get_containing_func(op)
+        assert parent is not None        
+        args = self.func_to_args[func]
+        symb_name = op.symbol_name.data
+
+        if symb_name not in args:
+            body = parent.body.blocks[0]
+            args[symb_name] = body.insert_arg(op.result.typ, 0)
+
+        op.result.replace_by(args[symb_name])
+        rewriter.erase_matched_op()
+        recalc_func_type(parent)
+        # attach information on parameter names to func
+        parent.attributes['param_names'] = builtin.ArrayAttr([
+            builtin.StringAttr(name) for name, _ in sorted(args.items(), key=lambda tpl: tpl[1].index)
+        ])
+
+
+
+
+
+
+
+
+
+if __name__ == '__main__':
+    import sys
+    from xdsl.parser import Parser, MLContext, Source
+    from devito.ir.ietxdsl.lowering import DropIetComments, LowerIetForToScfFor, ConvertForLoopVarToIndex, ConvertScfForArgsToIndex
+    ctx = MLContext()
+    ctx.register_dialect(stencil.Stencil)
+    ctx.register_dialect(builtin.Builtin)
+    ctx.register_dialect(arith.Arith)
+    ctx.register_dialect(func.Func)
+    ctx.register_dialect(iet_ssa.DEVITO_SSA)
+    ctx.register_dialect(iet_ssa.IET_SSA)
+    ctx.register_dialect(memref.MemRef)
+    ctx.register_dialect(math.Math)
+    p = Parser(ctx, open(sys.argv[-1], "r").read(), Source.MLIR, filename=sys.argv[-1])
+
+    module = p.parse_module()
+
+    grpa = GreedyRewritePatternApplier([
+        _DevitoStencilToStencilStencil(),
+        _LowerGetField(),
+        DropIetComments(),
+        _LowerLoadSymbolidToFuncArgs(),
+        LowerIetForToScfFor(),
+        ConvertScfForArgsToIndex(),
+        ConvertForLoopVarToIndex()
+    ])
+    PatternRewriteWalker(grpa).rewrite_module(module)
+
+    from xdsl.printer import Printer
+    Printer(target=Printer.Target.MLIR).print(module)
+
+

devito/ir/ietxdsl/iet_ssa.py

@@ -4,7 +4,7 @@
 from typing import Iterable, Tuple, List, Annotated, Union, Sequence
 from dataclasses import dataclass
 
-from xdsl.dialects.builtin import (IntegerType, StringAttr, ArrayAttr, OpAttr,
+from xdsl.dialects.builtin import (IntegerType, StringAttr, ArrayAttr, OpAttr, IntegerAttr,
                                    ContainerOf, IndexType, Float16Type, Float32Type,
                                    Float64Type, AnyIntegerAttr, FloatAttr, f32, IntAttr)
 
@@ -417,21 +417,23 @@ class Stencil(Operation):
     """
     Represents a cluster of expressions that should then be translated to a stencil.
     """
-    name = "iet.stencil"
+    name = "devito.stencil"
 
     input_indices: Annotated[VarOperand, AnyIntegerAttr]
     output: Annotated[Operand, AnyIntegerAttr]
 
-    shape: OpAttr[ArrayAttr[IntAttr]]
+    shape: OpAttr[ArrayAttr[IntegerAttr]]
     """
     shape without halo (int, int, int)
     """
-    halo: OpAttr[ArrayAttr[ArrayAttr[IntAttr]]]
+    halo: OpAttr[ArrayAttr[ArrayAttr[IntegerAttr]]]
     """
     how far each dimension is expanded
     ((int, int), (int, int), (int, int))
     """
-    time_buffers: OpAttr[IntAttr]
+    time_buffers: OpAttr[IntegerAttr]
+
+    field_type: OpAttr[Attribute]
 
     body: SingleBlockRegion
 
@@ -460,31 +462,50 @@ def get(time_indices: Sequence[SSAValue | Operation], shape: Sequence[int], halo
         return Stencil.build(
             operands=[inputs, output],
             attributes={
-                'shape': ArrayAttr(IntAttr(x) for x in shape),
-                'halo': ArrayAttr(ArrayAttr(IntAttr(x) for x in inner) for inner in halo),
-                'time_buffers': IntAttr(time_buffers)
+                'shape': ArrayAttr(IntegerAttr(x, 64) for x in shape),
+                'halo': ArrayAttr(ArrayAttr(IntegerAttr(x, 64) for x in inner) for inner in halo),
+                'time_buffers': IntegerAttr(time_buffers, 64),
+                'field_type': typ
             },
             regions=[Region.from_block_list([block])]
         )
 
 
 @irdl_op_definition
 class LoadSymbolic(Operation):
-    name = "iet.load_symbolic"
-    name: OpAttr[StringAttr]
+    name = "devito.load_symbolic"
+    symbol_name: OpAttr[StringAttr]
 
     result: OpResult
 
     @staticmethod
     def get(name: str, typ: Attribute):
         op = LoadSymbolic.build(
-            attributes={'name': StringAttr(name)},
+            attributes={'symbol_name': StringAttr(name)},
             result_types=[typ],
         )
         op.result.name = name
         return op
 
 
+@irdl_op_definition
+class GetField(Operation):
+    name = "devito.get_field"
+
+    data: Annotated[Operand, memref.MemRefType[stencil.FieldType]]
+    t_index: Annotated[Operand, IntegerType]
+    field: Annotated[OpResult, stencil.FieldType]
+
+    lb: OpAttr[stencil.IndexAttr]
+    ub: OpAttr[stencil.IndexAttr]
+
+    @staticmethod
+    def get(data: SSAValue | Operation, t_idx: SSAValue | Operation, field_t: stencil.FieldType, lb: stencil.IndexAttr, ub: stencil.IndexAttr):
+        return GetField.build(
+            operands=[data, t_idx],
+            attributes={'lb': lb, 'ub': ub},
+            result_types=[field_t],
+        )
 
 IET_SSA = Dialect([
     Statement,
@@ -499,4 +520,11 @@ def get(name: str, typ: Attribute):
 ], [
     Profiler,
     Dataobj
+])
+
+DEVITO_SSA = Dialect([
+    Stencil,
+    LoadSymbolic,
+    GetField
+], [
 ])

devito/ir/ietxdsl/lowering.py

@@ -14,12 +14,12 @@ def _generate_subindices(subindices: int, block: Block,
     arg_changes: list[tuple[SSAValue, SSAValue]] = []
 
     # keep track of the ops we want to insert
-    modulo = arith.Constant.from_int_and_width(subindices, builtin.i32)
+    modulo = arith.Constant.from_int_and_width(subindices, builtin.i64)
     new_ops = [modulo]
 
     # generate the new indices
     for i in range(subindices):
-        offset = arith.Constant.from_int_and_width(i, builtin.i32)
+        offset = arith.Constant.from_int_and_width(i, builtin.i64)
         index_off = arith.Addi.get(block.args[0], offset)
         index = arith.RemSI.get(index_off, modulo)
 
@@ -70,7 +70,7 @@ def match_and_rewrite(self, op: scf.ParallelOp, rewriter: PatternRewriter,
         if isinstance(op.lowerBound[0].typ, builtin.IndexType):
             return
         # increment upper bound
-        cst1 = arith.Constant.from_int_and_width(1, builtin.i32)
+        cst1 = arith.Constant.from_int_and_width(1, builtin.i64)
         rewriter.insert_op_before_matched_op(cst1)
         for ub in op.upperBound:
             new_ub = arith.Addi.get(ub, cst1)
@@ -95,21 +95,21 @@ def match_and_rewrite(self, op: Operation, rewriter: PatternRewriter, /):
             return
 
         block: Block = op.regions[0].blocks[0]
-        # change all loop vars to index and immediately cast to i32
+        # change all loop vars to index and immediately cast to i64
         for loop_var in block.args:
             if isinstance(loop_var.typ, builtin.IndexType):
                 continue
             loop_var.typ = builtin.IndexType()
-            # insert a cast from index to i32 at the start of the loop
+            # insert a cast from index to i64 at the start of the loop
 
             rewriter.insert_op_at_pos(
-                i32_val := arith.IndexCastOp.get(loop_var, builtin.i32),
+                i64_val := arith.IndexCastOp.get(loop_var, builtin.i64),
                 block,
                 0,
             )
 
-            loop_var.replace_by(i32_val.result)
-            i32_val.replace_operand(0, loop_var)
+            loop_var.replace_by(i64_val.result)
+            i64_val.replace_operand(0, loop_var)
 
 
 class LowerIetForToScfFor(RewritePattern):
@@ -227,32 +227,32 @@ def match_and_rewrite(self, op: iet_ssa.PointerCast,
         memref_typ = op.result.typ
         assert isinstance(memref_typ, memref.MemRefType)
 
-        # u_vec_size_ptr_addr : llvm.ptr<llvm.ptr<i32>>
+        # u_vec_size_ptr_addr : llvm.ptr<llvm.ptr<i64>>
         u_vec_size_ptr_addr = llvm.GetElementPtrOp.get(
             op.arg,
             [0, 1],
             result_type=llvm.LLVMPointerType.typed(
-                llvm.LLVMPointerType.typed(builtin.i32)),
+                llvm.LLVMPointerType.typed(builtin.i64)),
         )
         # dereference the ptr ptr to get a normal pointer
-        # u_vec_size_ptr : llvm.ptr<i32>
+        # u_vec_size_ptr : llvm.ptr<i64>
         u_vec_size_ptr = llvm.LoadOp.get(u_vec_size_ptr_addr)
 
         new_ops = [u_vec_size_ptr_addr, u_vec_size_ptr]
 
         for i in range(len(memref_typ.shape)):
-            # calculate ptr = (u_vec_size_ptr + i) : llvm.ptr<i32>
+            # calculate ptr = (u_vec_size_ptr + i) : llvm.ptr<i64>
             # TODO: since the actual type of `u_vec->size` is `ptr<i64>` we skip every
             #       second entry and only use the lower 32 bits of the size.
             #       this is okay as long as we don't work on data with more than
             #       2^32 cells, if each cell is 32bit long that is ~128GB of data.
             ptr_to_size_i = llvm.GetElementPtrOp.get(
                 u_vec_size_ptr,
                 [i*2],  # get the 2-i-th element
-                result_type=llvm.LLVMPointerType.typed(builtin.i32),
+                result_type=llvm.LLVMPointerType.typed(builtin.i64),
             )
             # dereference
-            # val = *(ptr) : i32
+            # val = *(ptr) : i64
             load = llvm.LoadOp.get(ptr_to_size_i)
             new_ops += [ptr_to_size_i, load]
             self.dimensions.append(load.dereferenced_value)
@@ -290,11 +290,14 @@ def match_and_rewrite(self, op: func.FuncOp, rewriter: PatternRewriter, /):
             elif isinstance(arg_typ, iet_ssa.Profiler):
                 op.body.blocks[0].args[i].typ = llvm.LLVMPointerType.opaque()
 
-        op.attributes['function_type'] = builtin.FunctionType.from_lists(
-            [arg.typ for arg in op.body.blocks[0].args],
-            op.function_type.outputs.data,
-        )
+        recalc_func_type(op)
+
 
+def recalc_func_type(op: func.FuncOp):
+    op.attributes['function_type'] = builtin.FunctionType.from_lists(
+        [arg.typ for arg in op.body.blocks[0].args],
+        op.function_type.outputs.data,
+    )
 
 @dataclass
 class LowerMemrefLoadToLLvmPointer(RewritePattern):