diff --git a/Dockerfile.base b/Dockerfile.base index 05444d41..87d5e5bb 100644 --- a/Dockerfile.base +++ b/Dockerfile.base @@ -76,6 +76,11 @@ RUN curl -L -H "Accept: application/octet-stream" https://api.github.com/repos/P # Store RISC-V LLVM for TorchSim ENV TORCHSIM_LLVM_PATH=/riscv-llvm/bin +# MLIR Python bindings shipped inside the LLVM release artifact (built by the +# llvm-project CI with -DMLIR_ENABLE_BINDINGS_PYTHON=ON). Lets PyTorchSim load +# mlir.ir / dialects for Python-side MLIR passes. The artifact must be built +# against this image's Python (3.11) or `import mlir` fails on ABI mismatch. +ENV PYTHONPATH=/riscv-llvm/python_packages/mlir_core:$PYTHONPATH ENV TORCHSIM_DIR=/workspace/PyTorchSim # Download Spike simulator diff --git a/PyTorchSimFrontend/extension_codecache.py b/PyTorchSimFrontend/extension_codecache.py index efd4d4cb..704162d9 100644 --- a/PyTorchSimFrontend/extension_codecache.py +++ b/PyTorchSimFrontend/extension_codecache.py @@ -43,24 +43,9 @@ def mlir_compile_command(filename, vectorlane_size, vlen=256): {extension_config.CONFIG_TORCHSIM_LLVM_PATH}/mlir-opt \ -test-loop-padding \ -dma-fine-grained='systolic-array-size={vectorlane_size}' \ - -global-idx='vlen={vlen}' \ -test-pytorchsim-to-vcix='systolic-array-size={vectorlane_size} vlen={vlen}' \ - -test-memref-to-gemmini="vectorlane={vectorlane_size}" \ - -convert-linalg-to-loops \ - -convert-vector-to-scf='full-unroll' \ - -lower-affine \ - -finalize-memref-to-llvm \ - -lower-vector-multi-reduction \ - -convert-vector-to-llvm \ - -convert-arith-to-llvm \ - -convert-math-to-llvm \ - -convert-scf-to-cf \ - -convert-cf-to-llvm \ - -convert-func-to-llvm \ - -convert-index-to-llvm \ - -reconcile-unrealized-casts \ {'--mlir-print-ir-after-all' if extension_config.CONFIG_TORCHSIM_DUMP_MLIR_IR else ''} \ - {filename}.mlir -o {filename}_llvm.mlir + {filename}.mlir -o {filename}_custom.mlir """, ).strip(), re.sub(r"[ \n]+", " ", @@ -93,25 +78,10 @@ def mlir_gem5_compile_command(filename, sample_filename, tog_file, vectorlane_si {extension_config.CONFIG_TORCHSIM_LLVM_PATH}/mlir-opt \ -test-loop-padding='timing_mode=1' \ -dma-fine-grained='systolic-array-size={vectorlane_size}' \ - -global-idx='vlen={vlen}' \ -test-pytorchsim-to-vcix='systolic-array-size={vectorlane_size} vlen={vlen}' \ -test-tile-operation-graph='vectorlane={vectorlane_size} sample-mode={extension_config.CONFIG_TLS_MODE}' \ - -test-memref-to-gemmini="vectorlane={vectorlane_size} timing=1" \ - -convert-linalg-to-loops \ - -convert-vector-to-scf='full-unroll' \ - -lower-affine \ - -finalize-memref-to-llvm \ - -lower-vector-multi-reduction \ - -convert-vector-to-llvm \ - -convert-arith-to-llvm \ - -convert-math-to-llvm \ - -convert-scf-to-cf \ - -convert-cf-to-llvm \ - -convert-func-to-llvm \ - -convert-index-to-llvm \ - -reconcile-unrealized-casts \ {'--mlir-print-ir-after-all' if extension_config.CONFIG_TORCHSIM_DUMP_MLIR_IR else ''} \ - {filename}.mlir -o {sample_filename}_llvm.mlir + {filename}.mlir -o {sample_filename}_custom.mlir """, ).strip(), re.sub(r"[ \n]+", " ", @@ -158,6 +128,11 @@ def load(cls, source_code, vlenb = vlen // 8 write_path = get_write_path(source_code) key, input_path = write(source_code, "mlir", specified_dir=write_path) + # Run the Python out-of-line MLIR passes (MLIR bindings) on the kernel + # .mlir in place, before mlir-opt. Currently lowers torchsim.vlane_idx + # (replaces the old C++ -global-idx pass); add more in passes/__init__.py. + from PyTorchSimFrontend.mlir.passes import run_python_passes, run_standard_lowering + run_python_passes(input_path) new_input_path = os.path.splitext(input_path)[0] raw_tog_path = new_input_path + "_tog.py" tog_path = os.path.join(write_path, "tile_graph.onnx") @@ -185,6 +160,10 @@ def load(cls, source_code, with lock: try: subprocess.check_call(opt_cmd) + # Standard MLIR -> LLVM-dialect lowering (registered upstream + # passes) runs in-process via the bindings PassManager, picking + # up after the custom mlir-opt passes (memref-to-gemmini). + run_standard_lowering(new_input_path + "_custom.mlir", new_input_path + "_llvm.mlir") subprocess.check_call(translate_cmd) subprocess.check_call(llc_cmd) subprocess.check_call(llc_asm_cmd) @@ -223,6 +202,8 @@ def load(cls, source_code, result = subprocess.check_output(gem5_sample_cmd) with open(raw_tog_path, "wb") as file: file.write(result) + # Standard MLIR -> LLVM-dialect lowering in-process (see functional path). + run_standard_lowering(sample_mlir_path + "_custom.mlir", sample_mlir_path + "_llvm.mlir", timing=True) subprocess.check_call(gem5_translate_cmd) subprocess.check_call(gem5_llc_cmd) except subprocess.CalledProcessError as e: diff --git a/PyTorchSimFrontend/mlir/axis_split.py b/PyTorchSimFrontend/mlir/axis_split.py new file mode 100644 index 00000000..1c33e021 --- /dev/null +++ b/PyTorchSimFrontend/mlir/axis_split.py @@ -0,0 +1,301 @@ +"""Aligned axis splitting at the Inductor scheduling layer. + +Goal: guarantee the MLIR codegen sees only per-axis affine index expressions +(no FloorDiv / ModularIndexing). When an index expr contains FloorDiv(v, k) or +ModularIndexing(v, k, m) where `v` is a single iteration variable of extent E +and the divisor (resp. k*m) divides E, the floor/mod is *aligned*: splitting the +loop axis v into (outer, inner) with v = outer*k + inner makes it collapse to a +plain affine term (outer), at zero data-movement cost. + +This is the cheap upstream tool of the affine-only contract. The misaligned case +(cat / non-factor reshape, divisor does not divide the extent) is NOT handled +here -- that needs graph-level copy insertion. + +The rebuild reuses Inductor's own LoopBody machinery, exactly like +MLIRScheduling.revert_group: feed a split var_ranges + iter_vars and re-trace the +node's store function so the index expressions are regenerated over the new +iteration domain. +""" +import sympy +from torch._inductor.ir import LoopBody +from torch._inductor.utils import sympy_index_symbol +from torch.utils._sympy.functions import FloorDiv, ModularIndexing + + +def _as_int(x): + try: + return int(x) + except (TypeError, ValueError): + return None + + +def collect_boundaries(exprs, var_to_axis, var_ranges): + """{axis_index: set(boundary cut points)} for the given index expressions. + + A FloorDiv(v, k) contributes boundary k; ModularIndexing(v, k, m) contributes + k and k*m. Only aligned terms count (boundary divides the var extent). Shared + by find_split_plan (fused LoopBody) and graph_copy (operand loaders). + """ + import collections + bset = collections.defaultdict(set) + for expr in exprs: + for fd in expr.atoms(FloorDiv): + base, div = fd.args + k = _as_int(div) + if base in var_to_axis and k and k > 1: + E = _as_int(var_ranges.get(base)) + if E and E % k == 0: + bset[var_to_axis[base]].add(k) + for mi in expr.atoms(ModularIndexing): + base, div, mod = mi.args + k, m = _as_int(div), _as_int(mod) + if base in var_to_axis and k and m: + E = _as_int(var_ranges.get(base)) + if E and E % (k * m) == 0: + ax = var_to_axis[base] + if k > 1: + bset[ax].add(k) + if k * m < E: + bset[ax].add(k * m) + return bset + + +def _is_chain(boundaries, E): + """True iff [1, sorted(boundaries in (1,E)), E] is a divisibility chain.""" + chain = [1] + sorted(b for b in boundaries if 1 < b < E) + [E] + return all(chain[i + 1] % chain[i] == 0 for i in range(len(chain) - 1)) + + +def ledger(nodes, plan): + """Classify every FloorDiv/ModularIndexing in the kernel against `plan`. + + Returns a list of (op_name, reason, term_str) for the terms NOT covered by + axis-split, so we can measure how often the graph-copy cases (incompatible + radix / non-dividing / multi-axis / dynamic) actually reach codegen. Read-only. + Reasons: covered terms are omitted; uncovered ones are + multi_axis_arg - floor/mod argument is not a single iter var (case 7) + non_dividing - divisor (or k*m) does not divide the extent (case 6) + incompatible_radix - single var, divides, but boundaries did not form a + divisibility chain so the axis was left unsplit (case 5) + dynamic - symbolic divisor/extent + """ + rows = [] + + def classify(base, k, m, var_to_axis, var_ranges): + if not (isinstance(base, sympy.Symbol) and base in var_to_axis): + return None if False else "multi_axis_arg" + ax = var_to_axis[base] + E = _as_int(var_ranges.get(base)) + if k is None or E is None or (m is not None and _as_int(m) is None): + return "dynamic" + if ax in plan: + return "covered" + period = k if m is None else k * _as_int(m) + if period and E % period != 0: + return "non_dividing" + return "incompatible_radix" + + for n in nodes: + body = getattr(n, "_body", None) + if body is None: + continue + op = n.get_name() if hasattr(n, "get_name") else "?" + var_to_axis = {v: i for i, v in enumerate(body.iter_vars)} + for expr in body.indexing_exprs.values(): + for fd in expr.atoms(FloorDiv): + r = classify(fd.args[0], _as_int(fd.args[1]), None, var_to_axis, body.var_ranges) + if r and r != "covered": + rows.append((op, r, str(fd))) + for mi in expr.atoms(ModularIndexing): + r = classify(mi.args[0], _as_int(mi.args[1]), mi.args[2], var_to_axis, body.var_ranges) + if r and r != "covered": + rows.append((op, r, str(mi))) + return rows + + +def find_split_plan(nodes): + """Inspect a group of scheduler nodes and return {axis_index: boundaries}. + + `boundaries` is an ascending divisibility chain [1, b1, ..., E] of cut points + for that axis: splitting the axis at these boundaries (mixed radix, + `v = sum_i d_i * b_i`) makes every FloorDiv/ModularIndexing on it collapse to + an affine combination of the split sub-vars. The cut points are gathered from + the terms on the axis: + - FloorDiv(v, k) -> boundary k + - ModularIndexing(v, k, m) -> boundaries k and k*m (the digit lives in [k, k*m)) + Only aligned terms count (the boundary must divide the extent E). If the + collected boundaries for an axis do NOT form a divisibility chain (e.g. + floor-by-2 and mod-by-3 on extent 6), the radices are incompatible -> the axis + is left unsplit (its floor/mod stays for the misaligned/recompile path). + + axis_index is positional in the group's iteration space, so the same plan + applies to every fused node sharing that space. + """ + import collections + bset = collections.defaultdict(set) # axis -> set of boundary cut points + ext_of = {} # axis -> extent + for n in nodes: + body = getattr(n, "_body", None) + if body is None: + continue + var_to_axis = {v: i for i, v in enumerate(body.iter_vars)} + nb = collect_boundaries(body.indexing_exprs.values(), var_to_axis, body.var_ranges) + for ax, bs in nb.items(): + bset[ax] |= bs + ext_of[ax] = _as_int(body.var_ranges[body.iter_vars[ax]]) + + plan = {} + for ax, bs in bset.items(): + E = ext_of[ax] + # require a real, divisibility-chain split (incompatible radices -> skip). + if E and any(1 < b < E for b in bs) and _is_chain(bs, E): + plan[ax] = [1] + sorted(b for b in bs if 1 < b < E) + [E] + + # Validation aid: force-split the first even index axis even without floor/mod. + # A floor-free index split is an identity transformation, so allclose must hold; + # used to exercise the reduction pass-through path (no natural op produces a + # floor on a reduction kernel's index axis). Off unless TORCHSIM_AXIS_SPLIT_FORCE. + import os as _os + if _os.environ.get("TORCHSIM_AXIS_SPLIT_FORCE"): + for n in nodes: + body = getattr(n, "_body", None) + if body is None or not body.reduce_vars: + continue + for ax, v in enumerate(body.iter_vars): + E = _as_int(body.var_ranges.get(v)) + if ax not in plan and E and E % 2 == 0 and E > 2: + plan[ax] = [1, 2, E] + break + + # Rank guard: if the split would push the index rank past 4, skip it and fall + # back to baseline. The >4D logical tile is *meant* to be peeled into <=4D + # physical descriptors by the decompose-transfer pass, and the #258 TOG crash + # (arith.addi DRAM offset) is now fixed -- but the peel still has a numerical + # correctness bug (pixel_shuffle -> MISMATCH; the peel was only ever isolation- + # validated for MLIR structure, never run end-to-end). Keep the guard until the + # peel numerics are fixed; then this guard can be removed and the recompile-dance + # retired for pixel. + base_rank = next((len(b.iter_vars) for n in nodes + for b in (getattr(n, "_body", None),) if b is not None), 0) + extra = sum(len(ch) - 2 for ch in plan.values()) + if base_rank + extra > 4: + return {} + return plan + + +def build_split_body(node, plan, prefix="z"): + """Rebuild node._body / sizes for the given split plan. + + Returns (body, (index_size, reduce_size)). Reindexes the EXISTING (already + collapsed/reordered) node._body via LoopBody's copy path instead of re-tracing + from the raw store function: pass the body as `fn` so LoopBody.__init__ takes + _init_with_copy, which substitutes each original iter var with our expression + and runs simplify_with_ranges. For a split axis the substitution + v -> sum_i d_i * b_i (mixed radix over the boundary chain) makes every + FloorDiv/ModularIndexing on it collapse to an affine combination of the d_i, + and reindexing the collapsed body keeps already-merged dims merged (no rank + blow-up). indexing_from_args requires exactly one replacement expr per original + var (index dims then reduce dims), flattened to len(body.var_ranges). + """ + body = node._body + orig_index_vars = list(body.iter_vars) + orig_reduce_vars = list(body.reduce_vars) + + iter_vars = [] + index_args = [] # one expr per ORIGINAL index dim (substituted in) + var_ranges = {} + index_size = [] + ctr = 0 + + for ax, v in enumerate(orig_index_vars): + ext = body.var_ranges[v] + if ax in plan: + bounds = plan[ax] # ascending chain [1, b1, ..., E] + # one sub-var per segment: d_i has extent b_{i+1}/b_i, significance b_i. + subs = [] # (symbol, extent, significance) low->high + expr = sympy.Integer(0) + for i in range(len(bounds) - 1): + seg_ext = bounds[i + 1] // bounds[i] + nv = sympy_index_symbol(f"{prefix}{ctr}"); ctr += 1 + subs.append((nv, seg_ext, bounds[i])) + expr = expr + nv * bounds[i] + # iteration nest: most-significant (outermost) dim first. + for nv, seg_ext, _sig in reversed(subs): + iter_vars.append(nv) + var_ranges[nv] = sympy.Integer(seg_ext) + index_size.append(sympy.Integer(seg_ext)) + index_args.append(expr) + else: + nv = sympy_index_symbol(f"{prefix}{ctr}"); ctr += 1 + iter_vars.append(nv) + var_ranges[nv] = ext + index_size.append(ext) + index_args.append(nv) + + # Reduction dims pass through unchanged (a fresh symbol with the same range), + # using the "r" prefix and kept after the index dims so the reduction axis + # stays innermost (var_ranges is ordered iter-then-reduce; sizes splits on + # len(iter_vars)). We do not split reduction dims here. + reduce_vars = [] + reduce_size = [] + reduce_args = [] + for rctr, v in enumerate(orig_reduce_vars): + ext = body.var_ranges[v] + nv = sympy_index_symbol(f"r{rctr}") + reduce_vars.append(nv) + var_ranges[nv] = ext + reduce_size.append(ext) + reduce_args.append(nv) + + args = [index_args, reduce_args] if orig_reduce_vars else [index_args] + new_body = LoopBody(body, args, var_ranges, iter_vars, reduce_vars) + new_body.indexing_exprs = { + name: _fold_with_ranges(e, var_ranges) + for name, e in new_body.indexing_exprs.items() + } + return new_body, (index_size, reduce_size) + + +def _fold_with_ranges(expr, var_ranges): + """Fold residual FloorDiv/ModularIndexing that simplify_with_ranges missed. + + A mixed-radix split leaves terms like FloorDiv(z1 + 4*z2, 12); these are 0 by + construction (the lower digits sum below the boundary), but the Inductor + simplifier cannot prove a multi-term numerator < divisor. We prove it directly + from the split sub-var ranges via bound_sympy: + FloorDiv(num, d) -> 0 if 0 <= num < d + ModularIndexing(num, k, m) -> num // k if 0 <= num < k*m (mod is a no-op) + Iterated to a fixpoint (folding a mod can expose a foldable floor). + """ + from torch.utils._sympy.value_ranges import bound_sympy, ValueRanges + ranges = {} + for v, sz in var_ranges.items(): + e = _as_int(sz) + if e is not None and e >= 1: + ranges[v] = ValueRanges(0, e - 1) + if not ranges: + return expr + + def vr(num): + try: + return bound_sympy(num, ranges) + except Exception: + return None + + for _ in range(8): + changed = False + for fd in list(expr.atoms(FloorDiv)): + num, div = fd.args + d = _as_int(div) + b = vr(num) if d else None + if b is not None and b.lower >= 0 and b.upper < d: + expr = expr.subs(fd, sympy.Integer(0)); changed = True + for mi in list(expr.atoms(ModularIndexing)): + num, k, m = mi.args + ki, mi_ = _as_int(k), _as_int(m) + b = vr(num) if (ki and mi_) else None + if b is not None and b.lower >= 0 and b.upper < ki * mi_: + expr = expr.subs(mi, FloorDiv(num, k)); changed = True + if not changed: + break + return expr diff --git a/PyTorchSimFrontend/mlir/graph_copy.py b/PyTorchSimFrontend/mlir/graph_copy.py new file mode 100644 index 00000000..51c2e9b6 --- /dev/null +++ b/PyTorchSimFrontend/mlir/graph_copy.py @@ -0,0 +1,163 @@ +"""Graph-copy (relayout) for incompatible-radix operands. + +When an elementwise consumer reads two operands whose floor/mod groupings on a +shared axis are incompatible (the boundary cut points do not form a divisibility +chain, e.g. floor-by-2 and mod-by-3 on extent 6), axis-split cannot linearize the +fused index. We `realize()` the cheaper operand at the consumer's lowering, which +materializes it as a contiguous buffer; the consumer then reads it affine and only +the other (single, compatible) grouping remains for axis-split to handle. + +Detection reuses axis_split.collect_boundaries on each operand's loader index, so +it is the same precise radix analysis used at the scheduling layer -- not an FX +view-chain heuristic. The hook wraps the already-registered lowering entries (the +make_pointwise results), so it sees every elementwise consumer in one place. The +realize() (not a clone, which Inductor inlines) is what actually forces the buffer +boundary; see the PoC notes in docs. + +Gated by TORCHSIM_GRAPH_COPY (install() is a no-op otherwise). Behavior-neutral +unless a genuine incompatible-radix conflict is detected. +""" +import os +from torch._inductor import lowering as L +from torch._inductor import dependencies +from torch._inductor import ir +from torch._inductor.ir import TensorBox +from torch.utils._sympy.functions import FloorDiv, ModularIndexing + +from . import axis_split + + +def _has_multivar_floormod(exprs): + """True if any FloorDiv/ModularIndexing argument spans >1 loop variable + (case 7: cross-axis floor/mod that axis-split cannot split).""" + for e in exprs: + for f in list(e.atoms(FloorDiv)) + list(e.atoms(ModularIndexing)): + if len(f.args[0].free_symbols) > 1: + return True + return False + + +def _numel(tb): + n = 1 + for s in tb.get_size(): + v = axis_split._as_int(s) + if v is None: + return float("inf") + n *= v + return n + + +def _relayout_args(args): + """Return a modified args list with one operand replaced by a forced copy when + it needs relayout, or None to leave args unchanged. The copy uses + ExternKernel.copy_input (a realized identity Pointwise) -- this materializes + *views* too, unlike StorageBox.realize() which is a no-op on a ReinterpretView. + The copy kernel iterates the operand's own (contiguous) shape, so its index + collapses to single-var and axis-split handles it; the consumer then reads the + copy affine.""" + pos = [i for i, x in enumerate(args) if isinstance(x, TensorBox)] + if not pos: + return None + tbs = [args[i] for i in pos] + # Output/iteration shape = the broadcast of all operands (the largest rank, + # max per dim). For a single-operand consumer (e.g. a reduction reading a + # multi-var-view input) this is just that operand's shape -- still enough to + # detect a multi-var floor and copy_input it (case 7); the 2-operand radix + # conflict (case 5) naturally needs >=2 operands. + ranges = max((t.get_size() for t in tbs), key=len) + extents = [axis_split._as_int(s) for s in ranges] + dbg = os.environ.get("TORCHSIM_GRAPH_COPY_DEBUG") + if dbg: + print(f"[GC] consumer ntbs={len(tbs)} ranges={extents} " + f"sizes={[[axis_split._as_int(s) for s in t.get_size()] for t in tbs]}") + if not extents or any(e is None for e in extents): + return None # scalar / dynamic -> skip + + # Only true elementwise consumers: each operand is broadcast-compatible with the + # output (same rank, every dim is 1 or == the output extent). This admits + # broadcasting operands (e.g. y[8,1] into [8,3]) while excluding mm/bmm/cat-style + # ops whose operands differ in a non-broadcast way. + for tb in tbs: + sz = [axis_split._as_int(s) for s in tb.get_size()] + if len(sz) != len(extents) or any( + d is not None and d != 1 and d != e for d, e in zip(sz, extents) + ): + return None + + # Trace each operand's loader to get its read indices (sympy) over the shared + # output iteration; make_loader returns a value, so extract_read_writes is what + # gives the index expressions. range_vars are positional per output axis, so the + # axis numbering is consistent across operands. + per_bnd = [] # [{axis: boundary set}] per operand + per_mv = [] # [bool] operand has multi-var floor/mod + for tb in tbs: + try: + rw = dependencies.extract_read_writes(tb.make_loader(), list(ranges)) + except Exception as e: + if dbg: + print(f"[GC] extract fail {type(e).__name__}: {repr(e)[:60]}") + per_bnd.append({}) + per_mv.append(False) + continue + v2a = {v: i for i, v in enumerate(rw.range_vars)} + exprs = [r.index for r in rw.reads if hasattr(r, "index")] + b = axis_split.collect_boundaries(exprs, v2a, rw.var_ranges) + mv = _has_multivar_floormod(exprs) + if dbg: + print(f"[GC] operand reads={[str(e) for e in exprs]} boundaries={dict(b)} multivar={mv}") + per_bnd.append(b) + per_mv.append(mv) + + victim = None + + # Case 5 -- incompatible radices on a shared axis between two operands. + for axis, E in enumerate(extents): + contrib = [(i, per_bnd[i][axis]) for i in range(len(tbs)) if per_bnd[i].get(axis)] + if len(contrib) < 2: + continue # single grouping -> axis-split handles + union = {b for _, s in contrib for b in s} + if axis_split._is_chain(union, E): + continue # compatible -> axis-split handles + victim = min(contrib, key=lambda c: _numel(tbs[c[0]]))[0] + break + + # Case 7 -- an operand whose floor/mod argument spans multiple consumer axes + # (e.g. (3*p0+p1)//4 from a transpose+reshape feeding a broadcast/softmax that + # keeps the dims separate). axis-split cannot split a multi-var argument. + if victim is None: + mv_ops = [i for i in range(len(tbs)) if per_mv[i]] + if mv_ops: + victim = min(mv_ops, key=lambda i: _numel(tbs[i])) + + if victim is None: + return None + new = list(args) + p = pos[victim] + new[p] = ir.ExternKernel.copy_input(args[p]) + if dbg: + print(f"[GC] relayout: copy_input operand #{victim} (arg {p})") + return new + + +def install(): + """Wrap registered lowering entries to insert relayout. Idempotent; ON by + default (set TORCHSIM_GRAPH_COPY=0 to disable). Call once at backend import + (after torch._inductor.lowering is populated -- make_pointwise runs at import + to build the entries, so we wrap the entries, not the factory).""" + if os.environ.get("TORCHSIM_GRAPH_COPY", "1") == "0": + return + if getattr(L, "_torchsim_relayout_installed", False): + return + for key, fn in list(L.lowerings.items()): + def wrap(orig): + def wrapped(*a, **k): + try: + na = _relayout_args(a) + except Exception: + na = None # detection must never break lowering + if na is not None: + a = na + return orig(*a, **k) + return wrapped + L.lowerings[key] = wrap(fn) + L._torchsim_relayout_installed = True diff --git a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py index b163ad1a..19ae3af5 100644 --- a/PyTorchSimFrontend/mlir/mlir_codegen_backend.py +++ b/PyTorchSimFrontend/mlir/mlir_codegen_backend.py @@ -322,6 +322,10 @@ def __init__(self, kernel_group, reason=None): self.spad_buffer_dict = dict() self.base_vector_initialized = False self.loop_size = None + # Set by get_dma_info when a DMA access cannot fit one <=4D Gemmini + # descriptor; load()/store() then emit a togsim.transfer for the + # decompose pass to peel into a loop of <=4D dma_start. + self._dma_needs_transfer = False def reset(self, reason): save = self.exit_stack, self._nested_context_depth @@ -537,9 +541,14 @@ def load(self, name: str, index: sympy.Expr): compute_index_var = ",".join(sram_index_var.split(",")[:-1] + [f"%{self.compute_idx}"]) # MVIN Encoding - attribute = mlir_common.format_dma_op_attributes(dram_stride, tile_stride, int(padding)) - code = self.get_dma_code("MVIN", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, attribute) + if self._dma_needs_transfer: + self._dma_needs_transfer = False + code = self.emit_transfer("MVIN", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, + dram_shape, tile_shape, dram_stride, tile_stride, int(padding)) + else: + attribute = mlir_common.format_dma_op_attributes(dram_stride, tile_stride, int(padding)) + code = self.get_dma_code("MVIN", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, + dram_shape, tile_shape, attribute) self.cse.generate(dma_buffer, code, assignment = False) # FIXME: assignment = False does not support caching if not comptute_depedency: @@ -608,9 +617,14 @@ def store(self, name: str, index: sympy.Expr, value, mode=None, *args, **kwargs) sram_index_var = self.spad_buffer_dict[str(value)][3] # Generate DMA instruction - attribute = mlir_common.format_dma_op_attributes(dram_stride, tile_stride, 0) - code = self.get_dma_code("MVOUT", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, - dram_shape, tile_shape, attribute) + if self._dma_needs_transfer: + self._dma_needs_transfer = False + code = self.emit_transfer("MVOUT", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, + dram_shape, tile_shape, dram_stride, tile_stride, 0) + else: + attribute = mlir_common.format_dma_op_attributes(dram_stride, tile_stride, 0) + code = self.get_dma_code("MVOUT", vlane_split_axis, vlane_stride, mlir_dtype, dram_var, index_var, sram_var, sram_index_var, + dram_shape, tile_shape, attribute) self.dma_stores.writeline(common.DeferredLine(name, code)) def reduction(self, dtype, src_dtype, reduction_type, value): @@ -1243,7 +1257,13 @@ def get_dma_info(self, name, index, broadcast=True, store_reduction=False, buffe local_tile_desc.vmap.vlane_split_axis = local_vlane_split_axis local_tile_desc.vmap.vlane_stride = kg_tile_desc.vmap.vlane_stride else: - raise NotImplementedError("Currently not implemented... ;)") + # >4D access: one Gemmini DMA descriptor (<=4D) cannot represent this. + # Build the full N-D tile and flag it for togsim.transfer; the decompose + # pass peels the excess dims into a loop of <=4D memref.dma_start. + local_tile_desc.set_tile_size([kg_tile_desc.get_dim_size(dim) for dim in local_dims]) + local_tile_desc.vmap.vlane_split_axis = local_vlane_split_axis + local_tile_desc.vmap.vlane_stride = kg_tile_desc.vmap.vlane_stride + self._dma_needs_transfer = True if len(implicit_local_dims)!=0 and len(local_dims) != len(implicit_local_dims) and self.is_modular_indexing(index): for axis_constraints in self.kernel_group.tile_desc.implicit_dim_size.values(): @@ -1426,6 +1446,53 @@ def get_dma_code(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtype return f"memref.dma_start {src_operand}, {dst_operand}, %{dma_type}, {tag_var}, {dma_attribute} : {src_shape}, {dst_shape}, {tag_shape} {attribute}" + def emit_transfer(self, dma_type_name, vlane_split_axis, vlane_stride, mlir_dtype, + dram_var, dram_index_var, sram_var, sram_index_var, + dram_shape, tile_shape, dram_stride, tile_stride, padding): + """Emit a generic togsim.transfer op for a DMA whose access exceeds the + 4D Gemmini descriptor limit. Carries the full N-D access (dram/tile + strides + shapes) plus the SSA operands a memref.dma_start needs + (dma_type / vlane_split_axis / vlane_stride), so the decompose pass + (passes/decompose_transfer.py) is purely mechanical: it peels the excess + dims into a loop of <=4D memref.dma_start, reusing these operands. + + The operand prep mirrors get_dma_code (dma_type enum via the read/write + cache+counter, vlane consts via CSE) so the transfer is self-contained; + togsim is an unregistered dialect -> generic form. + """ + dma_key = (vlane_split_axis, vlane_stride, mlir_dtype) + if dma_type_name == "MVIN" and dma_key in self.dma_read_cache: + dma_type, vsa, vst = self.dma_read_cache[dma_key] + elif dma_type_name == "MVOUT" and dma_key in self.dma_write_cache: + dma_type, vsa, vst = self.dma_write_cache[dma_key] + else: + vsa = self.get_const_cse(vlane_split_axis) + vst = self.get_const_cse(vlane_stride) + if dma_type_name == "MVIN": + dma_type = self.get_const_cse(DMA_TYPE[f"{dma_type_name}{self.dma_read_counter}"]) + self.dma_read_counter += 1 + self.dma_read_cache[dma_key] = [dma_type, vsa, vst] + else: + dma_type = self.get_const_cse(DMA_TYPE[f"{dma_type_name}{self.dma_write_counter}"]) + self.dma_write_cache[dma_key] = [dma_type, vsa, vst] + tag = self.get_tag_cse() + zero_cse = self.get_const_cse(0) + # vlane_split_axis is carried as a VALUE attr (not an SSA operand) because the + # decompose pass must remap it: collapsing unit tile dims renumbers the axes, + # so the descriptor's vlane axis index changes and the pass rebuilds the const. + attrs = ( + f'dma_kind = "{dma_type_name}", ' + f'vlane_split_axis = {int(vlane_split_axis)} : i64, ' + f'dram_stride = {dram_stride}, tile_stride = {tile_stride}, ' + f'padding = {int(padding)} : i64' + ) + # operands: dram, dram_idx, sram, sram_idx, tag, dma_type, vlane_stride + return ( + f'"togsim.transfer"(%{dram_var}, %{dram_index_var}, %{sram_var}, %{zero_cse}, ' + f'%{tag}, %{dma_type}, %{vst}) {{{attrs}}} : ' + f'({dram_shape}, index, {tile_shape}, index, memref<1xi32>, index, index) -> ()' + ) + def allocate_sram_buffer(self, dtype, dram_name, tile_desc, raw_index, buffer=None, forced_name=None): c_type = mlir_common.DTYPE_TO_C[dtype] mlir_dtype = mlir_common.DTYPE_TO_MLIR[dtype] diff --git a/PyTorchSimFrontend/mlir/mlir_common.py b/PyTorchSimFrontend/mlir/mlir_common.py index 734ca967..45bb144a 100644 --- a/PyTorchSimFrontend/mlir/mlir_common.py +++ b/PyTorchSimFrontend/mlir/mlir_common.py @@ -1,5 +1,6 @@ import dataclasses import math +import os import contextvars from contextlib import contextmanager from dataclasses import dataclass @@ -472,29 +473,25 @@ def apply_constraints(self, constraints, ranges): @staticmethod def init_tile_size(ranges, vlane_stride, vector_lane): + # Logical tile init for ANY rank. Only the innermost dims carry the + # vectorized tile; all further-outer dims stay 1. The physical Gemmini DMA + # descriptor is <=4D -- a higher-rank logical tile is mapped onto <=4D + # descriptors by togsim.transfer + the decompose pass (logical/physical + # tile split), so no rank cap here. nr_dim = len(ranges) + if nr_dim == 0: # scalar + return [1] tile_size = [1] * nr_dim - if len(tile_size) == 2: + if nr_dim == 1: + tile_size[0] = 1 if ranges[0] == 1 else 2 * vlane_stride * vector_lane + elif nr_dim == 2: tile_size[-1] = vlane_stride * vector_lane tile_size[-2] = 2 * vector_lane - elif len(tile_size) == 0: # Scalar - tile_size = [1] - ranges = [1] - elif len(tile_size) == 1 and ranges[0]==1: - tile_size[0] = 1 - elif len(tile_size) == 1: - tile_size[0] = 2 * vlane_stride * vector_lane - elif len(tile_size) == 3: + else: # 3D and up (general) tile_size[-1] = vector_lane tile_size[-2] = 4 * vector_lane tile_size[-3] = 2 - elif len(tile_size) == 4: - tile_size[-1] = vector_lane - tile_size[-2] = 4 * vector_lane - tile_size[-3] = 2 - tile_size[-4] = 1 - else: - raise NotImplementedError("dummy tile size fail!") + # tile_size[:-3] stay 1 (subsumes the old 4D [-4]=1 and any higher rank) return tile_size @staticmethod @@ -840,10 +837,14 @@ def codegen_nodes(self, nodes, kernel_name): node.run(vars, reduction_vars) except RecompileSignal as e: recompile_try += 1 + # Measure what still depends on the recompile-dance once axis-split + + # graph-copy are on by default (set TORCHSIM_RECOMPILE_LOG=1). + if os.environ.get("TORCHSIM_RECOMPILE_LOG"): + import sys as _sys + print(f"[RECOMPILE {recompile_try}/{max_retry_compile}] {e}", file=_sys.stderr) if recompile_try > max_retry_compile: raise RuntimeError("Failed to compile kernel after multiple attempts.") # Retry compile nodes - #print(f"Try recompile({recompile_try}/{max_retry_compile}). Reason: {e}") continue V.graph.removed_buffers |= self.removed_buffers # V.graph.inplaced_to_remove |= self.inplaced_to_remove diff --git a/PyTorchSimFrontend/mlir/mlir_ops.py b/PyTorchSimFrontend/mlir/mlir_ops.py index 217129e8..f1fb4186 100644 --- a/PyTorchSimFrontend/mlir/mlir_ops.py +++ b/PyTorchSimFrontend/mlir/mlir_ops.py @@ -1135,11 +1135,19 @@ def extract_strided_slice(operand, target_size, offsets=None, sizes=None, stride @staticmethod def vlane_offset(operand1, operand2, *args, **kwargs): + # Emit a dedicated torchsim.vlane_idx op (generic form; torchsim is an + # unregistered dialect) instead of overloading arith.addi with a + # vlane_offset attribute. A Python out-of-line pass lowers it to + # (vcix.v.i per-lane index * offset); see + # PyTorchSimFrontend/mlir/passes/lower_vlane_idx.py. tile_size, ret_type, operand1, operand2 = ExtensionOverrides.binary_elementwise_common(operand1, operand2) shape = f"vector<{tile_size}x{ret_type}>" if tile_size > 1 else ret_type - opcode = f'arith.add{ret_type[0]}' - op_str = f'{opcode} %{operand1}, %{operand2}' - return format_mlir_op(op_str, shape, **kwargs), [tile_size, ret_type] + offset = kwargs.get("attributes", {}).get("vlane_offset", 0) + op_str = '"torchsim.vlane_idx"()' + func_type = f'() -> {shape}' + return format_mlir_op(op_str, func_type, + attributes={"vlane_offset": f"{offset} : i64"}, + comment=kwargs.get("comment")), [tile_size, ret_type] @staticmethod def multi_reduction(acc, init, vec_size, red_size, red_shape, red_type, type_name, *args, **kwargs): diff --git a/PyTorchSimFrontend/mlir/mlir_scheduling.py b/PyTorchSimFrontend/mlir/mlir_scheduling.py index 22d1011b..48eead47 100644 --- a/PyTorchSimFrontend/mlir/mlir_scheduling.py +++ b/PyTorchSimFrontend/mlir/mlir_scheduling.py @@ -249,6 +249,44 @@ def codegen_node(self, _node): nodes, key=lambda x: int(x.is_reduction()) ).group + def _dump_axis(tag): + import sys as _sys + print(f"\n[AXIS_SPLIT:{tag}] group={group} reduction_group={reduction_group}", file=_sys.stderr) + for _n in nodes: + _body = getattr(_n, "_body", None) + if _body is None: + continue + print(f"[AXIS_SPLIT:{tag}] node={_n.get_name()} var_ranges={getattr(_body, 'var_ranges', None)}", file=_sys.stderr) + for _k, _e in getattr(_body, "indexing_exprs", {}).items(): + print(f"[AXIS_SPLIT:{tag}] idx[{_k}] = {_e}", file=_sys.stderr) + + if os.environ.get("TORCHSIM_DEBUG_AXIS_SPLIT"): + _dump_axis("before") + + if os.environ.get("TORCHSIM_AXIS_LEDGER"): + from . import axis_split + import sys as _sys + _plan = axis_split.find_split_plan(nodes) + for _op, _reason, _term in axis_split.ledger(nodes, _plan): + print(f"[AXIS_LEDGER] op={_op} reason={_reason} term={_term}", file=_sys.stderr) + + # axis-split is ON by default; set TORCHSIM_AXIS_SPLIT=0 to disable. + if os.environ.get("TORCHSIM_AXIS_SPLIT", "1") != "0": + from . import axis_split + plan = axis_split.find_split_plan(nodes) + if plan: + for _n in nodes: + if getattr(_n, "_body", None) is None: + continue + _body, _ranges = axis_split.build_split_body(_n, plan) + _n._sizes, _n._body, _n.group = _ranges, _body, (_n.get_device(), self.group_fn(_ranges)) + _, (group, reduction_group) = max( + nodes, key=lambda x: int(x.is_reduction()) + ).group + if os.environ.get("TORCHSIM_DEBUG_AXIS_SPLIT"): + print(f"[AXIS_SPLIT] applied plan={plan}", file=__import__("sys").stderr) + _dump_axis("after") + # Note: We assume that there is at least one loop in the nodes # But, inductor simplifies the group, there could be no loop # In that case, we add dummy loop(size=1) to the group @@ -353,3 +391,9 @@ def get_order(n): if origins: _, _, last = max(origins) V.graph.wrapper_code.enter_context(last) + + +# Install the graph-copy (incompatible-radix relayout) lowering hook once at import. +# No-op unless TORCHSIM_GRAPH_COPY is set; see graph_copy.py. +from . import graph_copy as _graph_copy +_graph_copy.install() diff --git a/PyTorchSimFrontend/mlir/passes/__init__.py b/PyTorchSimFrontend/mlir/passes/__init__.py new file mode 100644 index 00000000..e69bfe68 --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/__init__.py @@ -0,0 +1,70 @@ +"""Python out-of-line MLIR passes run on each kernel .mlir before mlir-opt. + +MLIR's PassManager only schedules *registered C++ passes*, not arbitrary Python +functions, so imperative Python rewrites are orchestrated here instead. The flow +is Module-centric: parse the .mlir once, run each registered pass on the shared +Module, print once. A text marker check skips parsing entirely when no pass's +target op is present (the common case). + +To add a pass, create a module exposing MARKERS (tuple of op-name strings) and +run(module) (mutates the Module in place), and append it to PASSES below. +""" +def _ensure_mlir_bindings_on_path(): + """Make `import mlir` work even when PYTHONPATH is not set, by deriving the + bindings location from TORCHSIM_LLVM_PATH (e.g. /riscv-llvm/bin -> + /riscv-llvm/python_packages/mlir_core). The container sets PYTHONPATH, but + plain local runs may not.""" + try: + import mlir.ir # noqa: F401 + return + except ModuleNotFoundError: + pass + import os + import sys + from PyTorchSimFrontend import extension_config + llvm_path = (extension_config.CONFIG_TORCHSIM_LLVM_PATH or "").rstrip("/") + cand = os.path.join(os.path.dirname(llvm_path), "python_packages", "mlir_core") + if os.path.isdir(cand) and cand not in sys.path: + sys.path.insert(0, cand) + + +_ensure_mlir_bindings_on_path() + +from . import lower_vlane_idx +from . import decompose_transfer +from .lower_to_llvm import run_standard_lowering # noqa: F401 (re-exported) + +# Ordered passes applied to each kernel .mlir before mlir-opt. +# decompose_transfer first: it lowers togsim.transfer -> memref.dma_start, which +# downstream passes (and the gemmini lowering) expect. +PASSES = [ + decompose_transfer, + lower_vlane_idx, +] + + +def run_python_passes(mlir_path): + """Apply all registered Python MLIR passes to the .mlir at `mlir_path`, in place. + + Returns True if the file was modified, False otherwise. + """ + with open(mlir_path) as f: + text = f.read() + + # Fast path: nothing to do if no pass's target op appears in the text. + active = [p for p in PASSES if any(mk in text for mk in p.MARKERS)] + if not active: + return False + + from mlir.ir import Context, Module, Location + ctx = Context() + ctx.allow_unregistered_dialects = True + with ctx, Location.unknown(): + module = Module.parse(text) + for p in active: + p.run(module) + out = str(module) + + with open(mlir_path, "w") as f: + f.write(out) + return True diff --git a/PyTorchSimFrontend/mlir/passes/decompose_transfer.py b/PyTorchSimFrontend/mlir/passes/decompose_transfer.py new file mode 100644 index 00000000..87b8aadf --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/decompose_transfer.py @@ -0,0 +1,216 @@ +"""Python out-of-line MLIR pass: decompose togsim.transfer -> <=4D memref.dma_start. + +A togsim.transfer carries a per-axis affine DMA whose descriptor rank may exceed +the 4D Gemmini limit. This pass is a **pure mechanical rank peel** of that +already-affine access (see docs/dma-transfer-lowering.md, "aligned-only peel"): + + - drop unit (extent-1) tile dims: they contribute no descriptor axis; + - if the remaining (effective) rank <= 4 -> emit one customized + memref.dma_start, reusing the transfer's operands (fast path); + - if effective rank > 4 -> peel the outer dims into a loop, adjusting the + base index by stride*iv per iteration, inner descriptor <=4D. + +It does NO floor/mod linearization (aligned split happens upstream at the +scheduling layer) and NO relayout (misaligned access is copy-inserted at the +graph level). A transfer whose access is not per-axis affine is a contract +violation -- but by construction codegen only emits affine transfers. + +togsim.transfer operands (see emit_transfer): + (dram, dram_idx, sram, sram_idx, tag, dma_type, vlane_split_axis, vlane_stride) +attrs: dma_kind ("MVIN"/"MVOUT"), dram_stride[], tile_stride[], padding. + +memref.dma_start (customized) operands: + src[idx], dst[idx], dma_type, tag[idx], vlane_split_axis, vlane_stride + : src_memref, dst_memref, memref<1xi32> {dram_stride, sram_stride, padding} + +Pass interface (passes/__init__.py): MARKERS + run(module). +""" + +OP_NAME = "togsim.transfer" +MARKERS = (OP_NAME,) + + +def _iter_ops(block): + for op in list(block.operations): + yield op + for region in op.operation.regions: + for b in region.blocks: + yield from _iter_ops(b) + + +def _int_array(attr): + from mlir.ir import ArrayAttr, IntegerAttr + return [IntegerAttr(a).value for a in ArrayAttr(attr)] + + +def _squeeze_reassociation(shape): + """Group source dims so each group's product is one effective (non-unit) dim; + unit dims attach to a neighbor. Returns (groups, target_shape).""" + groups, cur = [], [] + for i, e in enumerate(shape): + cur.append(i) + if e > 1: + groups.append(cur) + cur = [] + if cur: # trailing unit dims + if groups: + groups[-1] += cur + else: + groups.append(cur) # all-ones -> single dim of size 1 + import math + target = [math.prod(shape[d] for d in g) for g in groups] + return groups, target + + +def run(module): + """Lower every togsim.transfer in `module`, in place. Context must be active.""" + import itertools + from mlir.ir import (InsertionPoint, Operation, MemRefType, ArrayAttr, + IntegerAttr, IntegerType, IndexType, DenseI64ArrayAttr, + DenseI32ArrayAttr, StridedLayoutAttr, AffineMap, AffineMapAttr, + AffineExpr) + i64 = IntegerType.get_signless(64) + idx_ty = IndexType.get() + + targets = [] + for region in module.operation.regions: + for b in region.blocks: + for op in _iter_ops(b): + if op.operation.name == OP_NAME: + targets.append(op.operation) + + for op in targets: + dram, dram_idx, sram, sram_idx, tag, dma_type, vst = op.operands + kind = op.attributes["dma_kind"].value # StringAttr -> "MVIN"/"MVOUT" + vlane_axis = IntegerAttr(op.attributes["vlane_split_axis"]).value + dram_stride = _int_array(op.attributes["dram_stride"]) + tile_stride = _int_array(op.attributes["tile_stride"]) + padding = op.attributes["padding"] + + sram_ty = MemRefType(sram.type) + elem, space = sram_ty.element_type, sram_ty.memory_space + tile_shape = list(sram_ty.shape) + # effective (non-unit) dims carry the descriptor; unit dims drop out. + eff = [i for i, e in enumerate(tile_shape) if e > 1] + + def _const(v): + return Operation.create( + "arith.constant", results=[idx_ty], + attributes={"value": IntegerAttr.get(idx_ty, v)}).results[0] + + def _emit(sram_mem, sram_indices, dram_idx_val, vsa_val, dr_attr, tl_attr): + vsa = _const(vsa_val) + if kind == "MVIN": + operands = [dram, dram_idx_val, sram_mem, *sram_indices, + dma_type, tag, sram_idx, vsa, vst] + else: + operands = [sram_mem, *sram_indices, dram, dram_idx_val, + dma_type, tag, sram_idx, vsa, vst] + Operation.create( + "memref.dma_start", results=[], operands=operands, + attributes={"dram_stride": dr_attr, "sram_stride": tl_attr, + "padding": padding}) + + if len(eff) <= 4: + # Fast path: drop unit dims so the descriptor reaches <=4D. The customized + # dma_start convention requires SRAM rank == #indices == len(sram_stride), + # so collapse the unit tile dims away. DRAM stays flat rank-1 (its N-D + # structure is in dram_stride). + groups, target = _squeeze_reassociation(tile_shape) + reassoc = ArrayAttr.get( + [ArrayAttr.get([IntegerAttr.get(i64, d) for d in g]) for g in groups]) + collapsed_ty = MemRefType.get(target, elem, memory_space=space) + keep = [g[-1] for g in groups] # the non-unit dim in each group + dr_attr = ArrayAttr.get([IntegerAttr.get(i64, dram_stride[i]) for i in keep]) + tl_attr = ArrayAttr.get([IntegerAttr.get(i64, tile_stride[i]) for i in keep]) + # Remap vlane axis to the collapsed-dim index (the group containing it). + new_vlane = next(gi for gi, g in enumerate(groups) if vlane_axis in g) + with InsertionPoint(op): + sram_c = Operation.create( + "memref.collapse_shape", results=[collapsed_ty], operands=[sram], + attributes={"reassociation": reassoc}).results[0] + _emit(sram_c, [sram_idx] * len(target), dram_idx, new_vlane, + dr_attr, tl_attr) + op.erase() + continue + + # Peel path: >4 effective dims. Keep the inner 4 as the <=4D descriptor and + # peel the outer (len-4) effective dims into a fully-unrolled set of slices + # (one descriptor per outer index combo; base advances by stride*idx). The + # SRAM slice is a rank-reduced memref.subview at the slice offset; the DRAM + # base advances by a *constant* per slice. + # + # The constant DRAM offset must be folded into an affine.apply over the + # original dram_idx (NOT arith.addi): the TOG pass reads loop_idx_list by + # walking the DRAM index via processDramIndices, which understands + # affine.apply / block-arg / constant but NOT arith.addi -- an addi yields an + # empty loop_idx_list and the kernel fails ONNX serialization (#258). The + # peeled dim itself is a fixed constant in each unrolled slice (this DMA does + # not iterate it), so it correctly contributes no loop var; the surviving + # loop vars come from the original dram_idx affine.apply, into which + # processDramIndices recurses. + peeled, inner = eff[:-4], eff[-4:] + ndim = len(tile_shape) + inner_shape = [tile_shape[d] for d in inner] + inner_strides = [tile_stride[d] for d in inner] + dr_attr = ArrayAttr.get([IntegerAttr.get(i64, dram_stride[d]) for d in inner]) + tl_attr = ArrayAttr.get([IntegerAttr.get(i64, tile_stride[d]) for d in inner]) + # the vlane axis must survive into the inner descriptor (it is the lane dim). + new_vlane = inner.index(vlane_axis) if vlane_axis in inner else 0 + for combo in itertools.product(*[range(tile_shape[d]) for d in peeled]): + static_offsets = [0] * ndim + static_sizes = [1] * ndim + for k, d in enumerate(peeled): + static_offsets[d] = combo[k] + for d in inner: + static_sizes[d] = tile_shape[d] + sram_off = sum(combo[k] * tile_stride[peeled[k]] for k in range(len(peeled))) + dram_off = sum(combo[k] * dram_stride[peeled[k]] for k in range(len(peeled))) + res_ty = MemRefType.get( + inner_shape, elem, + layout=StridedLayoutAttr.get(sram_off, inner_strides), memory_space=space) + with InsertionPoint(op): + sub = Operation.create( + "memref.subview", results=[res_ty], operands=[sram], + attributes={"static_offsets": DenseI64ArrayAttr.get(static_offsets), + "static_sizes": DenseI64ArrayAttr.get(static_sizes), + "static_strides": DenseI64ArrayAttr.get([1] * ndim), + # operandSegmentSizes is an i32 property: [source, offsets, + # sizes, strides] dynamic-operand counts. All static here -> + # only the source operand. Must be i32, not i64 (i64 silently + # zeroes to [0,0,0,0] and fails verification). + "operandSegmentSizes": DenseI32ArrayAttr.get([1, 0, 0, 0])} + ).results[0] + if dram_off == 0: + dram_idx_val = dram_idx + else: + # affine.apply (d0) -> (d0 + dram_off) so TOG's processDramIndices + # recurses through it into the original dram_idx's loop vars. + amap = AffineMap.get(1, 0, [AffineExpr.get_dim(0) + dram_off]) + dram_idx_val = Operation.create( + "affine.apply", results=[idx_ty], operands=[dram_idx], + attributes={"map": AffineMapAttr.get(amap)}).results[0] + _emit(sub, [sram_idx] * 4, dram_idx_val, new_vlane, dr_attr, tl_attr) + op.erase() + + +def lower_text(text: str) -> str: + """Parse `text`, run this pass, return the printed module. CLI/testing helper.""" + if OP_NAME not in text: + return text + from mlir.ir import Context, Module, Location + ctx = Context() + ctx.allow_unregistered_dialects = True + with ctx, Location.unknown(): + m = Module.parse(text) + run(m) + return str(m) + + +if __name__ == "__main__": + import sys + out = lower_text(open(sys.argv[1]).read()) + if len(sys.argv) > 2: + open(sys.argv[2], "w").write(out) + else: + sys.stdout.write(out) diff --git a/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py b/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py new file mode 100644 index 00000000..f5b841bb --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/lower_dma_to_gemmini.py @@ -0,0 +1,227 @@ +"""Lower customized memref.dma_start ops to Gemmini RISC-V inline asm. + +Python port of the C++ test-memref-to-gemmini conversion. Each memref.dma_start +(carrying dram_stride / sram_stride / subtile_size attrs and vlane params encoded +in its stride / num_elements_per_stride / num_elements operands) becomes a +sequence of `llvm.inline_asm` ".insn r CUSTOM_1 ..." Gemmini instructions: +config_mvin/mvout, config2 (dram strides), config3 (spad strides), then the +mvin/mvout itself with the DRAM and scratchpad byte addresses. + +The conversion-framework coupling of the C++ pass (LLVMTypeConverter, +getStridedElementPtr, MemRefDescriptor) is avoided by working at the memref level: +addresses are computed with `memref.extract_aligned_pointer_as_index` + arith, +and the existing standard MLIR->LLVM lowering finalizes everything. Pass order: +this runs on memref-level IR (after test-pytorchsim-to-vcix), before +run_standard_lowering. + +NOTE: indirect-access (gather) dma_start is not yet handled (Phase 2); such ops +raise so they are caught rather than silently mishandled. +""" + +OP_NAME = "memref.dma_start" +WAIT_NAME = "memref.dma_wait" +MARKERS = (OP_NAME, WAIT_NAME) + +# func7 instruction codes (CustomDMAAttribute.h) +CONFIG, CONFIG2, CONFIG3, CONFIG4 = 0, 4, 5, 6 +MVIN, MVIN2, MVIN3, MVOUT = 2, 1, 14, 3 +CONFIG_TYPE = {MVIN: 0, MVIN2: 1, MVIN3: 2, MVOUT: 3} +MAX_TENSOR_DIM = 4 +CONSTRAINTS = "r,r,~{dirflag},~{fpsr},~{flags}" + + +def _asm(func7): + return f".insn r CUSTOM_1, 0x3, {func7}, x0, $0, $1" + + +def _i64_signed(v): + """Wrap an unsigned 64-bit packed value into signed int64 (matches C++ getI64IntegerAttr).""" + v &= 0xFFFFFFFFFFFFFFFF + return v - (1 << 64) if v >= (1 << 63) else v + + +def _row_major_strides(shape): + strides = [1] * len(shape) + for i in range(len(shape) - 2, -1, -1): + strides[i] = strides[i + 1] * shape[i + 1] + return strides + + +def run(module, timing=False): + """Lower memref.dma_start / dma_wait to Gemmini instructions. + + timing=False (functional/Spike): dma_start -> gemmini config + mvin/mvout asm. + timing=True (gem5 cycle path): dma_start is erased (the TOG already carries + DMA timing; the cycle binary needs no asm). + memref.dma_wait is erased in both modes (matches C++ DmaWaitOpLowering). + """ + from mlir.ir import (InsertionPoint, Operation, IntegerType, IndexType, + IntegerAttr, MemRefType) + from mlir.dialects import llvm, arith, memref + + i64 = IntegerType.get_signless(64) + idx = IndexType.get() + + def const_int(val): + return IntegerAttr(val.owner.attributes["value"]).value + + def i64_const(value): + return arith.ConstantOp(i64, IntegerAttr.get(i64, _i64_signed(value))).result + + def asm(func7, rs1, rs2): + llvm.InlineAsmOp(None, [rs1, rs2], _asm(func7), CONSTRAINTS, + has_side_effects=True, asm_dialect=0) + + def elem_addr_i64(memref_val, indices, mtype, elem_bytes): + """i64 byte address of memref_val[indices] (aligned ptr + linear elem offset).""" + base = memref.ExtractAlignedPointerAsIndexOp(memref_val).result # index = byte addr + strides = _row_major_strides(list(mtype.shape)) + off = None # element offset (index) + for k, ival in enumerate(indices): + if strides[k] == 0: + continue + term = ival + if strides[k] != 1: + term = arith.MulIOp(ival, arith.ConstantOp(idx, IntegerAttr.get(idx, strides[k])).result).result + off = term if off is None else arith.AddIOp(off, term).result + if off is not None: + byte = arith.MulIOp(off, arith.ConstantOp(idx, IntegerAttr.get(idx, elem_bytes)).result).result + base = arith.AddIOp(base, byte).result + return arith.IndexCastOp(i64, base).result + + starts, waits = [], [] + for region in module.operation.regions: + for b in region.blocks: + _collect(b, starts, waits) + + for op in waits: # dma_wait: erase in both modes + op.erase() + + for op in starts: + if timing: # gem5 cycle path: drop the dma_start (TOG has timing) + op.erase() + continue + operands = list(op.operands) + src, dst = operands[0], None + src_ty = MemRefType(src.type) + src_rank = len(src_ty.shape) + dst = operands[1 + src_rank] + dst_ty = MemRefType(dst.type) + dst_rank = len(dst_ty.shape) + src_idx = operands[1:1 + src_rank] + dst_idx = operands[1 + src_rank + 1:1 + src_rank + 1 + dst_rank] + + dma_type = const_int(operands[1 + src_rank + 1 + dst_rank]) # num_elements + vlane_split_axis = const_int(operands[-2]) # stride (always 2nd-to-last) + vlane_stride = const_int(operands[-1]) & 0x7FFF # num_elements_per_stride (last) + is_mvin = dma_type in (MVIN, MVIN2, MVIN3) + + elem_bytes = _elem_bytes(src_ty.element_type) + # Indirect (gather): the gather-side indices are src for mvin, dst for mvout. + gather_idx = src_idx if is_mvin else dst_idx + indirect, indirect_memref = _find_indirect(gather_idx) + + tile_shape = _subtile(op) + if tile_shape is None: + tile_shape = list(dst_ty.shape) if is_mvin else list(src_ty.shape) + dram_strides = _int_array(op, "dram_stride") + spad_strides = _int_array(op, "sram_stride") + assert len(tile_shape) == len(dram_strides) == len(spad_strides), \ + f"shape/stride rank mismatch: {tile_shape} {dram_strides} {spad_strides}" + + expand = MAX_TENSOR_DIM - len(tile_shape) + shape4 = [1] * expand + tile_shape + dram4 = [0] * expand + dram_strides + spad4 = [0] * expand + spad_strides + vlane_split_axis += expand + config_type = CONFIG_TYPE[dma_type] + + with InsertionPoint(op): + addrA = elem_addr_i64(src, src_idx, src_ty, elem_bytes) + addrB = elem_addr_i64(dst, dst_idx, dst_ty, elem_bytes) + dram_addr, spad_addr = (addrA, addrB) if is_mvin else (addrB, addrA) + + cfg_rs1 = i64_const(((shape4[0] & 0xFFFF) << 48) | ((shape4[1] & 0xFFFF) << 32) + | ((shape4[2] & 0xFFFF) << 16) | (shape4[3] & 0xFFFF)) + cfg_rs2 = i64_const((vlane_stride << 32) | ((config_type & 0x3) << 17) + | ((1 if indirect else 0) << 16) + | ((vlane_split_axis & 0x3) << 14) | elem_bytes) + asm(CONFIG, cfg_rs1, cfg_rs2) + asm(CONFIG2, i64_const((dram4[0] << 32) | (dram4[1] & 0xFFFFFFFF)), + i64_const((dram4[2] << 32) | (dram4[3] & 0xFFFFFFFF))) + asm(CONFIG3, i64_const((spad4[0] << 32) | (spad4[1] & 0xFFFFFFFF)), + i64_const((spad4[2] << 32) | (spad4[3] & 0xFFFFFFFF))) + if indirect: + # CONFIG4: rs1 = indirect index-spad base address, rs2 = (elem_size<<16)|stride(1) + ind_base = memref.ExtractAlignedPointerAsIndexOp(indirect_memref).result + ind_addr = arith.IndexCastOp(i64, ind_base).result + ind_esize = _elem_bytes(MemRefType(indirect_memref.type).element_type) + asm(CONFIG4, ind_addr, i64_const(((ind_esize & 0xFF) << 16) | (1 & 0xFFFF))) + asm(dma_type, dram_addr, spad_addr) + op.erase() + + +def _collect(block, starts, waits): + for op in list(block.operations): + name = op.operation.name + if name == OP_NAME: + starts.append(op.operation) + elif name == WAIT_NAME: + waits.append(op.operation) + for region in op.operation.regions: + for b in region.blocks: + _collect(b, starts, waits) + + +def _subtile(op): + from mlir.ir import ArrayAttr, IntegerAttr + if "subtile_size" not in op.attributes: + return None + return [IntegerAttr(a).value for a in ArrayAttr(op.attributes["subtile_size"])] + + +def _int_array(op, name): + from mlir.ir import ArrayAttr, IntegerAttr + return [IntegerAttr(a).value for a in ArrayAttr(op.attributes[name])] + + +def _elem_bytes(elem_type): + from mlir.ir import IntegerType, FloatType + bits = (IntegerType(elem_type).width if IntegerType.isinstance(elem_type) + else FloatType(elem_type).width) + return max(bits, 8) // 8 + + +def _find_indirect(indices): + """If a gather index is an affine.apply{indirect_access} whose operands include + index_cast(affine.load(%spad)), return (True, %spad memref); else (False, None).""" + for idx in indices: + ap = idx.owner + if getattr(ap, "name", None) != "affine.apply" or "indirect_access" not in ap.attributes: + continue + for operand in ap.operands: + ic = operand.owner + if getattr(ic, "name", None) != "arith.index_cast": + continue + ld = ic.operands[0].owner + if getattr(ld, "name", None) == "affine.load": + return True, ld.operands[0] # affine.load operand 0 == the index spad memref + return False, None + + +def lower_text(text): + if OP_NAME not in text: + return text + from mlir.ir import Context, Module, Location + ctx = Context() + ctx.allow_unregistered_dialects = True + with ctx, Location.unknown(): + m = Module.parse(text) + run(m) + return str(m) + + +if __name__ == "__main__": + import sys + out = lower_text(open(sys.argv[1]).read()) + (open(sys.argv[2], "w").write(out) if len(sys.argv) > 2 else sys.stdout.write(out)) diff --git a/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py b/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py new file mode 100644 index 00000000..5cd16e18 --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/lower_to_llvm.py @@ -0,0 +1,68 @@ +"""Standard MLIR -> LLVM-dialect lowering via the bindings PassManager. + +Runs the upstream *registered* lowering passes (convert-*-to-llvm, lower-affine, +reconcile-unrealized-casts, ...) in-process on the post-custom-pass IR, replacing +the tail of the mlir-opt pipeline. The custom passes (test-loop-padding, +dma-fine-grained, test-pytorchsim-to-vcix, test-tile-operation-graph, +test-memref-to-gemmini) still run in mlir-opt; this picks up right after +memref-to-gemmini. As those custom passes migrate to Python, mlir-opt shrinks and +eventually this becomes the whole back half of an all-in-process flow. + +Validated to produce byte-identical LLVM IR to running the same passes inside +mlir-opt. Note: only lower-vector-multi-reduction is func.func-scoped (the +bindings pass-pipeline parser does not auto-nest like the mlir-opt CLI, so it is +wrapped explicitly); order is preserved to match the original pipeline. +""" + +STANDARD_PIPELINE = ( + "builtin.module(" + "convert-linalg-to-loops," + "convert-vector-to-scf{full-unroll=true}," + "lower-affine," + "expand-strided-metadata," # decompose memref.collapse_shape/subview before LLVM + "finalize-memref-to-llvm," + "func.func(lower-vector-multi-reduction)," + "convert-vector-to-llvm," + "convert-arith-to-llvm," + "convert-math-to-llvm," + "convert-scf-to-cf," + "convert-cf-to-llvm," + "convert-func-to-llvm," + "convert-index-to-llvm," + "reconcile-unrealized-casts)" +) + + +def run_standard_lowering(in_path, out_path=None, timing=False): + """Lower the post-custom-pass MLIR at `in_path` to the LLVM dialect. + + Runs the imperative Gemmini lowering (memref.dma_start/dma_wait) then the + registered standard MLIR->LLVM passes. `timing` selects the Gemmini behavior: + False for the functional/Spike path (emit gemmini asm), True for the gem5 + cycle path (erase dma_start; the TOG already carries DMA timing) -- this + preserves the old test-memref-to-gemmini `timing=1` semantics. + + Writes the result to `out_path` (defaults to `in_path`, i.e. in place). + Requires the MLIR Python bindings on PYTHONPATH. + """ + if out_path is None: + out_path = in_path + from mlir.ir import Context, Module, Location + from mlir.passmanager import PassManager + from . import lower_dma_to_gemmini + ctx = Context() + ctx.allow_unregistered_dialects = True + with ctx, Location.unknown(): + with open(in_path) as f: + module = Module.parse(f.read()) + # Imperative Python pass: memref.dma_start/dma_wait -> Gemmini asm (replaces + # the C++ test-memref-to-gemmini), then the registered standard lowering. + lower_dma_to_gemmini.run(module, timing=timing) + PassManager.parse(STANDARD_PIPELINE, ctx).run(module.operation) + with open(out_path, "w") as f: + f.write(str(module)) + + +if __name__ == "__main__": + import sys + run_standard_lowering(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else None) diff --git a/PyTorchSimFrontend/mlir/passes/lower_vlane_idx.py b/PyTorchSimFrontend/mlir/passes/lower_vlane_idx.py new file mode 100644 index 00000000..c9898f4b --- /dev/null +++ b/PyTorchSimFrontend/mlir/passes/lower_vlane_idx.py @@ -0,0 +1,92 @@ +"""Python out-of-line MLIR pass: lower torchsim.vlane_idx -> per-lane index * offset. + +Codegen emits a dedicated `torchsim.vlane_idx` op (generic form, unregistered +dialect) carrying a `vlane_offset` integer attribute. This pass rewrites each +such op to: + + %v = "vcix.v.i"(%K) {opcode = 0, rs2 = 0, imm = 0} : (i64) -> vector // per-lane index + %n = arith.constant dense : vector + %r = arith.muli %v, %n : vector + +and replaces uses of the original op with %r. Replaces the former C++ +`-global-idx` pass (which overloaded arith.addi with a vlane_offset attribute). + +Pass interface (see passes/__init__.py): MARKERS + run(module). Also runnable +standalone as a CLI: + python PyTorchSimFrontend/mlir/passes/lower_vlane_idx.py in.mlir [out.mlir] + +Requires the MLIR Python bindings on PYTHONPATH +(/riscv-llvm/python_packages/mlir_core). The `vcix` dialect must be registered +in the consuming mlir-opt for the result to round-trip (see +registerVCIXDialectTranslation in mlir-opt.cpp). +""" + +OP_NAME = "torchsim.vlane_idx" +MARKERS = (OP_NAME,) + + +def _iter_ops(block): + for op in list(block.operations): + yield op + for region in op.operation.regions: + for b in region.blocks: + yield from _iter_ops(b) + + +def run(module): + """Lower every torchsim.vlane_idx op in `module`, in place. + + Must be called with the module's Context active (the orchestrator provides it). + """ + from mlir.ir import (InsertionPoint, Operation, IntegerType, IntegerAttr, + DenseElementsAttr, VectorType) + i64 = IntegerType.get_signless(64) + i32 = IntegerType.get_signless(32) + + targets = [] + for region in module.operation.regions: + for b in region.blocks: + for op in _iter_ops(b): + if op.operation.name == OP_NAME: + targets.append(op.operation) + + for op in targets: + res = op.results[0] + vt = VectorType(res.type) + k, et = vt.shape[0], vt.element_type + offset = IntegerAttr(op.attributes["vlane_offset"]).value + with InsertionPoint(op): + rvl = Operation.create("arith.constant", results=[i64], + attributes={"value": IntegerAttr.get(i64, k)}).results[0] + lane = Operation.create("vcix.v.i", results=[vt], operands=[rvl], + attributes={"opcode": IntegerAttr.get(i64, 0), + "rs2": IntegerAttr.get(i32, 0), + "imm": IntegerAttr.get(i32, 0)}).results[0] + ovec = Operation.create("arith.constant", results=[vt], + attributes={"value": DenseElementsAttr.get_splat( + vt, IntegerAttr.get(et, offset))}).results[0] + mul = Operation.create("arith.muli", results=[vt], operands=[lane, ovec]).results[0] + res.replace_all_uses_with(mul) + op.erase() + + +def lower_text(text: str) -> str: + """Parse `text`, run this pass, return the printed module. CLI/testing helper.""" + if OP_NAME not in text: + return text + from mlir.ir import Context, Module, Location + ctx = Context() + ctx.allow_unregistered_dialects = True + with ctx, Location.unknown(): + m = Module.parse(text) + run(m) + return str(m) + + +if __name__ == "__main__": + import sys + out = lower_text(open(sys.argv[1]).read()) + if len(sys.argv) > 2: + open(sys.argv[2], "w").write(out) + else: + sys.stdout.write(out) diff --git a/docs/axis-split-scheduling.md b/docs/axis-split-scheduling.md new file mode 100644 index 00000000..10171ab4 --- /dev/null +++ b/docs/axis-split-scheduling.md @@ -0,0 +1,211 @@ +# Aligned axis splitting at the Inductor scheduling layer + +Status: **prototype / proposed**. Companion to `dma-transfer-lowering.md`. This doc +covers the *upstream* half of the affine-only contract: removing aligned +`FloorDiv` / `ModularIndexing` from index expressions before they reach MLIR +codegen, by splitting loop axes at the Inductor scheduling layer. + +## Goal: the affine-only contract + +We want the MLIR codegen (`get_dma_info` in `mlir_codegen_backend.py`) to receive +only per-axis affine index expressions: + + off(i,j,k,...) = base + Sum_k stride_k * loop_var_k (stride_k constant int) + +with **zero** `FloorDiv` / `ModularIndexing`. If that invariant holds, codegen no +longer fights non-affine indices: the recompile dance (RecompileSignal, forced +tile sizes, max_retry_compile), the heuristic `TestLoopPadding` pass, and the +hard-fail-on-conflict path all become unnecessary. Codegen's only remaining job +is the *mechanical* rank<=4 peel for the Gemmini descriptor (orthogonal; see +`dma-transfer-lowering.md`), which operates on already-affine input. + +Two tools produce this invariant, matching the alignment theory: + +- **aligned floor/mod -> axis split** (this doc): loop transformation, free, no + data movement. +- **misaligned floor/mod -> graph copy insertion** (XLA-style): genuine data + movement; out of scope here. + +"Aligned" means the floor/mod argument is a *single* iteration variable `v` of +extent `E` and the divisor `k` (resp. `k*m` for ModularIndexing) divides `E`, so +splitting `v = outer*k + inner` lands the wrap point on a fixed axis boundary. + +## Where: the scheduling layer already rebuilds LoopBody + +`mlir_scheduling.py` already does loop-IR surgery at the scheduling layer: + +- `revert_group` (line ~219) rebuilds a `LoopBody` from `get_store_function()` + with a chosen `var_ranges` -- it undoes Inductor's `simplify_and_reorder`. +- `codegen_node` (line ~246) injects dummy size-1 loops when Inductor + over-simplified the group. + +Axis splitting is the same operation with a different `var_ranges`: split the +axes carrying aligned floor/mod, then rebuild. No new infrastructure -- reuse +`LoopBody`. This is "upstream" of MLIR codegen and native to Inductor's IR (sympy +ranges + index exprs), so we are not reverse-engineering MLIR text. + +## How: detect / rebuild / hook + +Implemented in `PyTorchSimFrontend/mlir/axis_split.py`, wired into +`codegen_node` behind `TORCHSIM_AXIS_SPLIT=1` (dump with +`TORCHSIM_DEBUG_AXIS_SPLIT=1`). + +1. **Detect -- `find_split_plan(nodes)`**: scan each node's + `_body.indexing_exprs` for `FloorDiv(v, k)` / `ModularIndexing(v, k, m)` where + `v` is a single iter var and the divisor divides `v`'s extent. Return + `{axis_index: divisor}`, keyed positionally so it applies to every fused node + sharing the iteration space. +2. **Rebuild -- `build_split_body(node, plan)`**: rebuild `node._body` / + `_sizes` with the split var_ranges; feed the store function the index + expression `outer*k + inner` at the split dim so the floor/mod collapses. +3. **Hook -- `codegen_node`**: apply the plan to every node + (`_sizes, _body, group = ...`), then recompute the group. + +## Empirical validation (group norm) + +`group_norm(x[2,6,4,4], num_groups=3)` normalize kernel, before vs after split: + + before var_ranges={p0:2, p1:6, p2:16} + idx0 = 96*p0 + 16*p1 + p2 # x input, affine + idx1 = 3*p0 + (p1//2) # mean/rstd <- FloorDiv(p1,2), 2|6 aligned + idx2 = p1 # weight/bias, affine + + after plan={1: 2}, var_ranges={s0:2, s1:3, s2:2, ...} + idx1 = 3*s0 + (s1//1) # FloorDiv collapsed to identity -> s1 + ... # mean now affine; s2/spatial broadcast (stride 0) + +The FloorDiv is eliminated. group `(2,6,16) -> (2,3,2,...)`. + +## Coverage (what this framework can and cannot do) + +| Case | Example | Status | +|---|---|---| +| aligned FloorDiv, single var | group norm `c//2` (2\|6) | DONE (prototype) | +| aligned ModularIndexing | `(v//k)%m`, k*m\|E | needs mixed-radix multi-split | +| multiple radices on one axis | `//2` + `%3`, E=6 | needs nested split (now: first divisor only) | +| reduction-axis floor/mod | `r//k` inside reduce | needs reduction-var splitting | +| divisor does not divide extent | C=8 groups of 3; uneven cat | IMPOSSIBLE by split -> graph copy | +| multi-axis argument | `(4p+q)//6` non-factor reshape | IMPOSSIBLE by split -> graph copy | +| dynamic / symbolic | `v//s`, symbolic extent | separate symbolic/guard path | + +The aligned class is the framework's domain (currently only single-split +FloorDiv); the misaligned class is structurally a graph-copy problem. + +## Resolved + +- **5D blow-up (fixed).** `build_split_body` now reindexes the already-collapsed + `node._body` via `LoopBody`'s copy path (pass the body as `fn` -> + `_init_with_copy`), instead of re-tracing the raw store function over + `inode.data.get_size()`. This keeps merged dims merged (spatial stays `16`), + so group_norm goes `(2,6,16) -> (2,3,2,16)` (4D, no cap hit), and + `_init_with_copy`'s `simplify_with_ranges` folds the split floor. +- **`floor//1` residue (fixed).** The fold only happened once the new symbols + carried integer/non-negative assumptions: build them with + `torch._inductor.utils.sympy_index_symbol` (not bare `sympy.Symbol`), which is + also why the index prefix must not be `s` (reserved for shape symbols). With + this, `idx1 = 3*p0 + (p1//2)` becomes `3*z0 + z1` -- the channel FloorDiv is + gone, not left as `z1//1`. +- **Symbol conventions.** Index dims use the `z` prefix; reduction dims use the + `r` prefix and are kept after the index dims so the reduction axis stays + innermost (`var_ranges` is ordered iter-then-reduce; `LoopBody.sizes` splits on + `len(iter_vars)`). LoopBody var names are remapped to `index` during MLIR + codegen, so the prefix is internal -- but it must not collide with the original + body's names (those are `p`/`q`, so `z`/`r` are safe). + +## Resolved (cont.) + +- **`floor//1` / residual floor on multi-level split (fixed).** `simplify_with_ranges` + cannot prove a *multi-term* numerator is below the divisor (e.g. + `FloorDiv(z1 + 4*z2, 12)` with `z1<4, z2<3`), so a 3-level mixed-radix split left + a residual floor that codegen rejected ("Not supporting this view operation"). + `_fold_with_ranges` now proves it directly from the split sub-var ranges via + `bound_sympy`: `FloorDiv(num,d)->0` when `0<=numnum//k` + when `0<=num 5D), which triggers the nascent + decompose-transfer peel + TOG path (see below). `find_split_plan` now has a rank + guard: if applying the plan would make the index rank exceed 4, the whole plan is + dropped and the kernel falls back to baseline. pixel_shuffle now passes (via + baseline); 3D group_norm still splits (rank 4, allowed). + +## Known issues / open + +- **decompose-transfer peel <-> TOG incompatibility**: the >4D peel emits + `memref.subview` + unrolled constant-offset `dma_start`, which the C++ TOG + generation pass cannot read (empty `loop_idx_list`). The rank guard above + side-steps it; the real fix is to rewrite the peel as an `affine.for` loop + (keeping a loop index TOG can read) instead of unrolling. **Tracked as a GitHub + issue + the `dma-transfer-lowering.md` TODO.** + +## Done + +- **Mixed-radix (ModularIndexing + multi-radix)**: `find_split_plan` returns a + per-axis divisibility-chain of boundaries; `build_split_body` splits into one + sub-var per segment (`v = sum_i d_i*b_i`). Validated allclose=True on group_norm + (FloorDiv, `[1,2,6]`) and `x.repeat(1,2)` (single-axis ModularIndexing, + `[1,8,16]`); pixel_shuffle (floor+mod on two axes) linearizes correctly. +- **Reduction pass-through**: reduction dims keep the `r` prefix and stay innermost + (after the index dims). Exercised via the `TORCHSIM_AXIS_SPLIT_FORCE` validation + gate (force-split a reduction kernel's index axis even without floor -- an + identity transform, so allclose must hold): layernorm `(512)->(256,2)` and + reduce `(68)->(34,2)` keep their reduction groups and pass. +- **Graph-copy for incompatible radices (case 5)** -- `graph_copy.py`, + `TORCHSIM_GRAPH_COPY`. When two operands of an elementwise consumer carry + incompatible-radix groupings on a shared axis (e.g. `a[c//2] + b[c%3]`, floor-by-2 + vs mod-by-3 on extent 6 -- not a divisibility chain), neither axis-split nor the + recompile-dance can express it. We wrap the registered lowering entries (the + make_pointwise results = every elementwise consumer, one place), trace each + operand's loader with `extract_read_writes` to get its read indices, run the same + `collect_boundaries` analysis, and if the union is not a chain, `realize()` the + cheaper operand. realize() (not clone -- Inductor inlines clone, confirmed) forces + a buffer: the consumer then reads it affine and the remaining single grouping is + handled by axis-split. Validated: `incompat` (`a.repeat_interleave(2)+b.repeat(2)`) + goes ERR -> allclose=True with `GRAPH_COPY+AXIS_SPLIT` (still ERR on default, + confirming graph-copy is the fix); no regression on the pattern battery, + test_add, resnet (compile overhead negligible). +- **Graph-copy for cross-axis floor/mod (case 7)** -- same hook. A transpose+reshape + feeding a consumer that keeps the output dims separate (broadcast / softmax / + layernorm / reduce-one-dim) produces a floor/mod whose argument spans *two* loop + vars, e.g. `(3*p0+p1)//4`; axis-split cannot split a multi-var argument. We detect + an operand whose read index has a floor/mod argument with >1 free symbol and + replace it with `ExternKernel.copy_input` (a realized identity Pointwise). This is + why copy_input and not `realize()`: `StorageBox.realize()` is a no-op on a + ReinterpretView (a reshape), so it does not materialize view operands; copy_input + forces the copy. The copy kernel iterates the operand's own contiguous shape, so + its index collapses to single-var for axis-split, and the consumer reads the copy + affine. Also covers single-operand consumers (a reduction reading a multi-var + view). Validated allclose=True: reshape+broadcast, softmax(reshape), + layernorm(reshape) (all ERR on default). NOTE the empirical correction: case 7 is + NOT rare -- it is the common attention/norm "reshape then reduce/broadcast" + shape; Inductor only avoids it when it can collapse the output to 1D (then the + floor is single-var). + +## Default-on + recompile-dance status + +axis-split and graph-copy are **ON by default** (disable with `TORCHSIM_AXIS_SPLIT=0` +/ `TORCHSIM_GRAPH_COPY=0`). With them on, the codegen recompile-dance (tile-forcing +for floor/mod divisibility) is demoted from primary mechanism to a rarely-hit +fallback. + +Measured under default-on (`TORCHSIM_RECOMPILE_LOG=1`), 33 tests, all pass: +- 16 core (elementwise/gemm/reduce/conv/view/fusion + mlp/resnet/transformer/vit): 0 recompiles. +- 7 broader families (cnn/pool/group_conv/sort/indirect_access/exponent/conv_fusion): 0 recompiles. +- 10 floor/mod patterns: 1 recompile total (an unrelated tile-divisibility in the + 3-level mixed-radix case). + +**Full retirement of the dance is deferred** (it is still a real dependency, not +just a safety net): removing the floor/mod recompile branches would break the +3-level mixed-radix case (1 recompile) and any case axis-split/graph-copy do not +yet cover (case 6, >4D rank-guard skips). attention/sdpa families were not run here +(too slow locally) and need CI validation before retirement. + +## Next steps + +1. Eliminate the last recompile dependency (the 3-level mixed-radix sub-kernel) so + the dance reaches 0/all -> then retire the floor/mod recompile branches (keep the + non-floor/mod ones: non-power-of-2 vec size, indirect). +2. Graph-copy coverage: case 6 (non-dividing divisor / uneven cat -> pad or gather), + and conflicts internal to templates (gemm/conv/sdpa). +3. High-rank interaction: cap split-induced rank or harden decompose-peel + TOG for + high-rank tiles (pixel_shuffle end-to-end, #258). +4. Dynamic shapes -> symbolic divisibility / guards. diff --git a/docs/dma-transfer-lowering.md b/docs/dma-transfer-lowering.md new file mode 100644 index 00000000..cbf875c0 --- /dev/null +++ b/docs/dma-transfer-lowering.md @@ -0,0 +1,478 @@ +# DMA transfer op + decomposition lowering + +Status: **design / proposed**. Captures the plan to fix the recompile-dance +fragility by representing DMA as a high-level declarative transfer op and +decomposing it into affine descriptors in a lowering pass. + +Companion docs: `linalg-codegen-migration.md` (Plan B, the full structured-ops +rewrite this is a narrow tactical slice of). The near-term graph-level padding +work is referred to here as Plan A. + +## TL;DR + +The MLIR codegen forces tile sizes so that non-affine index expressions +(`FloorDiv` / `ModularIndexing`, produced by view/reshape/cat) collapse into the +DMA's strictly-affine 4D integer-stride address model. That forcing is a lazy, +greedy, monotonic, restart-based search (the "recompile dance") capped at 5 +retries; when operand constraints conflict or exceed 4D it hard-fails and the +model does not compile. This blocks model coverage, which is the primary goal. + +Proposed fix: stop forcing one affine descriptor. Introduce a **high-level +`togsim.transfer` op** that carries an iteration domain plus `iter->src` / +`iter->dst` affine maps (which may legally contain floordiv/mod), and a +**decomposition pass** that lowers it to a loop of the **existing customized +`memref.dma_start` descriptors** (kept unchanged as the leaf). The non-affine / +high-rank part is peeled into a base-pointer loop instead of being crammed into +one descriptor. This inverts the tile<->DMA dependency (the DMA adapts to the +tile, not the reverse), removes the rank cap, and removes the recompile dance. + +## Problem + +### Root cause: an impedance mismatch + +The DMA address model is `base + sum_i stride_i * idx_i`, with **integer strides, +4D**, i.e. strictly affine/linear. Inductor's index expressions are not: views, +reshapes, `cat`, and broadcasts introduce `FloorDiv` / `ModularIndexing`, which +are non-affine. The codegen copes by searching for a tiling under which the +floor/mod collapses to a linear stride within a tile -- that is exactly what the +ModularIndexing tile constraints ("tile must be a multiple of the floordiv +divisor and a divisor of the modular divisor") encode. + +### The recompile dance (where it breaks) + +`codegen_nodes` (`mlir_common.py`) is a `while True` loop, `max_retry_compile = 5`. +During emission, `get_dma_info` (`mlir_codegen_backend.py`) inspects the index: + +- the **split path** (good): `apply_divisor(axis, divisor, "split")` peels an axis + into two affine dims to represent floor/mod, inserting a `0` into `dram_stride`; +- the **pad path** (fragile): when the tile is not divisible it mutates the tile + (`set_tile_size`, `tile_constraint.fixed = True`) and raises `RecompileSignal` + to restart emission. + +It breaks because: + +1. **One global tile must satisfy every operand's divisibility** on a shared axis. + Fused ops with conflicting constraints (common with reshape/modular indexing) + cannot be satisfied at once -- this is the loop<->tensor mismatch. +2. **Greedy + monotonic + no backtracking + 5-retry cap.** `tile_constraint.fixed` + persists across retries (the tile descriptor lives on `kernel_group`, survives + `reset`), so the search ratchets one way; conflicting fixes oscillate and hit + the cap -> `RuntimeError("Failed to compile kernel after multiple attempts")`. +3. **4D rank cap.** A reshape needing more than 4 affine dims after splitting + raises `NotImplementedError`. +4. **vlane / LMUL entanglement.** Pad-forcing moves `vlane_split_axis` / relaxes + `vlane_stride`, and `compute_vec` must be a power of two; these can be mutually + unsatisfiable with the divisibility constraints. + +Padding logic is currently spread across three places: the Python recompile/tile +-adjust dance (1), the Python `get_mask` vector-tail handling (2), and the MLIR +`TestLoopPadding` pass (3). Removing (3) alone does not fix the fragility; (1) is +the larger source. + +### We already have a de-facto custom DMA op + +`get_dma_code` emits `memref.dma_start` overloaded via string formatting with +extra operands (`dma_type` MVIN/MVOUT, tag, `vlane_split_axis`, `vlane_stride`) +and extra attributes (`dram_stride`, `tile_stride`, padding type). This is a +custom descriptor in all but name -- and it is what Spike / gem5 / TOGSim already +consume. + +## Proposed design + +Two op levels, with a pass bridging them. + +``` +[high] togsim.transfer iteration domain + iter->src / iter->dst affine maps + (maps MAY contain floordiv/mod; rank unbounded) + | decompose-transfer pass (cost-aware peel) + v +[low] scf.for { customized memref.dma_start } <- existing leaf, UNCHANGED +``` + +### Low-level descriptor (keep as-is) + +The existing customized `memref.dma_start` is the lowering target / leaf: affine, +4D, integer stride, simulator-understood. **Do not add maps or floor/mod to it** -- +that would re-create the representational limit and blur the boundary. Optionally +formalize it into a real op (`togsim.dma_descriptor`) with a verifier so the pass +rewrites real ops instead of strings; not required to start. + +### High-level transfer op (new) + +Strawman: + +```mlir +togsim.transfer + ins(%src : memref) // DRAM + outs(%dst : memref<...xf16, 1>) // scratchpad + iter_bounds = [%M, %N, %K] // iteration domain (dynamic via SSA operands) + attributes { + src_map = affine_map<(m,n,k)[s0] -> (m, (n floordiv s0), (n mod s0), k)>, // non-affine lives here + dst_map = affine_map<(m,n,k) -> (m, n, k)>, + vlane_split_axis = 1, vlane_stride = 4, + dma_kind = "MVIN", tag_policy = "async", + peel_plan = [0] // optional: which iter dims to peel (decided in Python; see below) + } +``` + +Design choices: + +- **Iteration domain + two maps, not a single src->dst map.** A direct src->dst + relation only works for bijections; broadcast (`cat([a, a])`) and non-bijective + access need the loop-mediated form. This is the `linalg.generic` model, and peel + becomes "tile the iteration domain." +- **floor/mod ride in the `AffineMap`.** MLIR `AffineMap` supports constant-divisor + `floordiv`/`mod`/`ceildiv` natively; the codegen already produces these as + strings in `convert_index`. Symbolic divisors are semi-affine -- representable, + handled by our pass. +- **`memref`, not raw pointers.** The memref carries base + shape + layout so the + pass can reason about strides; `src_ptr`/`dst_ptr` are inside it. Buffer shape is + the memref type; the transfer region is `iter_bounds` + maps. +- **Closest existing op is `linalg.generic` / `linalg.copy`** (same shape + maps + + body structure) but it lacks DMA/vlane/tag/scratchpad semantics and its tiling + lowers to subview+scf, not our descriptors -- so a custom op modeled on linalg's + design, reusing AffineMap utilities. + +### Decomposition pass (contract): aligned-only mechanical peel + +> **Scope decision (narrowed).** This pass is a **pure mechanical rank peel** of an +> already-affine access. It does **not** linearize floor/mod and does **not** do +> relayout. Those two responsibilities moved upstream (see "Division of labor" +> below): aligned floor/mod is removed by **axis splitting at the Inductor +> scheduling layer** (`axis-split-scheduling.md`), and misaligned access is +> resolved by **graph-level copy insertion**. So every `togsim.transfer` that +> reaches this pass is guaranteed per-axis affine; the only thing left is that its +> rank may exceed the 4D Gemmini descriptor. + +The DMA descriptor is an **affine map of rank <= 4 with integer strides** +(`base + sum_i stride_i * idx_i`). The pass sees affine input (rank `D`) and: + +1. **`D <= 4`** -> emit **one** customized `memref.dma_start`; the dims become the + descriptor's <=4D shape/strides. Identical to today's output (fast path). +2. **`D > 4`** -> peel `D - 4` dims into an outer `affine.for`; each iteration + computes a base with `affine.apply` (the peeled dims' linear contribution) and + issues the inner <=4D affine descriptor. SRAM offsets are computed symmetrically + in the same loop. + +That is the whole pass. There is **no linearization step** (upstream guarantees +affine) and **no relayout fallback** (upstream graph copy handles misalignment). + +**Fail loud, not silent.** If the pass encounters floor/mod that does not reduce to +per-axis affine (misaligned), or a genuinely non-affine / indirect / gather index, +that is a **contract violation** -- upstream did not normalize it. The pass +**asserts/errors** rather than silently inserting a relayout. A silent in-pass copy +would be a hidden performance cliff and would duplicate, at the wrong layer, a +global layout decision only the graph can make correctly. + +The decision point maps onto existing code: `get_dma_info` already raises at >4D. +That exact site becomes "emit `togsim.transfer`" (done, Phase 1), and this pass +consumes it. The recompile/tile-forcing dance is unnecessary because (a) aligned +floor/mod is gone before codegen and (b) the outer peel loop's `ceil` bound absorbs +non-divisible remainders. + +### Division of labor (the affine-only contract) + +| floor/mod source | handled by | cost | layer | +|---|---|---|---| +| aligned (single axis, divisor \| extent; group norm, broadcast) | axis split | free | Inductor scheduling | +| misaligned (uneven cat, non-factor reshape, multi-axis arg) | copy insertion | copy | FX graph | +| affine but rank > 4 (e.g. 5D permute) | mechanical peel | free | **this pass** | +| data-dependent / indirect / gather | indirect-indexing path | -- | out of scope | + +Only the third row is this pass. The first two produce the affine-only invariant +this pass relies on. + +### Relationship to memref-to-gemmini (ISA lowering) -- keep separate + +`memref.dma_start` is the boundary, not the endpoint. The layering is: + + togsim.transfer --[Python decompose]--> memref.dma_start --[Python lower_dma_to_gemmini]--> Gemmini ISA + +decompose-transfer stops at `memref.dma_start` and must **not** emit Gemmini +instructions directly; the ISA encoding is a separate pass +(`passes/lower_dma_to_gemmini.py`, which replaced the C++ test-memref-to-gemmini). +Rationale: + +- **Separation of concerns**: decompose does descriptor decomposition (affine + algebra: rank / peel); gemmini does instruction encoding (hardware). Different + axes; merging couples affine logic with ISA detail. They stay distinct passes. +- **`memref.dma_start` is a shared contract** with multiple consumers + (lower_dma_to_gemmini, dma-fine-grained, the TOG pass). Keeping it as the + interface lets all of them stay unchanged. +- **gemmini is now a Python out-of-line pass too** -- the conversion-framework + coupling (LLVMTypeConverter / getStridedElementPtr) was avoided by working at + the memref level (`memref.extract_aligned_pointer_as_index` + arith for + addresses, `llvm.inline_asm` for instructions; the existing standard lowering + finalizes to LLVM). So both decompose and gemmini live in Python; mlir-opt keeps + only the remaining custom passes. + +One constraint flows the other way: gemmini's ISA limits (max dims / size per MVIN) +set decompose's target inner-descriptor shape (the "<=4D" and max-extent bounds). +decompose must *respect* those limits when it picks what stays inner vs gets peeled +-- but respecting a constraint is not doing the lowering. + +### Cost-aware peeling (this is a cycle-accurate simulator) + +Descriptor count is a **modeled cost** (issue overhead + DRAM burst efficiency in +Ramulator). Rules: + +1. Peel the **outermost, lowest-trip-count** dims (descriptor count = product of + peeled extents). +2. Keep the inner descriptor **as large and contiguous as possible** (maximize + bytes per descriptor). + +(A pathological peel is not this pass's problem to fix: it means the operand's +layout is bad, which is a graph-level layout/copy decision, not an in-pass +relayout.) + +### Placement: hybrid (least burden) + +Keep the decision in Python (where shape/sympy info is available and iteration is +fast); keep the C++ pass purely mechanical. + +| Step | Where | +|---|---| +| peel-plan decision (which dims to peel, count estimate) | Python | +| encode plan as op attributes | Python -> MLIR | +| emit `scf.for { customized dma_start }` per the plan | C++ pass | + +The cost model can migrate into C++ later if desired. + +## Expected effects + +- **Removes recompile-dance hard-fails.** The `max_retry` `RuntimeError` path + disappears: access that does not linearize is peeled, not retried-then-killed. + This directly increases model coverage (the primary goal). +- **Removes the 4D rank cap.** Arbitrary-rank reshapes become expressible via the + base-pointer loop; the `NotImplementedError` for >4D goes away. +- **Inverts the tile<->DMA dependency.** Tile size is chosen for compute / vlane + efficiency only; the DMA conforms to whatever access results. No divisibility + forcing, no oscillation. Tile selection simplifies. +- **Shrinks the codegen.** The `FloorDiv` / `ModularIndexing` recompile branches in + `get_dma_info`, and the in-emission `RecompileSignal` paths, leave Python; the + codegen emits one declarative op instead of procedurally forcing tiles. +- **Collapses two of the three padding sites for the DMA case.** Once divisibility + is no longer required to represent access, the Python tile-adjust dance (1) is + unnecessary for DMA, and `get_mask` (2) shrinks. `TestLoopPadding` (3) is + addressed by Plan A. (Compute-side vectorization remainder is separate; see + Plan A.) +- **Behavior-preserving for the common case.** Access without floor/mod still emits + a single `dma_start` identical to today -> low-risk, incremental rollout. +- **Preserves the simulator contract.** The leaf is the existing customized + `dma_start`; Spike / gem5 / TOGSim see the same descriptor kind, just more of + them in a loop. +- **A clean tactical slice toward Plan B.** This factors out exactly the one piece + that is actually broken (DMA decomposition) into a lowering pass, without the + full linalg rewrite. +- **Cost-aware, so modeled performance is protected.** Peel small/outer, keep inner + contiguous. Pathological layouts are fixed upstream (graph copy), not by an + in-pass relayout. + +## Migration strategy + +1. Define `togsim.transfer` (op + verifier) above the existing descriptor. + Optionally formalize the descriptor as `togsim.dma_descriptor`. +2. Make the codegen emit `togsim.transfer` for loads/stores, carrying the access + maps and vlane attributes it already computes. +3. Implement `decompose-transfer` with the fast path first (<=4D affine -> one + `dma_start`), proving **bit-identical output** to today on a smoke test. +4. Add the **affine** peel path for >4D; validate end-to-end through all three + simulators (the loop-of-descriptors must satisfy the TOG / Spike / gem5 + contract). Make the pass **assert** on any non-affine residue (contract guard). +5. Land the upstream producers of the affine-only invariant: aligned axis split at + scheduling (`axis-split-scheduling.md`) and misaligned graph copy insertion. +6. Remove the `get_dma_info` recompile branches once the pass + upstream cover their + cases; use the failure ledger + assert-only `TestLoopPadding` to confirm nothing + regresses before deleting. + +## Relationship to Plan A and Plan B + +- **Plan A (graph-level padding)** reduces how often peeling/relayout is needed by + making dims granule-aligned, and retires `TestLoopPadding`. Complementary: this + op makes representation robust; Plan A reduces constraint frequency. +- **Plan B (linalg)** is the full structured-ops rewrite; `expand_shape` / + `collapse_shape` are the principled home for reshape, and the framework would + generate the same peel/relayout under the hood. This transfer op is the narrow, + now-achievable slice of that idea. + +## Risks / open questions + +- **C++ pass in the `PSAL-POSTECH/llvm-project` fork**: heavier iteration + (rebuild), logic split across two repos. Mitigated by the hybrid split (smarts in + Python, pass is mechanical). +- **TOG / Spike / gem5 contract on a loop of descriptors.** If TOG generation + assumes "one DMA = one node," the loop form needs handling. Validate at step 4. +- **Cost model accuracy** for the peel plan; start with a simple descriptor-count + threshold and refine against measured cycles. +- **Dynamic shapes**: `iter_bounds` as SSA operands; affine peel must handle + symbolic outer extents. (Symbolic-divisor floor/mod normalization is an upstream + concern, not this pass.) +- **Upstream completeness.** The pass's fail-loud contract is only safe if the + upstream producers (axis split + graph copy) actually normalize every misaligned + case. Until they do, the assert may fire on real models -- track which ops trip it + as the work-list for the upstream passes. +- **Async / tag management across the peel loop**: double-buffering / compute + overlap must survive decomposition (e.g. keep the inner large DMA async, sequence + the outer peel). + +## Appendix: alignment theory (when floor/mod is statically decomposable) + +This section records the math that decides, for a given DMA access, whether the +non-affine `FloorDiv`/`ModularIndexing` terms can be peeled into a *static* loop +of affine descriptors (free) or require a data movement (relayout / copy). + +### Setup + +A Gemmini-style descriptor addresses an element as + + addr(idx) = base + Σ_k stride_k · idx_k (integer strides, rank <= 4) + +i.e. each loop index `idx_k` contributes a **constant** stride. A DMA is +statically decomposable iff every index term it reads has constant stride over +the rectangular tile domain. Inductor index expressions, after fusion/view, carry +`FloorDiv(x, y)` and `ModularIndexing(x, y, z)` of the *flattened* loop variable +`x`. The question is when those reduce to constant-stride axes. + +### Mixed-radix decomposition + +Write the flattened index `x` (extent `E`) in mixed radix. For a `ModularIndexing` +with inner period `y` and modulus `z`, decompose uniquely as + + x = o·(y·z) + m·y + r, with 0 <= r < y, 0 <= m < z, o >= 0 + +Then `FloorDiv(x, y) = o·z + m`, and `ModularIndexing(x, y, z) = m`. Each of +`o, m, r` is a separate **implicit axis** with a constant per-axis stride — +*provided the axis boundaries do not move across the tile*. That holds iff the +period divides the extent it partitions: + +- `ModularIndexing(x, y, z)` is a valid rectangular axis **iff y·z | E**. +- `FloorDiv(x, y)` is a valid rectangular axis **iff y | E**. + +**Aligned** = the divisor (and modular period `y·z`) divides the extent, so the +wrap point lands on a fixed axis boundary -> constant stride -> peelable for free. +**Misaligned** = the wrap point falls at a loop-value-dependent position inside the +descriptor (e.g. uneven `cat`, ragged split) -> the stride is not constant -> +**not** statically decomposable; only a relayout (physical copy) fixes it. + +### One loop axis -> several implicit axes (complex fusion) + +When fusion merges many dims into one flattened loop variable, a *single* loop +axis can expand into **several** implicit axes through nested floor/mod, e.g. + + x in [0, D0·D1·D2): + a = FloorDiv(x, D1·D2) # outer + b = ModularIndexing(x, D2, D1) # middle + c = ModularIndexing(x, 1, D2) # inner + +That is three implicit descriptor axes coming from one loop axis. This is the +general case the un-flatten must handle: it is **not** limited to splitting one +axis into two. Key consequences: + +1. **The loop's own factorization is always aligned.** When the implicit axes + come from re-reading the loop's *own* contiguous factorization (the common + fusion case -- Inductor flattens contiguous dims then a consumer reads them + back via floor/mod), every period divides by construction (`D1·D2 | D0·D1·D2`, + etc.). So these un-flatten splits are **free** -- they just add descriptor + axes, never a copy. +2. **Rank blows past 4 fast.** k implicit axes per loop axis, across multiple + operands, means the descriptor rank exceeds the 4D Gemmini limit very quickly. + This is exactly why `togsim.transfer` + the peel pass matters *more* under + complex fusion, independent of any misalignment. The >4D branch in + `get_dma_info` already routes these to `togsim.transfer`. +3. **Misalignment is still only from non-factor views.** An implicit axis is + misaligned only when its period does not divide the extent -- i.e. the view + does not factor along the loop's factorization (uneven `cat`, ragged split, + group sizes that don't divide the channel count). Those, and only those, need + relayout. + +### Case-handling summary + +| Source of floor/mod | Aligned? | Handling | Cost | +|------------------------------------------------|----------|-----------------------------------|------| +| Broadcast / dim-merge (`[N,1]->[N,M]`, `i//M`) | always | un-merge (split loop axis back) | free | +| Reshape along the loop's own factorization | yes (`y·z\|E`) | un-flatten split, then peel for rank | free | +| >4D logical tile from complex fusion | yes | `togsim.transfer` -> peel into <=4D loop | free (extra DMA nodes) | +| Uneven `cat`, ragged split, non-dividing group | no | graph copy insertion (relayout, upstream) | copy = TPU `concatenate` | + +The TPU/XLA model is the reference: express only aligned views as +descriptor/bitcast (free reshape); never put a misaligned access in the +descriptor -- insert a copy (relayout) instead. Plan A (graph-level +force-contiguous / pad-to-granule, like XLA copy-insertion) is the upstream lever +that *reduces how often* the misaligned branch fires, keeping codegen affine-only. + +## Implementation status (Phase 1: codegen emission) + +Landed on branch `dma-transfer/codegen` (worktree), emission only -- the +decompose pass is deferred until explicitly signalled. A >4D access now emits a +`togsim.transfer` instead of hard-failing; without the pass it does not yet run +end-to-end (expected). + +- **`mlir_common.py` `init_tile_size`** generalized to any rank. Logical tile is + separated from the physical (<=4D) descriptor: only the innermost dims carry the + vectorized tile, all further-outer dims stay 1, and there is no rank cap. The + `nr_dim >= 3` formula reproduces the old 3D/4D values exactly (the old `[-4]=1` + is subsumed by "outer dims stay 1"); scalar/1D/2D keep their special cases. This + removes the old `raise NotImplementedError("dummy tile size fail!")` that + conflated logical and physical tile rank. +- **`mlir_codegen_backend.py`**: + - `__init__` adds `self._dma_needs_transfer = False`. + - `get_dma_info` >4D `else` branch (was + `raise NotImplementedError("Currently not implemented... ;)")`) now builds the + full N-D tile (`set_tile_size`, vlane split/stride) and sets + `self._dma_needs_transfer = True`. + - `emit_transfer(...)` emits the generic-form `"togsim.transfer"(...)` op + carrying `dma_kind`, `vlane_split_axis`, `vlane_stride`, `dram_stride`, + `tile_stride`, `padding`, with operands `(dram, dram_idx, sram, 0, tag)`. + `togsim` is an unregistered dialect, hence generic form. + - `load()` (MVIN) and `store()` (MVOUT) check the flag: if set, reset it and + call `emit_transfer`; otherwise the existing `get_dma_code` path is unchanged. + So aligned <=4D DMAs are **bit-identical** to before; only >4D accesses change. + +Validated: the 5D permute smoke test (`x.permute(4,3,2,1,0).contiguous() + 1.0`) +now emits MVIN/MVOUT `togsim.transfer` with 5D `dram_stride [1,6,30,120,360]` and a +`memref<1x1x2x4x2xf32,1>` tile, instead of crashing in `init_tile_size` or the +`get_dma_info` >4D branch. + +### Phase 2: aligned-only peel pass (landed: unit-collapse path) + +`passes/decompose_transfer.py` (registered in `passes/__init__.py`, runs before +`lower_vlane_idx`) lowers each `togsim.transfer` to a customized `memref.dma_start`: + +- **Unit-dim collapse (done, validated).** Drop extent-1 tile dims so the + descriptor reaches <=4D. The SRAM (spad) memref is collapsed to the effective + rank via `memref.collapse_shape` (the customized `dma_start` convention requires + SRAM rank == #indices == len(sram_stride)); DRAM stays flat rank-1 with its N-D + structure in `dram_stride`. The `vlane_split_axis` is **remapped** from the + original tile-dim index to the collapsed-dim index and rematerialized as a const + (carried as a value attr precisely so the pass can remap it). +- Supporting changes: `emit_transfer` now carries the SSA operands a `dma_start` + needs (`dma_type`, `vlane_stride`) + the `vlane_split_axis` value attr, so the + pass is mechanical. `lower_to_llvm.py` gains `expand-strided-metadata` to lower + `collapse_shape`. + +Validated end-to-end (Gem5 + Spike + TOGSim, `allclose=True`) on the 5D permute +`x.permute(4,3,2,1,0).contiguous() + 1.0`; no regression on 2D/3D/elementwise. + +- **Genuine >4 effective rank (isolation-only; INCOMPATIBLE with TOG -- see TODO).** + When >4 *non-unit* dims survive, the pass keeps the inner 4 as the <=4D descriptor + and peels the outer dims by **full unrolling**: one descriptor per outer-index + combo, the SRAM slice a rank-reduced `memref.subview` at the static slice offset, + the DRAM base `dram_idx + constant`. This passes `lower_text` / mlir-opt in + isolation, but **fails the full pipeline**: the C++ TOG generation pass cannot read + `memref.subview` + unrolled (constant-offset) DMAs and produces an empty + `loop_idx_list` (ValueError in `onnx_utility.py`). Surfaced once aligned axis-split + made the path reachable (pixel_shuffle -> 5D); axis-split now has a rank guard that + avoids triggering it. + +> **TODO (peel rework, tracked as GitHub issue #258).** Rewrite the >4D peel to emit +> a real `affine.for` over the peeled dims (so each DMA keeps an enclosing loop index +> the TOG pass can read) and index the spad directly instead of via `memref.subview`. +> Alternatively teach the C++ TOG pass to handle `subview` + unrolled DMAs. Until +> then the unroll path is isolation-only and the axis-split rank guard keeps it +> unreached. + +The input stays per-axis affine by upstream guarantee, so both paths are pure +mechanical peeling. A non-affine residue is a contract violation (aligned floor/mod +removal lives in `axis-split-scheduling.md`, misaligned relayout in graph copy +insertion -- see "Division of labor"); a genuinely non-affine / indirect index +would surface as a build failure here rather than being silently relaid out. diff --git a/docs/linalg-codegen-migration.md b/docs/linalg-codegen-migration.md new file mode 100644 index 00000000..9a60ba31 --- /dev/null +++ b/docs/linalg-codegen-migration.md @@ -0,0 +1,224 @@ +# Linalg-based codegen migration (Plan B) + +Status: **deferred / design only**. This is not scheduled work. It records *why* +a linalg-based rewrite of the MLIR codegen will eventually be worth doing, and a +rough plan, so the decision does not have to be re-derived from scratch later. + +For the near-term padding/dynamic-shape work, see Plan A (graph-level padding) in +the "Relationship to Plan A" section below. Plan A is the one to do first; it does +not depend on this document. + +## TL;DR + +The current MLIR codegen (`PyTorchSimFrontend/mlir/`) does not just emit loops — +it hand-implements the entire hardware mapping (tiling, vectorization, DMA, +scratchpad allocation, vector-lane distribution) as Python string emission. That +works, but it entangles three concerns that should be separable: + +1. **what to compute** (the op math), +2. **how to map it onto the NPU** (tile sizes, vlane layout, DMA/scratchpad), +3. **how to make shapes fit the hardware** (padding / divisibility). + +Concern 3 is currently spread across three places and is the source of the +"padding is heuristic and fragile" pain. Plan B factors concern 1 up to the +`linalg` dialect and rebuilds concern 2 as a set of MLIR lowering passes, so that +concern 3 falls out of the structured representation instead of being patched. + +Plan B is a multi-month, higher-risk effort because the hardware mapping (concern +2) is bespoke and has no upstream equivalent. Do it when the payoff (separation, +reuse of upstream tiling/vectorization/fusion, easier autotuning, lower codegen +maintenance) is worth that cost — not as a means to fix padding alone. + +## Where we are today + +Entry point: PyTorchSim is an **Inductor backend**. Inductor handles capture, +decomposition, lowering to Inductor IR, scheduling, and **fusion**. Our code runs +at the codegen (step-5) stage and turns scheduled `SchedulerNode`s into MLIR. + +`MLIRKernel` (`mlir/mlir_codegen_backend.py`) and friends emit, by hand: + +- explicit DMA: `memref.dma_start` with MVIN/MVOUT encoding, `vlane_split_axis`, + `vlane_stride` (`load`, `store`, `get_dma_code`); +- explicit scratchpad: `.spad` sections and `memref.global @bufN_spad` + (`allocate_sram_buffer`, `get_scratchpad_buffer`); +- explicit vector-lane (vlane) distribution: `vmap.vlane_split_axis` / + `vlane_stride`, `vlane_offset`, `get_used_vlane` (`_index_expr`, + `get_dma_info`); +- explicit tile descriptors: `MLIRMultiDimTile` with per-axis tile sizes; +- explicit reduction loops: manual accumulator/iterator vars + `affine_yield` + (`reduction`, `codegen_loops`); +- explicit vector-tail masking: `get_mask`. + +This is roughly 5,500 lines across `mlir_codegen_backend.py`, `mlir_common.py`, +`mlir_template.py`, and `mlir_ops.py`, plus per-op templates +(`mlir_gemm_template.py`, `mlir_conv_*`, `mlir_bmm_template.py`, +`mlir_sdpa_template.py`, `mlir_sort_template.py`, `mlir_cat_template.py`, +`mlir_maxpool_template.py`). Most of it encodes how this specific NPU's memory +hierarchy and vector unit work. **That accumulated hardware knowledge is the asset +and the cost center for any rewrite.** + +### Padding lives in three places + +This is the key observation motivating the separation. Divisibility / padding is +handled by: + +1. **Python recompile + tile adjustment** — `get_dma_info` FloorDiv / + ModularIndexing handling, `_index_expr`, `convert_indirect_indexing`: "if the + tile size does not divide the dim, bump the tile and raise `RecompileSignal` to + recompile." This is the actual heuristic-padding body, and it is in Python, not + MLIR. +2. **Python `get_mask`** — vector-tail masking for the innermost compute loop. +3. **MLIR `TestLoopPadding`** pass (in the `PSAL-POSTECH/llvm-project` fork) — + rounds affine loop bounds up to a multiple of the step and resizes buffers, + reverse-engineering the loop<->tensor mapping from affine maps. + +Removing only (3) does not fix the fragility; (1) is arguably the larger source of +"the loop and the tensor do not line up." Any real fix has to address all three. + +## Target architecture (Plan B) + +Split the codegen into two layers with a clean boundary: + +``` +L0 ATen / FX logical graph (dynamic dims as SymInt) <- Inductor, unchanged +L1 math layer: Inductor IR -> linalg.generic / named ops (untiled, unvectorized) + = iteration domain + affine indexing maps + scalar body. + Decides nothing about tiles / lanes / DMA. +L2 mapping layer: tiling(+pad) -> vectorize(+vl) -> bufferize + -> DMA / scratchpad / vlane lowering -> leaf replacement + = hardware mapping, parameterized by a target description. +L3 LLVM / RVV / systolic microkernel +``` + +Why this helps: + +- **The loop<->tensor mapping stops being reverse-engineered.** `linalg` ops carry + `indexing_maps` + `iterator_types`, i.e. exactly the information `TestLoopPadding` + tries to recover. Padding becomes a *parameter of the tiling transform* + (`tensor.pad` generated with full knowledge of the maps), not a separate + analysis pass. +- **The padding strategies we want become per-axis policy** in L2: systolic + operand axes -> pad-to-uniform-tile (keeps a single 128x128 microkernel and a + single gem5 latency entry); VPU / vector axes -> RVV `vl` (no padding); reduction + axes -> `affine.min` clamp. One mechanism, selected per axis, instead of three + scattered implementations. +- **Fusion policy stays in Inductor** (its scheduler decides what is one kernel), + while the *mechanism* is upstream `linalg` tile-and-fuse. Pointwise epilogue + fusion is essentially free because Inductor already composes `inner_fn`s into a + single fused body -> one `linalg.generic`. +- **Dynamic shapes** are carried as `?` dims + symbolic affine, uniformly handled + by L2 rather than by the recompile dance. +- **The cost model gets cleaner, not harder**: tile shape becomes an explicit + attribute, so the gem5 latency table / TOG key on it directly instead of on + inferred loop shapes. + +### What is reusable vs bespoke + +- **Reusable from upstream MLIR**: `linalg` ops, tiling + `tensor.pad`, + vectorization, bufferization, the TilingInterface. The L1 translation + (Inductor IR -> `linalg.generic`) is a *generic* translator (one path covers all + regular pointwise/reduction), not per-op work — Inductor IR is already in + structured iteration-domain + scalar-body form. +- **Bespoke, must be (re)written as MLIR passes**: MVIN/MVOUT DMA encoding, + `.spad` scratchpad assignment, and the vlane_split mapping. **These have no + upstream equivalent.** This is the bulk of the effort and the main risk: the + knowledge currently in ~5,500 lines of Python emission must be re-expressed as + custom bufferization-to-DMA / scratchpad / vlane-vectorization lowerings. + +## Expressibility boundary + +A regular (linalg.generic) op needs: a fixed rectangular iteration space; every +operand index an **affine** function of loop vars (no data-dependent indexing); +each axis purely `parallel` or a simple `reduction` (no scan/recurrence); a +statically-determined output shape (dynamic `?` ok, data-dependent shape not); and +a data-independent body (`arith.select` ok, data-dependent branching not). + +Expressible: elementwise, broadcast, transpose, reductions (incl. multiple +reduction axes), matmul / bmm / contractions, direct conv, pooling, and fused +chains of these (matmul+bias+activation, prologue cast/dequant/transpose, +pointwise->reduce). `slice` / `pad` / `cat` are structured `tensor` ops (not +`linalg.generic`) but are supported by the same pipeline. + +Not expressible -> stay as hand-written custom kernels: data-dependent indexing +(gather/scatter/embedding), sort/topk, data-dependent output shape +(nonzero/unique/masked_select), scan/recurrence (cumsum), and online/streaming +algorithms (flash-attention). In our op set this means **`sdpa` (online softmax) +and `sort` remain custom**; gemm/conv/bmm/maxpool are regular; `cat` is a +structured tensor op. + +## Migration strategy (when Plan B is scheduled) + +Incremental, op-by-op, with a numeric and a structural safety net. Do **not** +big-bang. + +1. **Stand up the L2 pipeline for one op (matmul first).** Emit `linalg.matmul` + from the matmul path; wire tiling(+pad, pad_value=0) -> bufferize -> a custom + pass that lowers the 128^3 leaf tile to the existing systolic intrinsic -> + LLVM. Milestone 1 is end-to-end correctness through all three simulators + (Spike functional, gem5 latency, TOGSim cycle) for a single matmul. +2. **Demote `TestLoopPadding` to assert-only** (check, do not modify; fail/log if a + loop bound is not a multiple of its step). Run the full test suite; anything it + flags is a case L1/L2 has not covered yet. +3. **Migrate the remaining regular ops** (conv, bmm, pointwise, reductions, + maxpool). Pointwise/reduction go through the generic L1 translator; VPU + remainder via `vl`. +4. **Delete `TestLoopPadding`** once the assert-only version is silent across the + suite, and retire the Python recompile/tile-adjust dance and most of `get_mask`. +5. Leave `sdpa` and `sort` as custom kernels that bypass L1/L2. + +### Risks + +- **Simulator-facing contract.** The current emission is tuned to produce exactly + the LLVM / TOG shape the three simulators expect. `linalg`'s standard lowering + emits different IR; re-validating the lowered artifact end-to-end (especially TOG + generation, which may assume specific loop/memory patterns) is the real + integration risk. This is why milestone 1 is "one matmul, end-to-end," not "all + ops, emission only." +- **Re-encoding the hardware mapping.** DMA/scratchpad/vlane lowerings are new code + with no upstream reference; budget for them dominating the schedule. +- **Inductor index expressions.** Inductor often collapses dims into one flat index + with `FloorDiv` / `ModularIndexing`, which are not affine; `linalg` indexing_maps + must be affine. Either keep dims uncollapsed or normalize div/mod back to + multi-dim affine. (We already convert these to affine strings today in + `_convert_sympy_to_mlir_expr`, but that path will need to be revisited for the + map-carrying representation.) +- **Fusion seams.** Not everything fuses cleanly (reductions with mismatched axes, + transpose/layout mismatches); expect some barriers, same as any framework. + +## Relationship to Plan A (graph-level padding) — do this first + +Plan A inserts padding at the FX/graph level (via Inductor's +`post_grad_custom_pass`) so that tiled dims arrive at codegen already aligned +(`tile_granule * symbol`). Under the hard constraint that we **keep the Inductor +spine**, Plan A is the high-ROI move: + +- It collapses all three padding sites at once: the recompile/tile-adjust dance (1) + becomes unnecessary (tiles always divide), `get_mask` (2) becomes trivial (no + tail), and `TestLoopPadding` (3) becomes unnecessary. +- It does **not** touch the bespoke DMA/scratchpad/vlane mapper. + +Key correctness facts that make Plan A tractable: + +- Matmul contraction (K) padding with zeros is *exact* (additive identity); weights + are constants, so they can be zero-padded once, offline, at no runtime cost. +- Padding only ever corrupts results when a *non-contraction* padded axis is later + reduced (softmax over keys; layernorm if hidden is padded). Those points need + masking; everything else is pad-transparent. +- Safety rule: default any op to slice-back-to-real-shape; only opt an op into + "propagate padded shape" once it is proven pad-transparent or given a mask + handler. Correct-by-construction; unknown ops cannot silently corrupt. + +Plan A and Plan B are compatible: Plan A's graph-level alignment makes the eventual +Plan B simpler (L2 tiling rarely needs to pad, because dims already divide). + +## Open questions + +- Does the current toolchain (the `PSAL-POSTECH/llvm-project` fork) already ship the + `linalg` + transform/tiling passes, or were they stripped? (Almost certainly + present if it tracks upstream — verify before committing.) +- Can the systolic leaf be expressed cleanly as a match-and-replace on a fixed-size + `linalg.matmul`, or does weight-stationary loading order force a more custom + representation? +- How much of `mlir_ops.py` (the scalar `OpsHandler`) survives? It currently emits + *vectorized* ops (compute_vec_size, broadcast) and is therefore entangled with + vlane; the linalg body should be scalar, with vectorization done in L2. diff --git a/docs/mlir-python-bindings.md b/docs/mlir-python-bindings.md new file mode 100644 index 00000000..6bb03339 --- /dev/null +++ b/docs/mlir-python-bindings.md @@ -0,0 +1,102 @@ +# Enabling MLIR Python bindings + +Goal: ship the MLIR Python bindings (`import mlir`, `mlir.ir`, `mlir.dialects`) +so we can write MLIR passes in Python (imperative IR rewriting via the bindings) +instead of only C++ passes in the `PSAL-POSTECH/llvm-project` fork. See +`dma-transfer-lowering.md` for the first intended use (a Python decompose pass). + +## How LLVM reaches the runtime (why this touches 3 places) + +``` +PSAL-POSTECH/llvm-project (fork, tag vX.Y.Z) + .github/workflows/build-torchsim.yaml -- CI builds + releases riscv-llvm-release.tar.gz + | (release asset) + v +thirdparty/github-releases.json -- pins llvm_project.release_tag + asset + | + v +Dockerfile.base -- downloads asset, extracts to /riscv-llvm, + sets TORCHSIM_LLVM_PATH (+ now PYTHONPATH) +``` + +`scripts/build_from_source.sh` is the alternative source-build path (not the +normal flow, but kept consistent). + +## The one real blocker: Python ABI must match + +The bindings are a native CPython extension (`_mlir.cpython-3XX-*.so`). They only +import under the **same Python minor version** they were built against. The +runtime base image uses **conda Python 3.11**. So the artifact must be built with +**Python 3.11**. Building with the build container's default (ubuntu-22.04 -> +3.10) produces bindings that fail to import at runtime with a confusing error +much later -- hence the fail-fast guard in the CI step. + +Patch version (3.11.x) does not matter; minor version (3.11 vs 3.10) does. + +## What was changed + +- **`scripts/build_from_source.sh`**: cmake gets + `-DMLIR_ENABLE_BINDINGS_PYTHON=ON -DPython3_EXECUTABLE=$(command -v python3)`; + build deps (nanobind/pybind11/numpy/PyYAML) pip-installed; after `make install` + the build-tree `tools/mlir/python_packages` is copied into `/riscv-llvm` + (install does not place it there). PYTHONPATH exported for the current shell. +- **`Dockerfile.base`**: `ENV PYTHONPATH=/riscv-llvm/python_packages/mlir_core:$PYTHONPATH` + after the LLVM artifact is extracted. +- **`llvm-project/.github/workflows/build-torchsim.yaml`** (fork): same cmake + flags + deps; copies `python_packages` into the `riscv-llvm` tree so the + existing `tar` includes it; fail-fast guard requiring `python3.11`. + +## Rollout sequence (must be done in order) + +1. **python3.11 in the build container: done, non-root.** The CI step keeps the + original `-u $(id -u):$(id -g)` (no root assumed) and fetches a standalone + CPython 3.11 with `uv` (`uv venv --python 3.11`), then points + `Python3_EXECUTABLE` at that venv. No apt / no root needed. ubuntu-22.04's + default 3.10 is not used for the bindings. + - ABI note: extensions built against a uv/python-build-standalone CPython 3.11 + are expected to import under the runtime conda CPython 3.11 (same minor + version, standard builds are C-ABI compatible). The verify step below is the + check; if it ever fails, build instead in the runtime image (`python:3.11` or + the pytorch base) so build Python == runtime Python by construction. +2. **Push the fork changes** to `PSAL-POSTECH/llvm-project` and cut a new tag + (e.g. `v1.0.9`). CI builds `riscv-llvm-release.tar.gz` now containing + `python_packages/`. +3. **Bump `thirdparty/github-releases.json`** -> `llvm_project.release_tag` to the + new tag (and `asset_name` unchanged). This triggers a new base image build. +4. **Rebuild the base image** (the fork CI already dispatches `build_base`; or run + the PyTorchSim docker-image workflow) so `Dockerfile.base` produces an image + with the bindings + PYTHONPATH. + +## Verify + +Inside the rebuilt container (or after `build_from_source.sh`): + +```bash +python -c "import mlir; print(mlir.__file__)" # -> /riscv-llvm/python_packages/mlir_core/mlir/__init__.py +python -c "from mlir.ir import Context; c=Context(); c.allow_unregistered_dialects=True; print('ok')" +python -c "from mlir.dialects import scf, affine, arith; print('dialects ok')" +``` + +`allow_unregistered_dialects=True` is what lets us read/write the custom ops +(`togsim.transfer`, the customized `memref.dma_start`) generically without +registering a dialect in the bindings. + +## Notes / gotchas + +- Keep the bindings statically linked (default, i.e. do NOT add + `-DBUILD_SHARED_LIBS=ON` / `-DLLVM_BUILD_LLVM_DYLIB=ON`); otherwise the `.so` + needs libMLIR/libLLVM at runtime and the artifact + LD_LIBRARY_PATH grow. +- Worktrees: add the same `PYTHONPATH` line to the worktree `.envrc` (see + `docs/worktrees.md`) if a worktree overrides paths. +- The bindings are an additive, optional dependency: text emission + C++ passes + keep working unchanged. Only new Python passes require the bindings present. +- This LLVM fork's MLIR bindings use **pybind11** (not nanobind) and require + **pybind11 <= 2.10.3**: newer pybind11 (3.x) fails to compile `IRCore.cpp` with + `def_property family does not currently support keep_alive`. Pin it + (`pybind11>=2.9.0,<=2.10.3`). See `mlir/python/requirements.txt` for the fork's + pins. pybind11 is build-time only; the runtime needs just the built `.so` + numpy. +- numpy: the fork's requirements pin `<=1.26`, but a local build against numpy 2.x + compiled and imported fine, so we keep numpy at the runtime version (2.x) to + avoid a numpy-1-built / numpy-2-runtime ABI mismatch. (Validated locally: + conda 3.11 + pybind11 2.10.3 + numpy 2.x -> `import mlir` and parsing a custom + `togsim.transfer` op with floordiv/mod affine maps both work.) diff --git a/scripts/build_from_source.sh b/scripts/build_from_source.sh index 4e7ff604..f23eab82 100644 --- a/scripts/build_from_source.sh +++ b/scripts/build_from_source.sh @@ -45,12 +45,23 @@ export GEM5_PATH="$home/gem5/build/RISCV/gem5.opt" cd "$home" # LLVM + MLIR (RISCV target) +# MLIR Python bindings are enabled so Python-side MLIR passes can run. The +# bindings are a native extension: they MUST be built against the same Python +# that runs PyTorchSim at runtime (the conda 3.11 here) or `import mlir` will +# fail with an ABI mismatch. nanobind/pybind11/numpy/PyYAML are build-time deps. +python3 -m pip install --user "pybind11>=2.9.0,<=2.10.3" numpy PyYAML git clone --depth 1 --branch "$LLVM_TAG" "https://github.com/${LLVM_REPO}.git" cd llvm-project && mkdir -p build && cd build && \ cmake -DLLVM_ENABLE_PROJECTS=mlir -DCMAKE_BUILD_TYPE=Release \ -DCMAKE_INSTALL_PREFIX=/riscv-llvm -DLLVM_TARGETS_TO_BUILD=RISCV \ + -DMLIR_ENABLE_BINDINGS_PYTHON=ON \ + -DPython3_EXECUTABLE="$(command -v python3)" \ -G "Unix Makefiles" ../llvm && \ - make -j && make install + make -j && make install && \ + rm -rf /riscv-llvm/python_packages && \ + cp -r tools/mlir/python_packages /riscv-llvm/python_packages +# Make the bindings importable in this shell (also set in .envrc / Dockerfile.base) +export PYTHONPATH="/riscv-llvm/python_packages/mlir_core:$PYTHONPATH" cd "$home" # Spike Simulator diff --git a/scripts/op_coverage.py b/scripts/op_coverage.py new file mode 100644 index 00000000..1f4567b6 --- /dev/null +++ b/scripts/op_coverage.py @@ -0,0 +1,540 @@ +"""Op-coverage diagnostic for new LLM models on PyTorchSim. + +Runs each model in two phases: + Phase 1 (enumerate): custom torch.compile backend captures the FX graph and + lists every aten op that appears, without touching NPU. + Phase 2 (run): torch.compile(model) on npu:0, real forward. On crash, + parses the traceback to identify the failing op. + +Usage: + python scripts/op_coverage.py # all models + python scripts/op_coverage.py --models qwen2 # subset + python scripts/op_coverage.py --enumerate-only # skip NPU compile (fast) +""" + +import argparse +import datetime as _dt +import os +import re +import sys +import traceback +from contextlib import contextmanager + +import torch + +REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +if REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + + +# --------------------------------------------------------------------------- +# Model registry: each entry returns (model, kwargs_for_forward) on CPU. +# Sizes follow "small but realistic" variants (1-layer) so a forward is cheap +# enough to actually drive through TOGSim. +# --------------------------------------------------------------------------- + +def _causal_mask(batch, seq_len, dtype): + min_v = torch.finfo(dtype).min + m = torch.full((seq_len, seq_len), min_v, dtype=dtype) + if seq_len > 1: + m = torch.triu(m, diagonal=1) + return m[None, None, :, :].expand(batch, 1, -1, -1).contiguous() + + +def build_qwen2(batch=1, seq_len=32, dtype=torch.float32): + from transformers.models.qwen2.configuration_qwen2 import Qwen2Config + from transformers.models.qwen2.modeling_qwen2 import Qwen2Model + cfg = Qwen2Config( + vocab_size=4096, + hidden_size=1536, + num_attention_heads=12, + num_key_value_heads=2, + intermediate_size=8960, + num_hidden_layers=2, + max_position_embeddings=4096, + rms_norm_eps=1e-6, + rope_theta=1000000.0, + torch_dtype=dtype, + use_cache=False, + _attn_implementation="eager", + ) + model = Qwen2Model(cfg).eval().to(dtype=dtype) + input_ids = torch.randint(0, cfg.vocab_size, (batch, seq_len)) + attn_mask = _causal_mask(batch, seq_len, dtype) + return model, {"input_ids": input_ids, "attention_mask": attn_mask} + + +def build_gemma(batch=1, seq_len=32, dtype=torch.float32): + from transformers.models.gemma.configuration_gemma import GemmaConfig + from transformers.models.gemma.modeling_gemma import GemmaModel + cfg = GemmaConfig( + vocab_size=4096, + hidden_size=2048, + num_attention_heads=8, + num_key_value_heads=1, + intermediate_size=16384, + num_hidden_layers=2, + head_dim=256, + max_position_embeddings=4096, + rms_norm_eps=1e-6, + rope_theta=10000.0, + torch_dtype=dtype, + use_cache=False, + _attn_implementation="eager", + ) + model = GemmaModel(cfg).eval().to(dtype=dtype) + input_ids = torch.randint(0, cfg.vocab_size, (batch, seq_len)) + attn_mask = _causal_mask(batch, seq_len, dtype) + return model, {"input_ids": input_ids, "attention_mask": attn_mask} + + +def build_gemma2(batch=1, seq_len=32, dtype=torch.float32): + from transformers.models.gemma2.configuration_gemma2 import Gemma2Config + from transformers.models.gemma2.modeling_gemma2 import Gemma2Model + cfg = Gemma2Config( + vocab_size=4096, + hidden_size=2304, + num_attention_heads=8, + num_key_value_heads=4, + intermediate_size=9216, + num_hidden_layers=2, + head_dim=256, + max_position_embeddings=4096, + rms_norm_eps=1e-6, + rope_theta=10000.0, + torch_dtype=dtype, + use_cache=False, + attn_logit_softcapping=50.0, + final_logit_softcapping=30.0, + sliding_window=16, + _attn_implementation="eager", + ) + model = Gemma2Model(cfg).eval().to(dtype=dtype) + input_ids = torch.randint(0, cfg.vocab_size, (batch, seq_len)) + attn_mask = _causal_mask(batch, seq_len, dtype) + return model, {"input_ids": input_ids, "attention_mask": attn_mask} + + +def build_phi3(batch=1, seq_len=32, dtype=torch.float32): + from transformers.models.phi3.configuration_phi3 import Phi3Config + from transformers.models.phi3.modeling_phi3 import Phi3Model + cfg = Phi3Config( + vocab_size=4096, + pad_token_id=0, + bos_token_id=1, + eos_token_id=2, + hidden_size=3072, + num_attention_heads=32, + num_key_value_heads=32, + intermediate_size=8192, + num_hidden_layers=2, + max_position_embeddings=4096, + rms_norm_eps=1e-5, + rope_theta=10000.0, + torch_dtype=dtype, + use_cache=False, + _attn_implementation="eager", + ) + model = Phi3Model(cfg).eval().to(dtype=dtype) + input_ids = torch.randint(0, cfg.vocab_size, (batch, seq_len)) + attn_mask = _causal_mask(batch, seq_len, dtype) + return model, {"input_ids": input_ids, "attention_mask": attn_mask} + + +def _build_lm(cfg, ModelCls, batch, seq_len, dtype): + """Shared helper: build a causal-LM-style model and matching token+mask inputs.""" + model = ModelCls(cfg).eval().to(dtype=dtype) + input_ids = torch.randint(0, cfg.vocab_size, (batch, seq_len)) + attn_mask = _causal_mask(batch, seq_len, dtype) + return model, {"input_ids": input_ids, "attention_mask": attn_mask} + + +def build_qwen3(batch=1, seq_len=32, dtype=torch.float32): + from transformers.models.qwen3.configuration_qwen3 import Qwen3Config + from transformers.models.qwen3.modeling_qwen3 import Qwen3Model + cfg = Qwen3Config( + vocab_size=4096, hidden_size=1024, num_attention_heads=8, num_key_value_heads=4, + intermediate_size=3072, num_hidden_layers=2, max_position_embeddings=4096, + rms_norm_eps=1e-6, rope_theta=1000000.0, torch_dtype=dtype, use_cache=False, + _attn_implementation="eager", + ) + return _build_lm(cfg, Qwen3Model, batch, seq_len, dtype) + + +def build_qwen3_moe(batch=1, seq_len=32, dtype=torch.float32): + from transformers.models.qwen3_moe.configuration_qwen3_moe import Qwen3MoeConfig + from transformers.models.qwen3_moe.modeling_qwen3_moe import Qwen3MoeModel + cfg = Qwen3MoeConfig( + vocab_size=4096, hidden_size=1024, num_attention_heads=8, num_key_value_heads=4, + intermediate_size=3072, moe_intermediate_size=768, num_experts=4, num_experts_per_tok=2, + decoder_sparse_step=1, num_hidden_layers=2, max_position_embeddings=4096, + rms_norm_eps=1e-6, rope_theta=1000000.0, torch_dtype=dtype, use_cache=False, + _attn_implementation="eager", + ) + return _build_lm(cfg, Qwen3MoeModel, batch, seq_len, dtype) + + +def build_gemma3(batch=1, seq_len=32, dtype=torch.float32): + from transformers.models.gemma3.configuration_gemma3 import Gemma3TextConfig + from transformers.models.gemma3.modeling_gemma3 import Gemma3TextModel + cfg = Gemma3TextConfig( + vocab_size=4096, hidden_size=2048, num_attention_heads=8, num_key_value_heads=4, + intermediate_size=8192, head_dim=256, num_hidden_layers=2, + sliding_window=16, sliding_window_pattern=2, + max_position_embeddings=4096, rms_norm_eps=1e-6, rope_theta=10000.0, + torch_dtype=dtype, use_cache=False, _attn_implementation="eager", + ) + return _build_lm(cfg, Gemma3TextModel, batch, seq_len, dtype) + + +def build_deepseek_v3(batch=1, seq_len=32, dtype=torch.float32): + from transformers.models.deepseek_v3.configuration_deepseek_v3 import DeepseekV3Config + from transformers.models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3Model + cfg = DeepseekV3Config( + vocab_size=4096, hidden_size=1024, num_attention_heads=16, num_key_value_heads=16, + intermediate_size=4096, moe_intermediate_size=512, + n_routed_experts=8, num_experts_per_tok=2, n_shared_experts=1, + n_group=2, topk_group=1, + q_lora_rank=512, kv_lora_rank=128, qk_rope_head_dim=32, qk_nope_head_dim=32, v_head_dim=64, + num_hidden_layers=2, first_k_dense_replace=1, + max_position_embeddings=4096, rms_norm_eps=1e-6, rope_theta=10000.0, + torch_dtype=dtype, use_cache=False, _attn_implementation="eager", + ) + return _build_lm(cfg, DeepseekV3Model, batch, seq_len, dtype) + + +def build_llama4(batch=1, seq_len=32, dtype=torch.float32): + from transformers.models.llama4.configuration_llama4 import Llama4TextConfig + from transformers.models.llama4.modeling_llama4 import Llama4TextModel + cfg = Llama4TextConfig( + vocab_size=4096, hidden_size=1024, num_attention_heads=8, num_key_value_heads=4, + intermediate_size=3072, intermediate_size_mlp=3072, + num_local_experts=4, num_experts_per_tok=1, num_hidden_layers=2, interleave_moe_layer_step=2, + max_position_embeddings=4096, rms_norm_eps=1e-6, rope_theta=10000.0, + torch_dtype=dtype, use_cache=False, _attn_implementation="eager", + ) + return _build_lm(cfg, Llama4TextModel, batch, seq_len, dtype) + + +def build_glm4(batch=1, seq_len=32, dtype=torch.float32): + from transformers.models.glm4.configuration_glm4 import Glm4Config + from transformers.models.glm4.modeling_glm4 import Glm4Model + cfg = Glm4Config( + vocab_size=4096, pad_token_id=0, bos_token_id=1, eos_token_id=2, + hidden_size=1536, num_attention_heads=12, num_key_value_heads=2, + intermediate_size=4096, num_hidden_layers=2, max_position_embeddings=4096, + rms_norm_eps=1e-5, rope_theta=10000.0, torch_dtype=dtype, use_cache=False, + _attn_implementation="eager", + ) + return _build_lm(cfg, Glm4Model, batch, seq_len, dtype) + + +def build_olmo2(batch=1, seq_len=32, dtype=torch.float32): + from transformers.models.olmo2.configuration_olmo2 import Olmo2Config + from transformers.models.olmo2.modeling_olmo2 import Olmo2Model + cfg = Olmo2Config( + vocab_size=4096, hidden_size=2048, num_attention_heads=16, num_key_value_heads=16, + intermediate_size=8192, num_hidden_layers=2, max_position_embeddings=4096, + rms_norm_eps=1e-6, rope_theta=10000.0, torch_dtype=dtype, use_cache=False, + _attn_implementation="eager", + ) + return _build_lm(cfg, Olmo2Model, batch, seq_len, dtype) + + +def build_granite(batch=1, seq_len=32, dtype=torch.float32): + from transformers.models.granite.configuration_granite import GraniteConfig + from transformers.models.granite.modeling_granite import GraniteModel + cfg = GraniteConfig( + vocab_size=4096, hidden_size=2048, num_attention_heads=16, num_key_value_heads=8, + intermediate_size=8192, num_hidden_layers=2, max_position_embeddings=4096, + rms_norm_eps=1e-5, rope_theta=10000.0, torch_dtype=dtype, use_cache=False, + _attn_implementation="eager", + ) + return _build_lm(cfg, GraniteModel, batch, seq_len, dtype) + + +def build_phimoe(batch=1, seq_len=32, dtype=torch.float32): + from transformers.models.phimoe.configuration_phimoe import PhimoeConfig + from transformers.models.phimoe.modeling_phimoe import PhimoeModel + cfg = PhimoeConfig( + vocab_size=4096, hidden_size=1024, num_attention_heads=8, num_key_value_heads=4, + intermediate_size=3072, num_local_experts=4, num_experts_per_tok=2, + num_hidden_layers=2, max_position_embeddings=4096, + rms_norm_eps=1e-5, rope_theta=10000.0, torch_dtype=dtype, use_cache=False, + _attn_implementation="eager", + ) + return _build_lm(cfg, PhimoeModel, batch, seq_len, dtype) + + +def build_mamba2(batch=1, seq_len=32, dtype=torch.float32): + # State-space model: no attention, no RoPE -- completely different op profile. + # Invariant: num_heads * head_dim == intermediate_size == expand * hidden_size + # (modeling_mamba2.py:171 + the view(B, num_heads*head_dim) at line 365). + from transformers.models.mamba2.configuration_mamba2 import Mamba2Config + from transformers.models.mamba2.modeling_mamba2 import Mamba2Model + cfg = Mamba2Config( + vocab_size=4096, hidden_size=512, + num_heads=16, head_dim=64, + state_size=16, chunk_size=16, + expand=2, n_groups=1, + num_hidden_layers=2, torch_dtype=dtype, use_cache=False, + ) + model = Mamba2Model(cfg).eval().to(dtype=dtype) + input_ids = torch.randint(0, cfg.vocab_size, (batch, seq_len)) + # Mamba has no attention mask; pass none. + return model, {"input_ids": input_ids} + + +def build_mllama(batch=1, seq_len=32, dtype=torch.float32): + # Llama 3.2 Vision -- text branch only (text-only call path). + # MllamaRotaryEmbedding requires config.rope_scaling["rope_type"]; pass default. + from transformers.models.mllama.configuration_mllama import MllamaTextConfig + from transformers.models.mllama.modeling_mllama import MllamaTextModel + cfg = MllamaTextConfig( + vocab_size=4096, pad_token_id=0, bos_token_id=1, eos_token_id=2, + hidden_size=1024, num_attention_heads=8, num_key_value_heads=4, + intermediate_size=3072, num_hidden_layers=2, + cross_attention_layers=[], + max_position_embeddings=4096, rms_norm_eps=1e-5, rope_theta=10000.0, + rope_scaling={"rope_type": "default"}, + torch_dtype=dtype, use_cache=False, _attn_implementation="eager", + ) + return _build_lm(cfg, MllamaTextModel, batch, seq_len, dtype) + + +BUILDERS = { + "qwen2": build_qwen2, + "gemma": build_gemma, + "gemma2": build_gemma2, + "phi3": build_phi3, + # Models newly available with transformers 4.51.3 + "qwen3": build_qwen3, + "qwen3_moe": build_qwen3_moe, + "gemma3": build_gemma3, + "deepseek_v3": build_deepseek_v3, + "llama4": build_llama4, + "glm4": build_glm4, + "olmo2": build_olmo2, + "granite": build_granite, + "phimoe": build_phimoe, + "mamba2": build_mamba2, + "mllama": build_mllama, +} + + +# --------------------------------------------------------------------------- +# Phase 1: enumerate aten ops by intercepting the FX graph from torch.compile. +# --------------------------------------------------------------------------- + +def _node_op_name(target): + # OpOverload / OpOverloadPacket: has a .name() method returning "aten::mm.default" etc. + if hasattr(target, "name") and callable(target.name): + try: + return target.name() + except Exception: + pass + if hasattr(target, "_schema"): + try: + return str(target._schema.name) + ( + "." + target._schema.overload_name if target._schema.overload_name else "" + ) + except Exception: + pass + # torch.* python builtins: use their __module__/__qualname__ + mod = getattr(target, "__module__", "") + qn = getattr(target, "__qualname__", None) or getattr(target, "__name__", "") + if mod and qn: + return f"{mod}.{qn}" + return str(target) + + +@torch.no_grad() +def enumerate_ops(model, inputs): + """Capture the post-AOTAutograd aten graph(s) via aot_module_simplified. + + This is the same level of IR TOGSim/Inductor consumes, so the op set + matches what the NPU backend actually has to lower. + """ + from functorch.compile import aot_module_simplified + + seen = set() + graph_sizes = [] + + def fw_compiler(gm, example_inputs): + graph_sizes.append(sum(1 for _ in gm.graph.nodes)) + for node in gm.graph.nodes: + if node.op == "call_function": + seen.add(_node_op_name(node.target)) + return gm.forward + + def dynamo_backend(gm, example_inputs): + return aot_module_simplified(gm, example_inputs, fw_compiler=fw_compiler) + + torch._dynamo.reset() + compiled = torch.compile(model, backend=dynamo_backend, dynamic=False) + compiled(**inputs) + return sorted(seen), graph_sizes + + +# --------------------------------------------------------------------------- +# Phase 2: real NPU compile + run. Capture and parse failure tracebacks. +# --------------------------------------------------------------------------- + +ATEN_RE = re.compile(r"aten[.:][a-zA-Z_][a-zA-Z0-9_.]*") +NOTIMPL_RE = re.compile(r"NotImplementedError[: ]+(.*)") + + +def parse_failure(tb_text): + aten_hits = [] + for m in ATEN_RE.finditer(tb_text): + op = m.group(0).replace("aten:", "aten.").lstrip(".") + if op not in aten_hits: + aten_hits.append(op) + msg = "" + nm = NOTIMPL_RE.search(tb_text) + if nm: + msg = nm.group(1).strip().splitlines()[0] + return aten_hits, msg + + +@torch.no_grad() +def run_on_npu(model, inputs): + device = torch.device("npu:0") + model = model.to(device) + inputs = {k: v.to(device) for k, v in inputs.items()} + torch._dynamo.reset() + compiled = torch.compile(model, dynamic=False) + out = compiled(**inputs) + # touch the output to force completion + if hasattr(out, "last_hidden_state"): + out.last_hidden_state.cpu() + elif isinstance(out, torch.Tensor): + out.cpu() + return "OK", None, None + + +# --------------------------------------------------------------------------- +# Driver +# --------------------------------------------------------------------------- + +def run_model(name, args, out_dir): + builder = BUILDERS[name] + log_path = os.path.join(out_dir, f"{name}.log") + with open(log_path, "w") as fh: + def w(s=""): + print(s) + fh.write(s + "\n") + + w(f"=== {name} ===") + w(f"batch={args.batch} seq_len={args.seq_len} dtype={args.dtype}") + + try: + model, inputs = builder(args.batch, args.seq_len, _DTYPE_MAP[args.dtype]) + except Exception as e: + w(f"[BUILD FAIL] {type(e).__name__}: {e}") + return {"name": name, "status": "BUILD_FAIL", "ops": [], "fail_op": str(e)} + + # Phase 1 + w("\n[Phase 1] FX op enumeration (eager backend, no NPU)") + try: + ops, graph_sizes = enumerate_ops(model, inputs) + w(f" graphs: {len(graph_sizes)} total_nodes_per_graph: {graph_sizes}") + w(f" unique aten ops: {len(ops)}") + for op in ops: + w(f" {op}") + except Exception: + tb = traceback.format_exc() + w("[Phase 1 FAIL]\n" + tb) + ops = [] + + if args.enumerate_only: + return {"name": name, "status": "ENUM_ONLY", "ops": ops, "fail_op": None} + + # Phase 2 + w("\n[Phase 2] torch.compile on npu:0 + forward") + try: + status, fail_op, msg = run_on_npu(model, inputs) + w(f" status: {status}") + return {"name": name, "status": status, "ops": ops, "fail_op": None} + except Exception: + tb = traceback.format_exc() + hits, msg = parse_failure(tb) + w(" status: FAIL") + if msg: + w(f" NotImplemented message: {msg}") + if hits: + w(f" aten ops in traceback (first = most likely culprit):") + for h in hits[:10]: + w(f" {h}") + w("\n----- traceback -----\n" + tb) + return { + "name": name, + "status": "FAIL", + "ops": ops, + "fail_op": hits[0] if hits else "?", + "msg": msg, + } + + +_DTYPE_MAP = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16} + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--models", nargs="+", default=list(BUILDERS.keys()), + choices=list(BUILDERS.keys())) + p.add_argument("--batch", type=int, default=1) + p.add_argument("--seq-len", type=int, default=32) + p.add_argument("--dtype", default="float32", choices=list(_DTYPE_MAP.keys())) + p.add_argument("--enumerate-only", action="store_true", + help="Skip NPU compile; just list aten ops per model (fast).") + p.add_argument("--out-dir", default=None) + args = p.parse_args() + + ts = _dt.datetime.now().strftime("%Y%m%d_%H%M%S") + out_dir = args.out_dir or os.path.join( + os.environ.get("TORCHSIM_LOG_PATH", os.path.join(REPO_ROOT, "togsim_results")), + "op_coverage", ts, + ) + os.makedirs(out_dir, exist_ok=True) + print(f"Output dir: {out_dir}") + + results = [] + for name in args.models: + try: + results.append(run_model(name, args, out_dir)) + except KeyboardInterrupt: + print(f"[interrupt] aborted during {name}") + break + except Exception: + traceback.print_exc() + results.append({"name": name, "status": "DRIVER_ERR", "ops": [], "fail_op": None}) + + # Summary + summary_path = os.path.join(out_dir, "summary.txt") + with open(summary_path, "w") as fh: + def w(s=""): + print(s) + fh.write(s + "\n") + w("\n========== SUMMARY ==========") + w(f"{'model':10s} {'ops':>5s} {'status':10s} first_fail") + for r in results: + w(f"{r['name']:10s} {len(r['ops']):>5d} {r['status']:10s} {r.get('fail_op') or '-'}") + # Union & overlap across models + all_ops = set() + for r in results: + all_ops.update(r["ops"]) + w(f"\nUnion of aten ops across all models: {len(all_ops)}") + w("Per-model op set diff (ops unique to this model):") + for r in results: + others = set().union(*(set(r2["ops"]) for r2 in results if r2 is not r)) + unique = sorted(set(r["ops"]) - others) + w(f" {r['name']}: {len(unique)} unique") + for op in unique: + w(f" {op}") + + print(f"\nWrote: {summary_path}") + + +if __name__ == "__main__": + main() diff --git a/tests/models/DeepSeek/test_deepseek_v3_base.py b/tests/models/DeepSeek/test_deepseek_v3_base.py index 5005b70b..84fb5cf8 100644 --- a/tests/models/DeepSeek/test_deepseek_v3_base.py +++ b/tests/models/DeepSeek/test_deepseek_v3_base.py @@ -199,6 +199,11 @@ def run_deepseek_v3_base( config.quantization_config = None config = _maybe_scale_config(config, scale=scale, max_layers=max_layers) + # Seed the global RNG so config-random weight init is deterministic. Without + # this every run builds a different network, so the worst-element NPU-vs-CPU + # error randomly crosses the (loose) allclose threshold and the test is flaky. + torch.manual_seed(0) + if init_mode == "config-random": model = AutoModelForCausalLM.from_config( config=config, diff --git a/tests/ops/view/test_floormod_axis_split.py b/tests/ops/view/test_floormod_axis_split.py new file mode 100644 index 00000000..10ebd114 --- /dev/null +++ b/tests/ops/view/test_floormod_axis_split.py @@ -0,0 +1,122 @@ +"""Floor/mod index handling: axis-split (aligned) + graph-copy (incompatible). + +Covers the index-expression shapes that view/reshape/tile/group ops produce and +how the frontend handles them: + + - aligned floor/mod (single iter var, divisor divides extent): removed by + axis-split at the scheduling layer (TORCHSIM_AXIS_SPLIT). group_norm, repeat, + repeat_interleave, permute+reshape (mixed-radix). + - incompatible radices on a shared axis (case 5, e.g. a[c//2] + b[c%3]): the + conflicting operand is realized by graph-copy (TORCHSIM_GRAPH_COPY) so the + consumer reads it affine and the remainder is axis-split's. + - cross-axis / multi-variable floor/mod argument (case 7, e.g. (3*p0+p1)//4 from + a transpose+reshape feeding a broadcast/softmax/layernorm that keeps the dims + separate): graph-copy materializes the multi-var operand with copy_input (which + forces a copy of a view, unlike realize()); the copy kernel iterates the + operand's own shape so its index collapses to single-var for axis-split. + +The features are env-gated; this test turns them on for itself. axis-split is read +per kernel from the env; graph-copy installs its lowering hook at import, so we +re-run install() after setting the flag. + +Not in the CI allowlist (pytorchsim_test.yml) -- local feature/regression test. +""" +import os +import sys + +import torch +import torch.nn.functional as F + +sys.path.insert(0, os.path.join(os.environ.get("TORCHSIM_DIR", default="/workspace/PyTorchSim"), "tests")) +from _pytorchsim_utils import test_result + +os.environ.setdefault("TORCHSIM_AXIS_SPLIT", "1") +os.environ.setdefault("TORCHSIM_GRAPH_COPY", "1") +from PyTorchSimFrontend.mlir import graph_copy +graph_copy.install() + + +def _run(device, name, fn, *inputs): + torch.manual_seed(0) + opt = torch.compile(dynamic=False)(fn) + res = opt(*[t.to(device=device) for t in inputs]) + ref = fn(*[t.cpu() for t in inputs]) + test_result(name, res, ref, rtol=1e-3, atol=1e-3) + + +# --- aligned floor/mod: handled by axis-split --------------------------------- +def test_group_norm(device): + _run(device, "group_norm c//(C/G)", lambda x: F.group_norm(x, 3), torch.randn(2, 6, 4, 4)) + + +def test_repeat(device): + # tile -> ModularIndexing(c, 1, n) + _run(device, "repeat (mod)", lambda x: x.repeat(1, 2) + 1.0, torch.randn(4, 8)) + + +def test_repeat_interleave(device): + # -> FloorDiv(c, k) + _run(device, "repeat_interleave (floor)", + lambda x: torch.repeat_interleave(x, 2, dim=1) + 1.0, torch.randn(2, 4, 8)) + + +def test_permute_reshape(device): + # permute+reshape -> single-var mixed-radix floor/mod + _run(device, "permute+reshape (mixed-radix)", + lambda x: x.permute(0, 2, 1).reshape(2, 12) + 1.0, torch.randn(2, 3, 4)) + + +def test_three_level_mixed_radix(device): + # reshape+permute+reshape -> chain [1,4,12,24]; the 3-level split leaves a + # residual FloorDiv that simplify_with_ranges cannot fold -> _fold_with_ranges. + _run(device, "3-level mixed-radix", + lambda x: x.reshape(2, 3, 2, 4).permute(0, 2, 1, 3).reshape(2, 24) + 1.0, + torch.randn(2, 6, 4)) + + +def test_pixel_shuffle(device): + # splits two spatial axes -> would be 5D; the rank guard skips the split and + # falls back to baseline (the >4D decompose-peel/TOG path is #258). + _run(device, "pixel_shuffle (rank guard)", + lambda x: F.pixel_shuffle(x, 2) + 1.0, torch.randn(1, 8, 4, 4)) + + +# --- incompatible radices (case 5): handled by graph-copy --------------------- +def test_incompatible_radix(device): + # a[c//2] + b[c%3] on axis c=6 : floor-by-2 vs mod-by-3 (not a chain) + _run(device, "incompat a[c//2]+b[c%3]", + lambda a, b: torch.repeat_interleave(a, 2, dim=1) + b.repeat(1, 2), + torch.randn(2, 3), torch.randn(2, 3)) + + +# --- cross-axis multi-var floor/mod (case 7): handled by graph-copy copy_input - +def test_case7_reshape_broadcast(device): + # (3*p0+p1)//4 from transpose+reshape feeding an elementwise broadcast consumer + _run(device, "case7 reshape+broadcast", + lambda x, y: x.t().reshape(8, 3) + y, torch.randn(4, 6), torch.randn(8, 1)) + + +def test_case7_softmax_reshape(device): + # same multi-var floor feeding a reduction (softmax over the kept-separate dim) + _run(device, "case7 softmax(reshape)", + lambda x: F.softmax(x.t().reshape(8, 3), dim=1), torch.randn(4, 6)) + + +def test_case7_layernorm_reshape(device): + _run(device, "case7 layernorm(reshape)", + lambda x: F.layer_norm(x.t().reshape(8, 3), (3,)), torch.randn(4, 6)) + + +if __name__ == "__main__": + device = torch.device("npu:0") + with torch.no_grad(): + test_group_norm(device) + test_repeat(device) + test_repeat_interleave(device) + test_permute_reshape(device) + test_three_level_mixed_radix(device) + test_pixel_shuffle(device) + test_incompatible_radix(device) + test_case7_reshape_broadcast(device) + test_case7_softmax_reshape(device) + test_case7_layernorm_reshape(device) diff --git a/tests/test_mlir_bindings.py b/tests/test_mlir_bindings.py new file mode 100644 index 00000000..a0e5055d --- /dev/null +++ b/tests/test_mlir_bindings.py @@ -0,0 +1,56 @@ +"""Exercise the MLIR Python bindings the way a decompose-transfer pass would: +parse a custom op, read its AffineMap attr, build an scf.for loop with +affine.apply + an inner (unregistered) DMA op, erase the original, re-verify. +""" +from mlir.ir import (Context, Module, Location, InsertionPoint, Operation, + IndexType, IntegerAttr, AffineMap) +from mlir.dialects import scf, affine, arith, func, memref + +ctx = Context() +ctx.allow_unregistered_dialects = True + +with ctx, Location.unknown(): + src = ''' + func.func @kernel(%dram: memref<256x256xf16>, %sram: memref<128x128xf16, 1>) { + "togsim.transfer"(%dram, %sram) { + dma_kind = "MVIN", + src_map = affine_map<(d0, d1) -> (d0, d1 floordiv 16, d1 mod 16)> + } : (memref<256x256xf16>, memref<128x128xf16, 1>) -> () + return + } + ''' + m = Module.parse(src) + print("[1] parsed module ok") + + fn = m.body.operations[0] + blk = fn.regions[0].blocks[0] + transfer = next(op.operation for op in blk.operations + if op.operation.name == "togsim.transfer") + print("[2] found op:", transfer.name) + + src_map = transfer.attributes["src_map"] + print("[3] src_map attr:", src_map) + + idx = IndexType.get() + def cst(v): + return Operation.create("arith.constant", results=[idx], + attributes={"value": IntegerAttr.get(idx, v)}).result + + with InsertionPoint(transfer): + lb, ub, step = cst(0), cst(2), cst(1) + loop = scf.ForOp(lb, ub, step) + with InsertionPoint(loop.body): + iv = loop.induction_variable + base = affine.AffineApplyOp(AffineMap.get_identity(1), [iv]) + Operation.create("togsim.dma_descriptor", + operands=[base.result], results=[]) + scf.YieldOp([]) + print("[4] built scf.for + affine.apply + inner op") + + transfer.erase() + print("[5] erased original transfer") + + print("[6] verify:", m.operation.verify()) + print("----- rewritten IR -----") + print(str(m)) +print("ALL GOOD") diff --git a/thirdparty/github-releases.json b/thirdparty/github-releases.json index ec89c24f..5b012178 100644 --- a/thirdparty/github-releases.json +++ b/thirdparty/github-releases.json @@ -8,7 +8,7 @@ }, "llvm_project": { "repository": "PSAL-POSTECH/llvm-project", - "release_tag": "v1.0.8", + "release_tag": "v1.0.10", "asset_name": "riscv-llvm-release.tar.gz" }, "spike": {