6. Real Bufferization: Closing Chapter 4's Own Honest Gap¶

Mountain goats in space helmets on Mercury

What you will understand: how to make MLIR's own real, built-in --one-shot-bufferize pass genuinely succeed on a custom dialect, by implementing the real BufferizableOpInterface -- the exact extension point Chapter 4 named but did not use, when it ran into a real, honest failure (error: op was not bufferized) and deferred the fix. This chapter closes that gap for real, and the result is a second, entirely independent way to get Mountain Goat's own tensor IR down to memrefs -- not a replacement for Chapter 4's hand-written pass, a second real path, compared directly against the first.

What you need to know first: Chapter 4's own ConvertMgToAffinePass and its real negative result against --one-shot-bufferize; Chapter 3's own constant-folding discussion; Chapter 5's own printMemrefF64 mechanism, reused directly in this chapter rather than reinvented.

Compile and run

cd docs/part6/code
./build.sh        # -> ./build/mg-opt (with the bufferization models)
./run.sh          # one-shot-bufferize on the examples, then lower and run

Same prerequisites as Chapter 5. Every listing and output on this page comes from these commands (and the chapter's own embedded files).

The real question this chapter answers¶

Chapter 4 cited MLIR's own real documentation stating plainly what BufferizableOpInterface requires (bufferizesToMemoryRead, bufferizesToMemoryWrite, at minimum) and predicted, without yet testing it, that implementing the real interface would be more work than the hand-written ConvertMgToAffinePass -- because One-Shot Bufferize's own whole-function analysis asks each op to correctly answer questions about aliasing and in-place reuse that a simple, always-allocate-fresh pass never has to consider. This chapter tests that prediction directly: implement the real interface for all four mg operations, and find out, by actually running it, whether --one-shot-bufferize now succeeds -- and what, concretely, the extra work turned out to be.

Four real external models, one per operation¶

MLIR's own real mechanism for attaching an interface to an operation after it is already defined -- exactly Mountain Goat's own situation, since mg's ODS definitions (Chapter 2) were written before this chapter existed -- is an "external model," registered once via a DialectRegistry extension rather than baked into the .td file itself.

`mg.constant`: the simplest real case¶

struct ConstantOpInterface
    : public BufferizableOpInterface::ExternalModel<ConstantOpInterface,
                                                     mg::ConstantOp> {
  bool bufferizesToAllocation(Operation *op, Value value) const {
    return true;
  }

  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                          const BufferizationOptions &options) const {
    auto constantOp = cast<mg::ConstantOp>(op);
    auto loc = op->getLoc();
    auto tensorType =
        cast<RankedTensorType>(constantOp.getResult().getType());
    auto memrefType =
        MemRefType::get(tensorType.getShape(), tensorType.getElementType());
    auto alloc = rewriter.create<memref::AllocOp>(loc, memrefType);

    llvm::SmallVector<double> values;
    for (auto v : constantOp.getValue().getValues<llvm::APFloat>())
      values.push_back(v.convertToDouble());

    auto shape = tensorType.getShape();
    int64_t idx = 0;
    for (int64_t i = 0; i < shape[0]; ++i) {
      for (int64_t j = 0; j < shape[1]; ++j) {
        Value cst = rewriter.create<arith::ConstantOp>(
            loc, rewriter.getF64FloatAttr(values[idx++]));
        Value iIdx = rewriter.create<arith::ConstantIndexOp>(loc, i);
        Value jIdx = rewriter.create<arith::ConstantIndexOp>(loc, j);
        rewriter.create<affine::AffineStoreOp>(loc, cst, alloc,
                                               ValueRange{iIdx, jIdx});
      }
    }

    replaceOpWithBufferizedValues(rewriter, op, alloc.getResult());
    return success();
  }
};

mg.constant has no tensor operands at all -- only a result, produced from an attribute it already owns outright. bufferizesToAllocation returning true is this chapter's own first real, required answer: it tells One-Shot Bufferize's own analysis that this value is a genuinely fresh buffer, never aliasing anything else in the function. The bufferize() body itself is, element for element, the same real logic as Chapter 4's own ConstantOpLowering -- a different real entry point into MLIR's own machinery, not different work.

`mg.add` and `mg.transpose`: real answers about reading, writing, and aliasing¶

