Every transform is a .so with a stable C ABI

A SASS→SASS transform in sass2mlir is not a branch of the toolchain — it’s a shared library you build against the installed headers and drop into a directory. sass2mlir apply lifts the cubin once, dlopens every .so in --passes, and calls each plugin per matched function. The plugin never sees bytes, ELF sections, or the config file; it sees one lifted sass.func at a time, through an opaque handle.

The ABI is deliberately small — three entry points, C linkage, no MLIR types on the boundary:

// sass2mlir/plugin.h — the entire stable ABI a plugin implements.
#include <stdint.h>

typedef struct sass_func_s*        sass_func;         // opaque: one lifted sass.func
typedef struct sass_pass_config_s  sass_pass_config;  // opaque: this run's targeting

enum { SASS_PASS_REPORT = 0, SASS_PASS_APPLY = 1 };

typedef struct {
  int32_t     rewritten;   // ops this run rewrote (or would rewrite, in report mode)
  const char* message;     // optional note for the apply log; plugin-owned
} sass_pass_result;

const char*      sass_pass_name(void);        // must match the [[pass]] name in pass.toml
uint32_t         sass_pass_abi_version(void); // return SASS_PASS_ABI_VERSION
sass_pass_result sass_pass_run(sass_func f, const sass_pass_config* cfg);

apply unwraps the opaque sass_func to a sass::FuncOp on the plugin side of the boundary — the headers ship that cast — so inside sass_pass_run you write an ordinary MLIR walk. The config handle answers one question that matters: report or apply. In report mode the pass counts where it would fire and changes nothing; in apply mode it edits. The engine enforces the default — a pass not named mode = "apply" in the config never sees apply — but the mode is passed to the plugin so the pass itself can stay honest.

Because the ABI is versioned and opaque, a plugin built against one release keeps loading across dialect changes that don’t touch the version; when the boundary does change, the version check fails loudly at load instead of corrupting a kernel at runtime.

natural-yield: fix the yield bit, nothing else

The running example is the scheduling transform from the overview roadmap: the DeepGEMM v1.x post-pass, and the Natural Yield strategy of Yan, Wang, and Chu (PPoPP ‘20), made composable. Volta-and-later SASS carries a per-instruction control code with a 1-bit warp-scheduler yield flag: yield=1 keeps issuing from the current warp, yield=0 prefers switching warps at the cost of a cycle and the register reuse cache. ptxas scatters yield=0 on a blind every-7–8-instruction heuristic. The natural strategy: keep issuing by default, and yield only where the instruction’s own stall count says the warp actually stalls. Yan et al. report about 10% higher throughput from tuning this flag alone — their number, from their Winograd kernels; this page adds no benchmark of its own.

In sass2mlir, yield, stall count, barriers, and reuse are first-class fields of the control-code attribute, so the whole pass is a one-field edit per instruction:

// natural_yield.cpp — a complete apply plugin.
#include <sass2mlir/plugin.h>   // the C ABI above
#include <sass2mlir/dialect.h>  // sass dialect + control-code API

using namespace mlir;

// The transform: an ordinary walk over one lifted function. Returns the
// number of instructions whose yield bit disagrees with the natural rule;
// edits them only in apply mode.
static unsigned naturalYield(sass::FuncOp fn, bool apply) {
  unsigned hits = 0;

  fn.walk([&](Operation* op) {
    auto ctrl = sass::getControlCode(op);
    if (!ctrl) return;                    // sass.inst carrier (undocumented
                                          // encoding): nothing to reason
                                          // about — leave it byte-exact.

    // Natural yield: keep issuing (yield=1) unless a real stall follows
    // this instruction, in which case prefer switching warps (yield=0).
    bool keepIssuing = (ctrl->stallCount() == 0);
    if (ctrl->yield() == keepIssuing) return;  // already right; don't touch.

    ++hits;
    if (apply)
      sass::setYield(op, keepIssuing);    // One-field edit on the control-
                                          // code attribute. This invalidates
                                          // $raw on exactly this op: at lower
                                          // time nvisa_encode re-encodes it
                                          // structurally, while every
                                          // untouched neighbor emits its
                                          // $raw verbatim.
  });
  return hits;
}

