add liveness analysis demo

Author: makslevental

.github/workflows/test.yml

@@ -75,6 +75,7 @@ jobs:
           
           python examples/mwe.py
           python examples/flash_attention.py
+          python examples/liveness_analysis.py
 
   test-other-host-bindings:
 

examples/liveness_analysis.py

@@ -0,0 +1,167 @@
+from mlir import ir
+from pathlib import Path
+
+import mlir.extras.types as T
+import numpy as np
+from mlir.ir import InsertionPoint, IntegerAttr, UnitAttr
+
+from mlir.extras.ast.canonicalize import canonicalize
+from mlir.extras.context import RAIIMLIRContextModule
+from mlir.extras.dialects.ext import memref, scf, arith, gpu, llvm
+from mlir.dialects import math
+
+# noinspection PyUnresolvedReferences
+from mlir.extras.dialects.ext.gpu import (
+    block_idx,
+    thread_idx,
+    grid_dim,
+    func as gpu_func,
+    set_container_module,
+    module,
+    get_compile_object_bytes,
+)
+from mlir.extras.runtime.passes import run_pipeline, Pipeline
+from mlir.extras.util import find_ops, walk_blocks_in_operation, walk_operations
+from mlir.extras.util.liveness import BlockInfoBuilder, Liveness
+
+# just so it doesn't get DCE'd by black/reformat
+# TypeError: 'mlir._mlir_libs._mlir.ir.BlockArgument' object is not subscriptable
+_ = memref
+
+ctx = RAIIMLIRContextModule()
+set_container_module(ctx.module)
+
+
+# just a default attr - actual target is set blow
+@module("kernels", [f'#rocdl.target<abi = "500">'])
+def gpu_module():
+    pass
+
+
+ip = InsertionPoint.at_block_begin(gpu_module.regions[0].blocks[0])
+ip.__enter__()
+
+Bc = 32
+Br = 32
+
+B = 16
+nh = 12
+N = 128
+d = 128
+
+softmax_scale = 1.0 / float(np.sqrt(d))
+
+
+rank_reduce = memref.rank_reduce
+
+
+# https://github.com/tspeterkim/flash-attention-minimal/blob/main/flash.cu
+@gpu_func(emit=True)
+@canonicalize(using=[scf.canonicalizer, arith.canonicalizer])
+def flash_attention(
+    Q: T.memref(B, nh, N, d, T.f32()),
+    K: T.memref(B, nh, N, d, T.f32()),
+    V: T.memref(B, nh, N, d, T.f32()),
+    l: T.memref(B, nh, N, T.f32()),
+    m: T.memref(B, nh, N, T.f32()),
+    O: T.memref(B, nh, N, d, T.f32()),
+):
+    tx = thread_idx.x
+    # batch idx, head_idx
+    bx, by = block_idx.x, block_idx.y
+    # gpu.printf("bx %ld, by %ld\n", bx, by)
+
+    # Offset into Q,K,V,O,l,m - different for each batch and head
+    K = K[bx, by, :, :, rank_reduce]
+    V = V[bx, by, :, :, rank_reduce]
+    Q = Q[bx, by, :, :, rank_reduce]
+    O = O[bx, by, :, :, rank_reduce]
+    l = l[bx, by, :, rank_reduce]
+    m = m[bx, by, :, rank_reduce]
+
+    # Define SRAM for Q,K,V,S
+    sram = gpu.dynamic_shared_memory()
+    Qi = memref.view(sram, (Br, d), dtype=T.f32())
+    Kj = memref.view(sram, (Bc, d), dtype=T.f32(), shift=Qi.n_elements)
+    Vj = memref.view(sram, (Bc, d), dtype=T.f32(), shift=Qi.n_elements + Kj.n_elements)
+    S = memref.view(
+        sram,
+        (Br, Bc),
+        dtype=T.f32(),
+        shift=Qi.n_elements + Kj.n_elements + Vj.n_elements,
+    )
+
+    for bc in scf.range_(0, N, Bc):
+        # Load Kj, Vj to SRAM
+        K_ = K[bc : bc + 1, :]
+        V_ = V[bc : bc + 1, :]
+        for x in scf.range_(0, d):
+            Kj[tx, x] = K_[tx, x]
+            Vj[tx, x] = V_[tx, x]
+
+        for br in scf.range_(0, N, Br):
+            # Load Qi to SRAM, l and m to registers
+            Q_ = Q[br : br + 1, :]
+            for x in scf.range_(0, d):
+                Qi[tx, x] = Q_[tx, x]
+
+            l_ = l[br : br + 1]
+            m_ = m[br : br + 1]
+            row_l_prev = l_[tx]
+            row_m_prev = m_[tx]
+
+            # S = QK^T, row_m = rowmax(S)
+            row_m: T.f32() = float(np.finfo(np.float32).min)
+            for y, row_m, _ in scf.range_(0, Bc, iter_args=[row_m]):
+                sum: T.f32() = 0.0
+                for x, sum, _ in scf.range_(0, d, iter_args=[sum]):
+                    sum += Qi[tx, x] * Kj[y, x]
+                    sum = yield sum
+
+                sum *= softmax_scale
+                S[tx, y] = sum
+
+                if sum > row_m:
+                    row_m_ = yield sum
+                else:
+                    row_m_ = yield row_m
+
+                row_m = yield row_m_
+
+            # P = exp(S - row_m), row_l = rowsum(P)
+            row_l: T.f32() = 0.0
+            for y, row_l, _ in scf.range_(0, Bc, iter_args=[row_l]):
+                S[tx, y] = math.exp(S[tx, y] - row_m)
+                row_l += S[tx, y]
+                row_l = yield row_l
+
+            # Compute new m and l
+            row_m_new = arith.maximumf(row_m_prev, row_m)
+            row_l_new = (
+                math.exp(row_m_prev - row_m_new) * row_l_prev
+                + math.exp(row_m - row_m_new) * row_l
+            )
+            div = 1.0 / row_l_new
+            f1 = row_l_prev * math.exp(row_m_prev - row_m_new)
+            f2 = math.exp(row_m - row_m_new)
+
+            # Write O, l, m to HBM
+            O_ = O[br : br + 1, :]
+            for x in scf.range_(0, d):
+                pv: T.f32() = 0.0  # Pij * Vj
+                for y, pv, _ in scf.range_(0, Bc, iter_args=[pv]):
+                    pv += S[tx, y] * Vj[y, x]
+                    pv = yield pv
+
+                O_[tx, x] = div * (f1 * O_[tx, x] + f2 * pv)
+
+            l_[tx] = row_l_new
+            m_[tx] = row_m_new
+
+            gpu.barrier()
+
+
+ip.__exit__(None, None, None)
+
+assert gpu_module.operation.verify()
+Liveness(gpu_module)