struct AddOpInterface
    : public BufferizableOpInterface::ExternalModel<AddOpInterface,
                                                     mg::AddOp> {
  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                              const AnalysisState &state) const {
    return true;
  }
  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
    return false;
  }
  bool bufferizesToAllocation(Operation *op, Value value) const {
    return true;
  }
  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                      const AnalysisState &state) const {
    return {};
  }

  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                          const BufferizationOptions &options) const {
    auto addOp = cast<mg::AddOp>(op);
    auto loc = op->getLoc();

    FailureOr<Value> lhsBuffer = getBuffer(rewriter, addOp.getLhs(), options);
    if (failed(lhsBuffer))
      return failure();
    FailureOr<Value> rhsBuffer = getBuffer(rewriter, addOp.getRhs(), options);
    if (failed(rhsBuffer))
      return failure();

    auto tensorType = cast<RankedTensorType>(addOp.getResult().getType());
    auto memrefType =
        MemRefType::get(tensorType.getShape(), tensorType.getElementType());
    auto alloc = rewriter.create<memref::AllocOp>(loc, memrefType);
    auto shape = tensorType.getShape();

    affine::buildAffineLoopNest(
        rewriter, loc, {0, 0}, {shape[0], shape[1]}, {1, 1},
        [&](OpBuilder &nestedBuilder, Location nestedLoc, ValueRange ivs) {
          Value lhsVal = nestedBuilder.create<affine::AffineLoadOp>(
              nestedLoc, *lhsBuffer, ivs);
          Value rhsVal = nestedBuilder.create<affine::AffineLoadOp>(
              nestedLoc, *rhsBuffer, ivs);
          Value sum =
              nestedBuilder.create<arith::AddFOp>(nestedLoc, lhsVal, rhsVal);
          nestedBuilder.create<affine::AffineStoreOp>(nestedLoc, sum, alloc,
                                                      ivs);
        });

    replaceOpWithBufferizedValues(rewriter, op, alloc.getResult());
    return success();
  }
};