// --- ABI glue: the only code that knows about the plugin boundary. ---

extern "C" const char* sass_pass_name(void) { return "natural-yield"; }
extern "C" uint32_t sass_pass_abi_version(void) { return SASS_PASS_ABI_VERSION; }

extern "C" sass_pass_result sass_pass_run(sass_func f, const sass_pass_config* cfg) {
  sass::FuncOp fn = sass::unwrap(f);      // opaque handle → dialect op
  bool apply = cfg && sass::passMode(cfg) == SASS_PASS_APPLY;
  unsigned n = naturalYield(fn, apply);
  return { static_cast<int32_t>(n), nullptr };
}

Two things are doing the safety work here. getControlCode returns nothing for sass.inst carriers — instructions the toolchain hasn’t positively identified — so the pass physically cannot touch an encoding it doesn’t understand. And setYield is the only mutation: it invalidates $raw on that one op, which forces nvisa_encode to prove it can reproduce the bytes from the semantic fields, while the other 80-odd percent of the function lowers verbatim from $raw. An edit can never silently corrupt its neighbors.

%41 = sass.FFMA_0x223.FTZ guard(%38) %27, %14, %40
        : (!sass.reg, !sass.reg, !sass.reg) -> (!sass.reg)
        { ctrl = #sass.ctrl<stall = 0, yield = 0, wr_bar = 0x0, rd_bar = 0x0, reuse = 0x1>,
          raw  = #sass.raw<0x00000000e22a0223, 0x004e200000710229> }

In disassembly terms, the whole transform is one bit per fired site:

- /* 02f0 */ FFMA.FTZ R4, R5, R6, R4   ; // ctrl: stall=0 yield=0 reuse=1  ← heuristic, no stall
+ /* 02f0 */ FFMA.FTZ R4, R5, R6, R4   ; // ctrl: stall=0 yield=1 reuse=1  ← keep issuing

One .so against the installed headers

A plugin is an ordinary shared library linked against the sass2mlir::pass-api interface target the toolchain installs. The whole build file:

cmake_minimum_required(VERSION 3.20)
project(natural-yield CXX)

find_package(sass2mlir REQUIRED)   # installed with the toolchain

add_library(natural-yield SHARED natural_yield.cpp)
target_link_libraries(natural-yield PRIVATE sass2mlir::pass-api)
target_compile_features(natural-yield PRIVATE cxx_std_17)

# apply dlopens every .so in one directory — install there:
install(TARGETS natural-yield
        LIBRARY DESTINATION lib/sass2mlir/passes)

cmake -B build -DCMAKE_PREFIX_PATH=/opt/sass2mlir
cmake --build build
cmake --install build --prefix /opt/sass2mlir
ls /opt/sass2mlir/lib/sass2mlir/passes/
# libscalar-to-uniform.so  libffma-deftz.so  libnatural-yield.so  ...

No symbol visibility flags, no registration macros: dlopen plus the three extern "C" entry points is the whole contract. If sass_pass_abi_version() doesn’t match the host binary, apply prints the plugin’s path and the mismatch and refuses to load it — it does not fall back, guess, or run the pass anyway.

Report first, apply on purpose

Passes are selected and targeted by config, per architecture and per function — the same [[pass]] blocks the CLI reference defines. Target a Winograd-style kernel on an Ampere cubin:

# yield.toml
[[pass]]
name = "natural-yield"          # must match sass_pass_name()
sm   = ["sm_8*"]                # arch globs
fn   = ["*winograd*"]           # function globs
mode = "report"                 # the default; shown here to be explicit

Run the engine. In report mode the pass walks every matched function and counts its hits, but the mode it receives is report, so $raw is never invalidated and no output cubin is written:

sass2mlir apply winograd.cubin \
  --passes /opt/sass2mlir/lib/sass2mlir/passes \
  --config yield.toml \
  -o winograd.patched.cubin
# apply: sm_86, 4 functions, 1 pass loaded (natural-yield, ABI v3)
# pass natural-yield [report]: winograd_4x4_kernel — would rewrite 212 of 1,394 instructions
# pass natural-yield [report]: batched_winograd_2x2 — would rewrite 96 of 602 instructions
# apply: 2 functions matched, 308 instructions flagged, 0 bytes changed (report mode)
# apply: no output written — flip mode to "apply" to transform

The report is the review step: 308 instructions across two kernels, all of them single-bit control-code edits. If the count or the sites look wrong, the cubin on disk is untouched. When it looks right, flip the one word:

mode = "apply"

sass2mlir apply winograd.cubin \
  --passes /opt/sass2mlir/lib/sass2mlir/passes \
  --config yield.toml \
  -o winograd.patched.cubin
# pass natural-yield [apply]: winograd_4x4_kernel — rewrote 212 of 1,394 instructions (nvisa_encode)
# pass natural-yield [apply]: batched_winograd_2x2 — rewrote 96 of 602 instructions (nvisa_encode)
# apply: patched .text + .nv.info for 2 functions → winograd.patched.cubin
# apply: 2 functions matched by no pass — byte-identical in the output

That last line is the targeting guarantee: the two kernels no fn glob named are copied through byte-identical, and within the matched kernels every instruction the pass didn’t touch lowers from its original $raw. The output is a self-contained cubin in the original kernel slots — same ABI, same launch sites.

Prove the edit, then hold the bar

After apply, roundtrip against the original cubin is no longer the check — the bytes are supposed to differ now, which is exactly why report mode is the default. The validation that replaces it has three parts.

1. Keep the IR, audit the encoder. --keep-mlir writes the module before and after the pass; crosscheck on the transformed module re-encodes every instruction from its structural fields alone and compares against the bytes it will actually emit. For the 308 edited ops this is the proof that nvisa_encode reproduces the control-code edit exactly:

sass2mlir apply winograd.cubin \
  --passes /opt/sass2mlir/lib/sass2mlir/passes \
  --config yield.toml \
  --keep-mlir keep/ \
  -o winograd.patched.cubin
ls keep/
# winograd_4x4_kernel.lifted.mlir       winograd_4x4_kernel.transformed.mlir
# batched_winograd_2x2.lifted.mlir      batched_winograd_2x2.transformed.mlir

sass2mlir crosscheck keep/winograd_4x4_kernel.transformed.mlir
# crosscheck: sm_86 winograd_4x4_kernel: 1394/1394 instructions re-encode exactly (0 mismatches)

2. Byte-diff the streams. Every differing byte must sit inside an instruction the report named — nothing else may move:

sass2mlir harvest winograd.cubin         -o before/
sass2mlir harvest winograd.patched.cubin -o after/

cmp -l before/winograd_4x4_kernel.code after/winograd_4x4_kernel.code | wc -l
# 212
# one byte per rewritten instruction — the control-code byte carrying the
# yield bit — at exactly the 212 offsets the report listed

3. Hold the production bar. Byte-diff and crosscheck prove this run did what it said. They don’t prove the transform sound — that a warp yielding only on real stalls preserves behavior for every kernel the glob will ever match. The bar for that in this project is the corpus: model the transform, validate it against every kernel the glob can match, then ship it as a plugin — the treatment scalar-to-uniform got and the treatment each roadmap pass gets before it’s called done. natural-yield as written here is a tutorial pass: validated by report, byte-diff, and crosscheck. The engine underneath it is regression-gated across 489,344 patched functions; the pass-specific correctness argument is yours to make at corpus scale.