mlir/extras/util/__init__.py

@@ -0,0 +1,7 @@
+from .util import *
+from .util import (
+    _get_previous_frame_idents,
+    _get_sym_name,
+    _update_caller_vars,
+    _unpack_sizes_element_type,
+)

mlir/extras/util/liveness.py

@@ -0,0 +1,146 @@
+from collections import OrderedDict
+
+from ...ir import Block, Value, Operation
+from .util import (
+    walk_blocks_in_operation,
+    walk_operations,
+    find_ancestor_block_in_region,
+    find_ancestor_op_in_block,
+)
+
+# based on https://github.com/llvm/llvm-project/blob/07ae19c132e1b0adbdb3cc036b9f50624e2ed1b7/mlir/lib/Analysis/Liveness.cpp
+
+
+def escapes_block(v: Value, b: Block):
+    # Check if value escapes, i.e., if there's a use
+    # which is in a block that is not "our" block
+    for use in v.uses:
+        user = use.owner
+        owner_block = user.block
+        owner_block = find_ancestor_block_in_region(owner_block)
+        if owner_block != b:
+            return True
+    return False
+
+
+class BlockInfoBuilder:
+    block = None
+    in_values = None
+    out_values = None
+    def_values = None
+    use_values = None
+
+    def __init__(self, block):
+        self.block = block
+        self.in_values = set()
+        self.out_values = set()
+        self.def_values = set()
+        self.use_values = set()
+
+        for arg in block.arguments:
+            self.def_values.add(arg)
+            if escapes_block(arg, self.block):
+                self.out_values.add(arg)
+        for op in block.operations:
+            for r in op.results:
+                if escapes_block(r, self.block):
+                    self.out_values.add(r)
+
+        # Mark all nested operation results as defined, and nested operation
+        # operands as used. All defined value will be removed from the used set
+        # at the end.
+        for op in block.operations:
+            for nested_op in walk_operations(op):
+                self.def_values |= set(nested_op.results)
+                self.use_values |= set(nested_op.operands)
+                for b in walk_blocks_in_operation(nested_op):
+                    self.def_values |= set(b.arguments)
+
+        self.use_values -= self.def_values
+
+    # newIn = use U out - def
+    def update_livein(self):
+        new_in = (self.use_values | self.out_values) - self.def_values
+        # It is sufficient to check the set sizes (instead of their contents) since
+        # the live-in set can only grow monotonically during all update operations.
+        if len(new_in) == len(self.in_values):
+            return set()
+        self.in_values = new_in
+        return new_in
+
+    def update_liveout(self, builders):
+        for succ in self.block.successors:
+            self.out_values |= builders[succ].in_values
+
+
+def build_block_mapping(op) -> dict[Block, BlockInfoBuilder]:
+    to_process = OrderedDict()
+    builders = {}
+    for b in walk_blocks_in_operation(op):
+        builder = builders[b] = BlockInfoBuilder(b)
+        if not builder.update_livein():
+            continue
+        for p in b.predecessors:
+            if p in to_process:
+                continue
+            to_process[p] = True
+
+    while to_process:
+        # Pairs are returned in LIFO order if last is true or FIFO order if false.
+        current, _ = to_process.popitem(last=False)
+        builder = builders[current]
+        builder.update_liveout(builders)
+        if not builder.update_livein():
+            continue
+        for p in current.predecessors:
+            if p in to_process:
+                continue
+            to_process[p] = True
+
+    return builders
+
+
+class LivenessBlockInfo:
+    block: Block = None
+    in_values = None
+    out_values = None
+
+    def __init__(self, block, in_values, out_values):
+        self.block = block
+        self.in_values = in_values
+        self.out_values = out_values
+
+    def is_livein(self, v: Value):
+        return v in self.in_values
+
+    def is_liveout(self, v: Value):
+        return v in self.out_values
+
+    def get_start_operation(self, v: Value):
+        owner = v.owner
+        if self.is_livein(v) or not owner:
+            return self.block.operations[0]
+        return owner
+
+    def get_end_operation(self, v: Value, start_op: Operation):
+        if self.is_livein(v):
+            return self.block.operations[-1]
+        end_op = start_op
+        for use in v.uses:
+            use_op = find_ancestor_op_in_block(use.owner)
+
+
+class Liveness:
+    operation = None
+    block_mapping = None
+
+    def __init__(self, op):
+        self.operation = op
+        self.block_mapping = {}
+
+        builders = build_block_mapping(self.operation)
+        for block, builder in builders.items():
+            assert block == builder.block
+            self.block_mapping[block] = LivenessBlockInfo(
+                builder.block, builder.in_values, builder.out_values
+            )