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,183 @@
+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))
+
+
+def softmax(x, axis=None):
+    x_max = np.amax(x, axis=axis, keepdims=True)
+    exp_x_shifted = np.exp(x - x_max)
+    return exp_x_shifted / np.sum(exp_x_shifted, axis=axis, keepdims=True)
+
+
+def manual_attn(q, k, v):
+    att = q @ k.transpose(0, 1, 3, 2) * (1.0 / float(np.sqrt(k.shape[-1])))
+    att = softmax(att, axis=-1)
+    y = att @ v
+    return y
+
+
+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()
+print(gpu_module)
+
+
+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,113 @@
+from collections import deque
+
+from ...ir import Block, Value, Operation
+from .util import (
+    walk_blocks_in_operation,
+    walk_operations,
+    find_ancestor_block_in_region,
+)
+
+
+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()
+
+        def gather_out_values(v: Value):
+            for use in v.uses:
+                user = use.owner.operation
+                owner_block = user.block
+                owner_block = find_ancestor_block_in_region(owner_block)
+                if owner_block != block:
+                    self.out_values.add(v)
+                    break
+
+        for arg in block.arguments:
+            gather_out_values(arg)
+        for op in block.operations:
+            for r in op.results:
+                gather_out_values(r)
+
+        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
+
+    def update_livein(self):
+        new_in = self.use_values
+        new_in |= self.out_values
+        new_in -= self.def_values
+
+        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):
+    visited = set()
+    to_process = deque()
+    builders = {}
+    for b in walk_blocks_in_operation(op):
+        builder = builders[b] = BlockInfoBuilder(b)
+        if builder.update_livein():
+            for p in b.predecessors:
+                if p not in visited:
+                    to_process.append(p)
+                    visited.add(p)
+
+    while to_process:
+        current = to_process.popleft()
+        builder = builders[current]
+        builder.update_liveout(builders)
+        if builder.update_livein():
+            for p in current.predecessors:
+                if p not in visited:
+                    to_process.append(p)
+                    visited.add(p)
+
+    return builders
+
+
+class LivenessBlockInfo:
+    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
+
+
+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
+            )