Four real, specific claims, each one load-bearing for the analysis, not boilerplate: bufferizesToMemoryRead (true -- mg.add genuinely reads through both operands), bufferizesToMemoryWrite (false -- it never writes through either operand's own buffer), bufferizesToAllocation (true -- its result is always a fresh buffer), and getAliasingValues (empty -- the result aliases neither operand). One real, concrete consequence of getting any of these wrong: a false answer where true was required could let One-Shot Bufferize's own analysis permit an in-place reuse that silently corrupts a buffer still needed elsewhere -- exactly why MLIR's own documentation, cited in Chapter 3, insists these methods "must be implemented conservatively."

getBuffer(rewriter, addOp.getLhs(), options) is the one real, new API this chapter needed that Chapter 4's own ConversionPatternRewriter-based code never did -- Chapter 4's adaptor.getLhs() already handed back the converted operand directly, a convenience of the dialect-conversion framework specifically; One-Shot Bufferize's own bufferize() hook has to ask explicitly for each operand's own buffer.

mg.transpose's own real interface implementation is the same shape, with one operand instead of two and the same index-swap between load and store Chapter 4 already established -- omitted here for space, included in full in this chapter's own real code.

`mg.print`: reusing Chapter 5's own real mechanism, not reinventing it¶

struct PrintOpInterface
    : public BufferizableOpInterface::ExternalModel<PrintOpInterface,
                                                     mg::PrintOp> {
  bool bufferizesToMemoryRead(Operation *op, OpOperand &opOperand,
                              const AnalysisState &state) const {
    return true;
  }
  bool bufferizesToMemoryWrite(Operation *op, OpOperand &opOperand,
                               const AnalysisState &state) const {
    return false;
  }
  AliasingValueList getAliasingValues(Operation *op, OpOperand &opOperand,
                                      const AnalysisState &state) const {
    return {};
  }

  LogicalResult bufferize(Operation *op, RewriterBase &rewriter,
                          const BufferizationOptions &options) const {
    auto printOp = cast<mg::PrintOp>(op);
    auto loc = op->getLoc();
    FailureOr<Value> inputBuffer =
        getBuffer(rewriter, printOp.getInput(), options);
    if (failed(inputBuffer))
      return failure();

    auto module = op->getParentOfType<ModuleOp>();
    auto printRef = getOrInsertPrintMemrefF64(rewriter, module);
    auto memrefType = cast<MemRefType>(inputBuffer->getType());
    auto unrankedType =
        UnrankedMemRefType::get(memrefType.getElementType(), /*memorySpace=*/0);
    Value cast =
        rewriter.create<memref::CastOp>(loc, unrankedType, *inputBuffer);
    rewriter.create<func::CallOp>(loc, printRef, TypeRange{}, ValueRange{cast});
    rewriter.eraseOp(op);
    return success();
  }
};

getOrInsertPrintMemrefF64 is, verbatim, Chapter 5's own real helper -- declaring @printMemrefF64 with the real llvm.emit_c_interface attribute. Reaching for it again here, rather than writing a second version, is itself the real point: once a real mechanism for calling into MLIR's own runtime support library exists, any later real pass -- hand-written conversion or interface-based bufferization alike -- can reuse it directly.

Registering the four real external models¶

namespace mg {
void registerBufferizableOpInterfaceExternalModels(DialectRegistry &registry) {
  registry.addExtension(+[](MLIRContext *ctx, MgDialect *dialect) {
    ConstantOp::attachInterface<ConstantOpInterface>(*ctx);
    AddOp::attachInterface<AddOpInterface>(*ctx);
    TransposeOp::attachInterface<TransposeOpInterface>(*ctx);
    PrintOp::attachInterface<PrintOpInterface>(*ctx);
  });
}
} // namespace mg

mg::registerBufferizableOpInterfaceExternalModels(registry);
mlir::bufferization::func_ext::registerBufferizableOpInterfaceExternalModels(
    registry);

The second line is a real gap this chapter hit directly, the same way Chapter 4 hit the func.func signature-conversion gap: registering mg's own four interfaces was not enough to bufferize a function taking real tensor arguments -- func.func itself needs its own real external model too, which MLIR already ships, in mlir/Dialect/Bufferization/Transforms/FuncBufferizableOpInterfaceImpl.h, simply not registered by default.

Real worked example: the negative result from Chapter 4, reversed¶

Chapter 4's own exact real test, re-run against this chapter's own rebuilt mg-opt:

func.func @main() {
  %0 = mg.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf64>
  %1 = mg.constant dense<[[5.0, 6.0], [7.0, 8.0]]> : tensor<2x2xf64>
  %2 = mg.add %0, %1 : tensor<2x2xf64>, tensor<2x2xf64> -> tensor<2x2xf64>
  %3 = mg.transpose %2 : tensor<2x2xf64> to tensor<2x2xf64>
  mg.print %3 : tensor<2x2xf64>
  func.return
}

./mg-opt mountain_goat_const.mlir --one-shot-bufferize

Real captured output (relevant excerpt -- full output is longer):

6. Real Bufferization: Closing Chapter 4's Own Honest Gap¶

The real question this chapter answers¶

Four real external models, one per operation¶

`mg.constant`: the simplest real case¶

`mg.add` and `mg.transpose`: real answers about reading, writing, and aliasing¶

`mg.print`: reusing Chapter 5's own real mechanism, not reinventing it¶

Registering the four real external models¶

Real worked example: the negative result from Chapter 4, reversed¶

A real, honest surprise: no constant-folding this time¶

A second real, honest difference: function-argument buffer types¶

Real worked example: genuinely re-running the full pipeline to the same real answer¶

The complete source files¶

What later chapters changed¶

Chapter summary¶

Self-check questions¶

6. Real Bufferization: Closing Chapter 4's Own Honest Gap¶

The real question this chapter answers¶

Four real external models, one per operation¶

mg.constant: the simplest real case¶

mg.add and mg.transpose: real answers about reading, writing, and aliasing¶

mg.print: reusing Chapter 5's own real mechanism, not reinventing it¶

Registering the four real external models¶

Real worked example: the negative result from Chapter 4, reversed¶

A real, honest surprise: no constant-folding this time¶

A second real, honest difference: function-argument buffer types¶

Real worked example: genuinely re-running the full pipeline to the same real answer¶

The complete source files¶

What later chapters changed¶

Chapter summary¶

Self-check questions¶

`mg.constant`: the simplest real case¶

`mg.add` and `mg.transpose`: real answers about reading, writing, and aliasing¶

`mg.print`: reusing Chapter 5's own real mechanism, not reinventing it¶