Skip to content

Stencil storage materialization

stencil_storage_materialization

ApplyOpMaterialization

Bases: RewritePattern

Adds stencil.buffer to any used output of a stencil.apply that is not otherwised mapped to storage.

Source code in xdsl/transforms/experimental/stencil_storage_materialization.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class ApplyOpMaterialization(RewritePattern):
    """
    Adds stencil.buffer to any used output of a stencil.apply that is not otherwised
    mapped to storage.
    """

    @op_type_rewrite_pattern
    def match_and_rewrite(self, op: ApplyOp | CombineOp, rewriter: PatternRewriter, /):
        clone = op.clone()
        new_res: list[OpResult] = []
        buffers: list[BufferOp] = []
        for i, out in enumerate(op.results):
            if should_materialize(out):
                buffer = BufferOp(clone.results[i])
                buffers.append(buffer)
                new_res.append(buffer.res)
            else:
                new_res.append(out)
        if buffers:
            rewriter.replace(op, [clone, *buffers], new_res)
        else:
            clone.erase()

match_and_rewrite(op: ApplyOp | CombineOp, rewriter: PatternRewriter)

Source code in xdsl/transforms/experimental/stencil_storage_materialization.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
@op_type_rewrite_pattern
def match_and_rewrite(self, op: ApplyOp | CombineOp, rewriter: PatternRewriter, /):
    clone = op.clone()
    new_res: list[OpResult] = []
    buffers: list[BufferOp] = []
    for i, out in enumerate(op.results):
        if should_materialize(out):
            buffer = BufferOp(clone.results[i])
            buffers.append(buffer)
            new_res.append(buffer.res)
        else:
            new_res.append(out)
    if buffers:
        rewriter.replace(op, [clone, *buffers], new_res)
    else:
        clone.erase()

StencilStorageMaterializationPass dataclass

Bases: ModulePass

Pass adding stencil.buffer whenever necessary to lower a stencil dialect IR, by adding stencil.buffer on any used stencil.apply output not otherwise mapped to storage.

Source code in xdsl/transforms/experimental/stencil_storage_materialization.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
class StencilStorageMaterializationPass(ModulePass):
    """
    Pass adding stencil.buffer whenever necessary to lower a stencil dialect IR,
    by adding stencil.buffer on any used stencil.apply output not otherwise mapped
    to storage.
    """

    name = "stencil-storage-materialization"

    def apply(self, ctx: Context, op: builtin.ModuleOp) -> None:
        PatternRewriteWalker(
            GreedyRewritePatternApplier(
                [
                    ApplyOpMaterialization(),
                ]
            )
        ).rewrite_module(op)

name = 'stencil-storage-materialization' class-attribute instance-attribute

apply(ctx: Context, op: builtin.ModuleOp) -> None

Source code in xdsl/transforms/experimental/stencil_storage_materialization.py 
58
59
60
61
62
63
64
65
def apply(self, ctx: Context, op: builtin.ModuleOp) -> None:
    PatternRewriteWalker(
        GreedyRewritePatternApplier(
            [
                ApplyOpMaterialization(),
            ]
        )
    ).rewrite_module(op)

should_materialize(temp: SSAValue)

Predicates if a specific stencil.apply output should be buffered. It should if it is used by another stencil.apply and not already buffered or stored.

Source code in xdsl/transforms/experimental/stencil_storage_materialization.py 
15
16
17
18
19
20
21
22
def should_materialize(temp: SSAValue):
    """
    Predicates if a specific stencil.apply output should be buffered.
    It should if it is used by another stencil.apply and not already buffered or stored.
    """
    return any(isinstance(u.operation, ApplyOp) for u in temp.uses) and not any(
        isinstance(u.operation, StoreOp | BufferOp) for u in temp.uses
    )