21. More Operations: Subtract, Hadamard, Divide, Matrix Product, and Scalars¶
What you will understand: how to grow a dialect by four operations and one kind of operand, end to end: the operation definitions, their verifiers, their lowerings, the front end's syntax for them, and the tests that prove each one. You will also meet the three different things "multiply" can mean for matrices (elementwise, matrix product, and by a scalar), learn why division by zero is not an error, and read the PTX for a reduction.
What you need to know first: the language tour for the syntax, then Chapters 3 to 5 (defining and lowering an op; mg.add is the model for most of this chapter), 13 and 14 (dynamic shapes and the run-time check), and 20 (the language and the mgc driver). Each new idea is explained where it first appears.
Compile and run: the commands for this chapter¶
Every code listing and sample output below was produced by these commands, run from docs/part21/code/:
./build.sh # compile mg-opt with the new ops -> build/mg-opt
./mgc run examples/04_matmul.mg # compile one example to native code and run it
./mgc ptx examples/13_gpu_matmul.mg # compile one example to PTX (not run: no GPU here)
./demo.sh > demo_out.txt # run all examples; the output below is this file
./show.sh # write the lowered MLIR shown below -> show/
cpp/run.sh # build and run the C++ program -> cpp/run_out.txt
cd ../../part15/code && ./run_lit.sh # the whole test suite (78 tests)
Why this page shows FAIL lines, and why that is good
Parts of this page come from negative controls: the same tests run against an older build, or against deliberately broken code (a "mutation"). In those runs a FAIL is the expected, wanted result: it means a test noticed the problem, which is how we know the tests are worth anything. What would be wrong is the opposite, a deliberately broken build that passes everything. The baseline (the unmodified current build) must always show every test passing, and it does. Each script that does this says so in its own header and prints a reminder at the top of its output.
Primer: three meanings of "multiply"¶
Given two matrices A and B, three different operations are all called multiplication.
Elementwise (also the Hadamard product, after Jacques Hadamard) multiplies matching elements: result[i][j] = A[i][j] * B[i][j]. It needs the two matrices to have the same shape. It is the same kind of operation as the mg.add of Chapter 3, with * in place of +. Subtraction and division are the same again: A - B and A / B, taken element by element.
The matrix product is a different operation. For A of shape m by k and B of shape k by n, the result has shape m by n, and result[i][j] = A[i][0]*B[0][j] + A[i][1]*B[1][j] + ... + A[i][k-1]*B[k-1][j]: the dot product of row i of A with column j of B. The two shapes do not need to match; only the inner dimensions (the two k's) must. It is the operation behind most of machine learning and graphics. The summed term is called a reduction: many values are combined into one.
By a scalar combines every element with one number: A * 2 doubles every element and 10 - A subtracts every element from 10. A scalar has no shape, so there is nothing to check at run time.
Two more facts matter. First, order matters for subtraction, division and the matrix product: A - B is not B - A, and A @ B usually has a different shape than B @ A. Second, floating-point division follows the IEEE 754 standard: x / 0 is positive or negative infinity (inf), and 0 / 0 is not a number (nan). Neither is an error, and neither stops the program; the numbers simply propagate. Mountain Goat inherits this because it lowers to the hardware's own division.
Design decisions¶
One operator, one meaning. The language already used +. For the new ones: - and / are elementwise; * is the Hadamard (elementwise) product; and @ is the matrix product (the spelling Python uses). A language that overloads * to mean "elementwise or matrix product, depending on shapes" makes every reader work out which one it is; here it is always visible in the operator.
Scalars are compile-time numbers. A number like 2 or 0.5, or a let bound to one (let k = 2 + 3 * 4, folded to 14 by the front end), combines with a matrix through one new operation, mg.scalar. Scalars are not run-time values in this language: you cannot read one out of a matrix, and you cannot print or pass one to a function.
One operation for all scalar forms. There are four arithmetic operations and two operand orders (A - 1 versus 1 - A), eight combinations. Rather than eight operations, mg.scalar carries three attributes: op (one of "add", "sub", "mul", "div"), value, and reversed (false computes x OP value, true computes value OP x). The reversed flag exists because 10 - A and A - 10 differ, as do 1 / A and A / 1. For + and * the flag is harmless but unnecessary; the front end sets it whenever the scalar is written on the left.
Precedence. Tightest first: unary minus, then *, / and @ (left to right, all at one level), then + and -. So a + a * a is a + (a * a), and a - b @ a is a - (b @ a).
The operations (MgOps.td)¶
The new definitions, in the same ODS language as Chapter 3:
let results = (outs F64Tensor:$result);
let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)";
let hasVerifier = 1;
let hasFolder = 1;
}
def SubOp : Mg_Op<"sub", [Pure]> {
let summary = "elementwise subtract of two compatible-shaped tensors (lhs - rhs)";
let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);
let results = (outs F64Tensor:$result);
let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)";
let hasVerifier = 1;
}
def MulOp : Mg_Op<"mul", [Pure]> {
let summary = "elementwise (Hadamard) product of two compatible-shaped tensors";
let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);
let results = (outs F64Tensor:$result);
let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)";
let hasVerifier = 1;
}
def DivOp : Mg_Op<"div", [Pure]> {
let summary = "elementwise division of two compatible-shaped tensors (lhs / rhs, IEEE: x/0 is inf or nan)";
let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);
let results = (outs F64Tensor:$result);
let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)";
let hasVerifier = 1;
}
def ScalarOp : Mg_Op<"scalar", [Pure]> {
let summary = "combine every element with one scalar constant";
let description = [{
`op` is one of "add", "sub", "mul", "div". With `reversed = false` each element becomes
`x OP value`; with `reversed = true` it becomes `value OP x` (so `10 - x` and `1 / x` exist).
Works for any shape, including dynamic ones, because no second tensor is involved.
}];
let arguments = (ins F64Tensor:$input, StrAttr:$op, F64Attr:$value, BoolAttr:$reversed);
let results = (outs F64Tensor:$result);
let assemblyFormat = "$input attr-dict `:` type($input) `->` type($result)";
let hasVerifier = 1;
}
def NegOp : Mg_Op<"neg", [Pure]> {
let summary = "elementwise negation";
let arguments = (ins F64Tensor:$input);
let results = (outs F64Tensor:$result);
let assemblyFormat = "$input attr-dict `:` type($input) `->` type($result)";
let hasVerifier = 1;
}
def MatmulOp : Mg_Op<"matmul", [Pure]> {
let summary = "matrix product: (m x k) times (k x n) gives (m x n)";
let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);
let results = (outs F64Tensor:$result);
let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)";
let hasVerifier = 1;
}
def TransposeOp : Mg_Op<"transpose", [Pure]> {
let summary = "transpose of a 2-D tensor";
let arguments = (ins F64Tensor:$input);
let results = (outs F64Tensor:$result);
let assemblyFormat = "$input attr-dict `:` type($input) `to` type($result)";
let hasVerifier = 1;
let hasCanonicalizer = 1;
}
def PrintOp : Mg_Op<"print"> {
let summary = "print a tensor at runtime";
let arguments = (ins F64Tensor:$input);
let assemblyFormat = "$input attr-dict `:` type($input)";
}
SubOp, MulOp and DivOp are copies of AddOp's shape: two operands, one result, the same printed format. AddOp also has a constant folder; the new ones do not (the lowering still produces correct results for constant operands, just without compile-time folding). NegOp has one operand. MatmulOp has the same shape of declaration as AddOp but a different rule, in its verifier below. ScalarOp carries its three attributes in its arguments.
The verifiers (MgDialect.cpp)¶
SubOp, MulOp, DivOp, NegOp and ScalarOp all share one rule: same rank, and each dimension compatible between every pair of types, where compatible is Chapter 13's rule (equal, or ? on either side). Rather than copy AddOp::verify five times, one helper does it:
namespace {
/// Shared by mg.add-style ops: same rank, every dimension compatible (equal, or `?` on either side).
mlir::LogicalResult verifyElementwise(mlir::Operation *op, llvm::ArrayRef<mlir::Type> types, llvm::StringRef name) {
llvm::SmallVector<mlir::RankedTensorType, 3> t;
for (mlir::Type ty : types) {
auto r = llvm::dyn_cast<mlir::RankedTensorType>(ty);
if (!r) return op->emitOpError(name + " only supports ranked tensors");
t.push_back(r);
}
for (auto x : t)
if (x.getRank() != t[0].getRank())
return op->emitOpError(name + " operands and result must have the same rank");
for (int64_t d = 0; d < t[0].getRank(); ++d)
for (auto x : t)
for (auto y : t)
if (!mg::dyn::compatibleDim(x.getDimSize(d), y.getDimSize(d)))
return op->emitOpError(name + " operands must have compatible shapes, got ") << x << " and " << y;
return mlir::success();
}
} // namespace
mlir::LogicalResult SubOp::verify() { return verifyElementwise(*this, {getLhs().getType(), getRhs().getType(), getResult().getType()}, "mg.sub"); }
mlir::LogicalResult MulOp::verify() { return verifyElementwise(*this, {getLhs().getType(), getRhs().getType(), getResult().getType()}, "mg.mul"); }
mlir::LogicalResult DivOp::verify() { return verifyElementwise(*this, {getLhs().getType(), getRhs().getType(), getResult().getType()}, "mg.div"); }
mlir::LogicalResult ScalarOp::verify() {
llvm::StringRef k = getOp();
if (k != "add" && k != "sub" && k != "mul" && k != "div")
return emitOpError("mg.scalar: op must be one of add, sub, mul, div, got '") << k << "'";
return verifyElementwise(*this, {getInput().getType(), getResult().getType()}, "mg.scalar");
}
mlir::LogicalResult NegOp::verify() { return verifyElementwise(*this, {getInput().getType(), getResult().getType()}, "mg.neg"); }
mlir::LogicalResult MatmulOp::verify() {
auto l = llvm::dyn_cast<mlir::RankedTensorType>(getLhs().getType());
auto r = llvm::dyn_cast<mlir::RankedTensorType>(getRhs().getType());
auto o = llvm::dyn_cast<mlir::RankedTensorType>(getResult().getType());
if (!l || !r || !o) return emitOpError("mg.matmul only supports ranked tensors");
if (l.getRank() != 2 || r.getRank() != 2 || o.getRank() != 2)
return emitOpError("mg.matmul only supports rank-2 tensors");
if (!mg::dyn::compatibleDim(l.getDimSize(1), r.getDimSize(0)))
return emitOpError("mg.matmul inner dimensions differ: ") << l << " times " << r;
if (!mg::dyn::compatibleDim(o.getDimSize(0), l.getDimSize(0)) || !mg::dyn::compatibleDim(o.getDimSize(1), r.getDimSize(1)))
return emitOpError("mg.matmul result shape must be (rows of lhs) x (columns of rhs), got ") << o;
return mlir::success();
}
mlir::OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
auto lhsAttr = llvm::dyn_cast_if_present<mlir::DenseElementsAttr>(adaptor.getLhs());
auto rhsAttr = llvm::dyn_cast_if_present<mlir::DenseElementsAttr>(adaptor.getRhs());
mg.matmul is the interesting one. It is not "same shape". It checks that the inner dimensions agree (lhs columns against rhs rows) and that the result is lhs rows by rhs columns, again using compatibleDim so ? is accepted. mg.scalar additionally rejects any op that is not one of the four names. (A string attribute is used instead of an enum to keep the definition short; the verifier is the guard. An enum attribute would catch the same mistake earlier, and would be the better design for a larger dialect.)
The lowerings (LowerToAffine.cpp)¶
Subtract, multiply and divide reuse the whole lowering of mg.add (Chapter 4 and 14): check shapes at run time where a dimension is ?, allocate the result, run a two-deep loop nest, load both operands, apply one arithmetic operation, store. The only thing that differs is the arithmetic operation, so one C++ template takes it as a parameter:
/// Chapter 21: sub, mul and div share add's whole lowering; only the arithmetic op differs.
template <typename MgOp, typename ArithOp>
struct ElementwiseBinaryLowering : public OpConversionPattern<MgOp> {
using OpConversionPattern<MgOp>::OpConversionPattern;
using OpAdaptor = typename OpConversionPattern<MgOp>::OpAdaptor;
LogicalResult matchAndRewrite(MgOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
auto tensorType = llvm::cast<RankedTensorType>(op.getResult().getType());
auto memrefType = convertTensorToMemref(tensorType);
Value lhs = adaptor.getLhs(), rhs = adaptor.getRhs();
mg::dyn::assertSameShape(rewriter, loc, lhs, rhs,
llvm::cast<RankedTensorType>(op.getLhs().getType()),
llvm::cast<RankedTensorType>(op.getRhs().getType()),
op.getOperationName());
llvm::SmallVector<Value, 2> extents;
for (int64_t d = 0; d < 2; ++d)
extents.push_back(tensorType.isDynamicDim(d) ? mg::dyn::extentOf(rewriter, loc, lhs, d) : Value());
auto alloc = mg::dyn::allocFor(rewriter, loc, memrefType, extents);
mg::dyn::nest2D(rewriter, loc, tensorType, extents,
[&](OpBuilder &b, Location l, ValueRange ivs) {
Value x = b.create<affine::AffineLoadOp>(l, lhs, ivs);
Value y = b.create<affine::AffineLoadOp>(l, rhs, ivs);
b.create<affine::AffineStoreOp>(l, b.create<ArithOp>(l, x, y), alloc, ivs);
});
rewriter.replaceOp(op, alloc.getResult());
return success();
}
};
The patterns are registered with ElementwiseBinaryLowering<mg::SubOp, arith::SubFOp>, and likewise for MulOp/MulFOp and DivOp/DivFOp. The run-time check's message names the operation (mg.sub: operand shapes differ at runtime in dimension 0), because the template passes op.getOperationName().
Scalar and negate have one input, so there is no second operand to check:
struct ScalarOpLowering : public OpConversionPattern<mg::ScalarOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult matchAndRewrite(mg::ScalarOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
auto tensorType = llvm::cast<RankedTensorType>(op.getResult().getType());
auto memrefType = convertTensorToMemref(tensorType);
Value in = adaptor.getInput();
llvm::SmallVector<Value, 2> extents;
for (int64_t d = 0; d < 2; ++d)
extents.push_back(tensorType.isDynamicDim(d) ? mg::dyn::extentOf(rewriter, loc, in, d) : Value());
auto alloc = mg::dyn::allocFor(rewriter, loc, memrefType, extents);
Value c = rewriter.create<arith::ConstantOp>(loc, rewriter.getF64FloatAttr(op.getValue().convertToDouble()));
llvm::StringRef kind = op.getOp();
bool rev = op.getReversed();
mg::dyn::nest2D(rewriter, loc, tensorType, extents, [&](OpBuilder &b, Location l, ValueRange ivs) {
Value x = b.create<affine::AffineLoadOp>(l, in, ivs);
Value lhs = rev ? c : x, rhs = rev ? x : c, r;
if (kind == "add") r = b.create<arith::AddFOp>(l, lhs, rhs);
else if (kind == "sub") r = b.create<arith::SubFOp>(l, lhs, rhs);
else if (kind == "mul") r = b.create<arith::MulFOp>(l, lhs, rhs);
else r = b.create<arith::DivFOp>(l, lhs, rhs);
b.create<affine::AffineStoreOp>(l, r, alloc, ivs);
});
rewriter.replaceOp(op, alloc.getResult());
return success();
}
};
struct NegOpLowering : public OpConversionPattern<mg::NegOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult matchAndRewrite(mg::NegOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
auto tensorType = llvm::cast<RankedTensorType>(op.getResult().getType());
auto memrefType = convertTensorToMemref(tensorType);
Value in = adaptor.getInput();
llvm::SmallVector<Value, 2> extents;
for (int64_t d = 0; d < 2; ++d)
extents.push_back(tensorType.isDynamicDim(d) ? mg::dyn::extentOf(rewriter, loc, in, d) : Value());
auto alloc = mg::dyn::allocFor(rewriter, loc, memrefType, extents);
mg::dyn::nest2D(rewriter, loc, tensorType, extents,
[&](OpBuilder &b, Location l, ValueRange ivs) {
Value x = b.create<affine::AffineLoadOp>(l, in, ivs);
b.create<affine::AffineStoreOp>(l, b.create<arith::NegFOp>(l, x), alloc, ivs);
});
rewriter.replaceOp(op, alloc.getResult());
return success();
}
};
The scalar is turned into one arith.constant outside the loop, and reversed decides which side of the arithmetic operation it goes on. You can see the result in the lowered code for 10 - x:
module {
func.func @f(%arg0: memref<2x2xf64>) -> memref<2x2xf64> {
%alloc = memref.alloc() : memref<2x2xf64>
%cst = arith.constant 1.000000e+01 : f64
affine.for %arg1 = 0 to 2 {
affine.for %arg2 = 0 to 2 {
%0 = affine.load %arg0[%arg1, %arg2] : memref<2x2xf64>
%1 = arith.subf %cst, %0 : f64
affine.store %1, %alloc[%arg1, %arg2] : memref<2x2xf64>
}
}
return %alloc : memref<2x2xf64>
}
}
(arith.subf %cst, %0 with the constant first: that is the reversed form.)
Matrix product needs three nested loops and a run-time check of the inner dimension. The result is first filled with zeros, because the main loop accumulates into it (result[i][j] += lhs[i][k] * rhs[k][j]), and an uninitialized buffer holds garbage:
/// mg.matmul: result[i,j] = sum over k of lhs[i,k] * rhs[k,j]. A zero-fill nest, then a 3-deep
/// accumulate nest. For `?` dimensions the inner extents must agree at run time: one cf.assert.
struct MatmulOpLowering : public OpConversionPattern<mg::MatmulOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult matchAndRewrite(mg::MatmulOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
auto loc = op.getLoc();
auto outType = llvm::cast<RankedTensorType>(op.getResult().getType());
auto lt = llvm::cast<RankedTensorType>(op.getLhs().getType());
auto rt = llvm::cast<RankedTensorType>(op.getRhs().getType());
auto memrefType = convertTensorToMemref(outType);
Value lhs = adaptor.getLhs(), rhs = adaptor.getRhs();
if (lt.isDynamicDim(1) || rt.isDynamicDim(0)) {
Value l = mg::dyn::extentOf(rewriter, loc, lhs, 1), r = mg::dyn::extentOf(rewriter, loc, rhs, 0);
Value eq = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, l, r);
rewriter.create<cf::AssertOp>(loc, eq, rewriter.getStringAttr("mg.matmul: inner dimensions differ at runtime"));
}
// m = rows of lhs, n = columns of rhs, k = columns of lhs.
llvm::SmallVector<Value, 2> extents = {outType.isDynamicDim(0) ? mg::dyn::extentOf(rewriter, loc, lhs, 0) : Value(),
outType.isDynamicDim(1) ? mg::dyn::extentOf(rewriter, loc, rhs, 1) : Value()};
auto alloc = mg::dyn::allocFor(rewriter, loc, memrefType, extents);
Value zero = rewriter.create<arith::ConstantOp>(loc, rewriter.getF64FloatAttr(0.0));
mg::dyn::nest2D(rewriter, loc, outType, extents, [&](OpBuilder &b, Location l, ValueRange ivs) {
b.create<affine::AffineStoreOp>(l, zero, alloc, ivs);
});
auto bound = [&](int64_t staticSize, Value lhsOrRhs, int64_t dim, Value known) -> Value {
if (!ShapedType::isDynamic(staticSize))
return rewriter.create<arith::ConstantIndexOp>(loc, staticSize);
return known ? known : mg::dyn::extentOf(rewriter, loc, lhsOrRhs, dim);
};
Value m = bound(outType.getDimSize(0), lhs, 0, Value());
Value n = bound(outType.getDimSize(1), rhs, 1, Value());
int64_t kStatic = !lt.isDynamicDim(1) ? lt.getDimSize(1) : rt.getDimSize(0);
Value k = bound(kStatic, lhs, 1, Value());
Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
affine::buildAffineLoopNest(rewriter, loc, ValueRange{c0, c0, c0}, ValueRange{m, n, k}, {1, 1, 1},
[&](OpBuilder &b, Location l, ValueRange ivs) {
Value i = ivs[0], j = ivs[1], kk = ivs[2];
Value acc = b.create<affine::AffineLoadOp>(l, alloc, ValueRange{i, j});
Value x = b.create<affine::AffineLoadOp>(l, lhs, ValueRange{i, kk});
Value y = b.create<affine::AffineLoadOp>(l, rhs, ValueRange{kk, j});
Value prod = b.create<arith::MulFOp>(l, x, y);
b.create<affine::AffineStoreOp>(l, b.create<arith::AddFOp>(l, acc, prod), alloc, ValueRange{i, j});
});
rewriter.replaceOp(op, alloc.getResult());
return success();
}
};
This is the lowered code for a static 2x3 @ 3x4. Read it in two halves: the first loop nest zero-fills the 2x4 result; the second is the three-deep accumulate. The k loop (%arg4) is the reduction:
module {
func.func @mm(%arg0: memref<2x3xf64>, %arg1: memref<3x4xf64>) -> memref<2x4xf64> {
%alloc = memref.alloc() : memref<2x4xf64>
%cst = arith.constant 0.000000e+00 : f64
affine.for %arg2 = 0 to 2 {
affine.for %arg3 = 0 to 4 {
affine.store %cst, %alloc[%arg2, %arg3] : memref<2x4xf64>
}
}
%c2 = arith.constant 2 : index
%c4 = arith.constant 4 : index
%c3 = arith.constant 3 : index
%c0 = arith.constant 0 : index
affine.for %arg2 = 0 to 2 {
affine.for %arg3 = 0 to 4 {
affine.for %arg4 = 0 to 3 {
%0 = affine.load %alloc[%arg2, %arg3] : memref<2x4xf64>
%1 = affine.load %arg0[%arg2, %arg4] : memref<2x3xf64>
%2 = affine.load %arg1[%arg4, %arg3] : memref<3x4xf64>
%3 = arith.mulf %1, %2 : f64
%4 = arith.addf %0, %3 : f64
affine.store %4, %alloc[%arg2, %arg3] : memref<2x4xf64>
}
}
}
return %alloc : memref<2x4xf64>
}
}
(Four arith.constant operations in the middle are unused leftovers of building the loop bounds; they are harmless and a later canonicalization pass would remove them.)
For a dynamic matmul the lowered code begins with the check. These are the first lines, showing the two memref.dim reads, the comparison, and the cf.assert with its message:
%c1 = arith.constant 1 : index
%dim = memref.dim %arg0, %c1 : memref<?x?xf64>
%c0 = arith.constant 0 : index
%dim_0 = memref.dim %arg1, %c0 : memref<?x?xf64>
%0 = arith.cmpi eq, %dim, %dim_0 : index
cf.assert %0, "mg.matmul: inner dimensions differ at runtime"
%c0_1 = arith.constant 0 : index
%dim_2 = memref.dim %arg0, %c0_1 : memref<?x?xf64>
The complete change against the earlier versions of these files (the operations, verifiers, lowerings, front end, driver and the PTX decoder):
The full diff for this chapter (click to expand)
--- ../../part3/code/MgOps.td 2026-10-01 21:21:01.991785523 +0000
+++ MgOps.td 2026-10-02 03:06:25.081013492 +0000
@@ -45,6 +45,59 @@
let hasFolder = 1;
}
+def SubOp : Mg_Op<"sub", [Pure]> {
+ let summary = "elementwise subtract of two compatible-shaped tensors (lhs - rhs)";
+ let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);
+ let results = (outs F64Tensor:$result);
+ let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)";
+ let hasVerifier = 1;
+}
+
+def MulOp : Mg_Op<"mul", [Pure]> {
+ let summary = "elementwise (Hadamard) product of two compatible-shaped tensors";
+ let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);
+ let results = (outs F64Tensor:$result);
+ let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)";
+ let hasVerifier = 1;
+}
+
+def DivOp : Mg_Op<"div", [Pure]> {
+ let summary = "elementwise division of two compatible-shaped tensors (lhs / rhs, IEEE: x/0 is inf or nan)";
+ let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);
+ let results = (outs F64Tensor:$result);
+ let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)";
+ let hasVerifier = 1;
+}
+
+def ScalarOp : Mg_Op<"scalar", [Pure]> {
+ let summary = "combine every element with one scalar constant";
+ let description = [{
+ `op` is one of "add", "sub", "mul", "div". With `reversed = false` each element becomes
+ `x OP value`; with `reversed = true` it becomes `value OP x` (so `10 - x` and `1 / x` exist).
+ Works for any shape, including dynamic ones, because no second tensor is involved.
+ }];
+ let arguments = (ins F64Tensor:$input, StrAttr:$op, F64Attr:$value, BoolAttr:$reversed);
+ let results = (outs F64Tensor:$result);
+ let assemblyFormat = "$input attr-dict `:` type($input) `->` type($result)";
+ let hasVerifier = 1;
+}
+
+def NegOp : Mg_Op<"neg", [Pure]> {
+ let summary = "elementwise negation";
+ let arguments = (ins F64Tensor:$input);
+ let results = (outs F64Tensor:$result);
+ let assemblyFormat = "$input attr-dict `:` type($input) `->` type($result)";
+ let hasVerifier = 1;
+}
+
+def MatmulOp : Mg_Op<"matmul", [Pure]> {
+ let summary = "matrix product: (m x k) times (k x n) gives (m x n)";
+ let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);
+ let results = (outs F64Tensor:$result);
+ let assemblyFormat = "$lhs `,` $rhs attr-dict `:` type($lhs) `,` type($rhs) `->` type($result)";
+ let hasVerifier = 1;
+}
+
def TransposeOp : Mg_Op<"transpose", [Pure]> {
let summary = "transpose of a 2-D tensor";
let arguments = (ins F64Tensor:$input);
--- ../../part13/code/MgDialect.cpp 2026-10-01 22:53:13.160361103 +0000
+++ MgDialect.cpp 2026-10-02 03:06:25.081495534 +0000
@@ -75,6 +75,52 @@
return mlir::success();
}
+namespace {
+/// Shared by mg.add-style ops: same rank, every dimension compatible (equal, or `?` on either side).
+mlir::LogicalResult verifyElementwise(mlir::Operation *op, llvm::ArrayRef<mlir::Type> types, llvm::StringRef name) {
+ llvm::SmallVector<mlir::RankedTensorType, 3> t;
+ for (mlir::Type ty : types) {
+ auto r = llvm::dyn_cast<mlir::RankedTensorType>(ty);
+ if (!r) return op->emitOpError(name + " only supports ranked tensors");
+ t.push_back(r);
+ }
+ for (auto x : t)
+ if (x.getRank() != t[0].getRank())
+ return op->emitOpError(name + " operands and result must have the same rank");
+ for (int64_t d = 0; d < t[0].getRank(); ++d)
+ for (auto x : t)
+ for (auto y : t)
+ if (!mg::dyn::compatibleDim(x.getDimSize(d), y.getDimSize(d)))
+ return op->emitOpError(name + " operands must have compatible shapes, got ") << x << " and " << y;
+ return mlir::success();
+}
+} // namespace
+
+mlir::LogicalResult SubOp::verify() { return verifyElementwise(*this, {getLhs().getType(), getRhs().getType(), getResult().getType()}, "mg.sub"); }
+mlir::LogicalResult MulOp::verify() { return verifyElementwise(*this, {getLhs().getType(), getRhs().getType(), getResult().getType()}, "mg.mul"); }
+mlir::LogicalResult DivOp::verify() { return verifyElementwise(*this, {getLhs().getType(), getRhs().getType(), getResult().getType()}, "mg.div"); }
+mlir::LogicalResult ScalarOp::verify() {
+ llvm::StringRef k = getOp();
+ if (k != "add" && k != "sub" && k != "mul" && k != "div")
+ return emitOpError("mg.scalar: op must be one of add, sub, mul, div, got '") << k << "'";
+ return verifyElementwise(*this, {getInput().getType(), getResult().getType()}, "mg.scalar");
+}
+mlir::LogicalResult NegOp::verify() { return verifyElementwise(*this, {getInput().getType(), getResult().getType()}, "mg.neg"); }
+
+mlir::LogicalResult MatmulOp::verify() {
+ auto l = llvm::dyn_cast<mlir::RankedTensorType>(getLhs().getType());
+ auto r = llvm::dyn_cast<mlir::RankedTensorType>(getRhs().getType());
+ auto o = llvm::dyn_cast<mlir::RankedTensorType>(getResult().getType());
+ if (!l || !r || !o) return emitOpError("mg.matmul only supports ranked tensors");
+ if (l.getRank() != 2 || r.getRank() != 2 || o.getRank() != 2)
+ return emitOpError("mg.matmul only supports rank-2 tensors");
+ if (!mg::dyn::compatibleDim(l.getDimSize(1), r.getDimSize(0)))
+ return emitOpError("mg.matmul inner dimensions differ: ") << l << " times " << r;
+ if (!mg::dyn::compatibleDim(o.getDimSize(0), l.getDimSize(0)) || !mg::dyn::compatibleDim(o.getDimSize(1), r.getDimSize(1)))
+ return emitOpError("mg.matmul result shape must be (rows of lhs) x (columns of rhs), got ") << o;
+ return mlir::success();
+}
+
mlir::OpFoldResult AddOp::fold(FoldAdaptor adaptor) {
auto lhsAttr = llvm::dyn_cast_if_present<mlir::DenseElementsAttr>(adaptor.getLhs());
auto rhsAttr = llvm::dyn_cast_if_present<mlir::DenseElementsAttr>(adaptor.getRhs());
--- ../../part20/code/LowerToAffine.cpp 2026-10-02 02:52:35.252838267 +0000
+++ LowerToAffine.cpp 2026-10-02 03:12:27.030941608 +0000
@@ -110,6 +110,139 @@
}
};
+/// Chapter 21: sub, mul and div share add's whole lowering; only the arithmetic op differs.
+template <typename MgOp, typename ArithOp>
+struct ElementwiseBinaryLowering : public OpConversionPattern<MgOp> {
+ using OpConversionPattern<MgOp>::OpConversionPattern;
+ using OpAdaptor = typename OpConversionPattern<MgOp>::OpAdaptor;
+
+ LogicalResult matchAndRewrite(MgOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto tensorType = llvm::cast<RankedTensorType>(op.getResult().getType());
+ auto memrefType = convertTensorToMemref(tensorType);
+ Value lhs = adaptor.getLhs(), rhs = adaptor.getRhs();
+ mg::dyn::assertSameShape(rewriter, loc, lhs, rhs,
+ llvm::cast<RankedTensorType>(op.getLhs().getType()),
+ llvm::cast<RankedTensorType>(op.getRhs().getType()),
+ op.getOperationName());
+ llvm::SmallVector<Value, 2> extents;
+ for (int64_t d = 0; d < 2; ++d)
+ extents.push_back(tensorType.isDynamicDim(d) ? mg::dyn::extentOf(rewriter, loc, lhs, d) : Value());
+ auto alloc = mg::dyn::allocFor(rewriter, loc, memrefType, extents);
+ mg::dyn::nest2D(rewriter, loc, tensorType, extents,
+ [&](OpBuilder &b, Location l, ValueRange ivs) {
+ Value x = b.create<affine::AffineLoadOp>(l, lhs, ivs);
+ Value y = b.create<affine::AffineLoadOp>(l, rhs, ivs);
+ b.create<affine::AffineStoreOp>(l, b.create<ArithOp>(l, x, y), alloc, ivs);
+ });
+ rewriter.replaceOp(op, alloc.getResult());
+ return success();
+ }
+};
+
+struct ScalarOpLowering : public OpConversionPattern<mg::ScalarOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(mg::ScalarOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto tensorType = llvm::cast<RankedTensorType>(op.getResult().getType());
+ auto memrefType = convertTensorToMemref(tensorType);
+ Value in = adaptor.getInput();
+ llvm::SmallVector<Value, 2> extents;
+ for (int64_t d = 0; d < 2; ++d)
+ extents.push_back(tensorType.isDynamicDim(d) ? mg::dyn::extentOf(rewriter, loc, in, d) : Value());
+ auto alloc = mg::dyn::allocFor(rewriter, loc, memrefType, extents);
+ Value c = rewriter.create<arith::ConstantOp>(loc, rewriter.getF64FloatAttr(op.getValue().convertToDouble()));
+ llvm::StringRef kind = op.getOp();
+ bool rev = op.getReversed();
+ mg::dyn::nest2D(rewriter, loc, tensorType, extents, [&](OpBuilder &b, Location l, ValueRange ivs) {
+ Value x = b.create<affine::AffineLoadOp>(l, in, ivs);
+ Value lhs = rev ? c : x, rhs = rev ? x : c, r;
+ if (kind == "add") r = b.create<arith::AddFOp>(l, lhs, rhs);
+ else if (kind == "sub") r = b.create<arith::SubFOp>(l, lhs, rhs);
+ else if (kind == "mul") r = b.create<arith::MulFOp>(l, lhs, rhs);
+ else r = b.create<arith::DivFOp>(l, lhs, rhs);
+ b.create<affine::AffineStoreOp>(l, r, alloc, ivs);
+ });
+ rewriter.replaceOp(op, alloc.getResult());
+ return success();
+ }
+};
+
+struct NegOpLowering : public OpConversionPattern<mg::NegOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(mg::NegOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto tensorType = llvm::cast<RankedTensorType>(op.getResult().getType());
+ auto memrefType = convertTensorToMemref(tensorType);
+ Value in = adaptor.getInput();
+ llvm::SmallVector<Value, 2> extents;
+ for (int64_t d = 0; d < 2; ++d)
+ extents.push_back(tensorType.isDynamicDim(d) ? mg::dyn::extentOf(rewriter, loc, in, d) : Value());
+ auto alloc = mg::dyn::allocFor(rewriter, loc, memrefType, extents);
+ mg::dyn::nest2D(rewriter, loc, tensorType, extents,
+ [&](OpBuilder &b, Location l, ValueRange ivs) {
+ Value x = b.create<affine::AffineLoadOp>(l, in, ivs);
+ b.create<affine::AffineStoreOp>(l, b.create<arith::NegFOp>(l, x), alloc, ivs);
+ });
+ rewriter.replaceOp(op, alloc.getResult());
+ return success();
+ }
+};
+
+/// mg.matmul: result[i,j] = sum over k of lhs[i,k] * rhs[k,j]. A zero-fill nest, then a 3-deep
+/// accumulate nest. For `?` dimensions the inner extents must agree at run time: one cf.assert.
+struct MatmulOpLowering : public OpConversionPattern<mg::MatmulOp> {
+ using OpConversionPattern::OpConversionPattern;
+ LogicalResult matchAndRewrite(mg::MatmulOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ auto loc = op.getLoc();
+ auto outType = llvm::cast<RankedTensorType>(op.getResult().getType());
+ auto lt = llvm::cast<RankedTensorType>(op.getLhs().getType());
+ auto rt = llvm::cast<RankedTensorType>(op.getRhs().getType());
+ auto memrefType = convertTensorToMemref(outType);
+ Value lhs = adaptor.getLhs(), rhs = adaptor.getRhs();
+
+ if (lt.isDynamicDim(1) || rt.isDynamicDim(0)) {
+ Value l = mg::dyn::extentOf(rewriter, loc, lhs, 1), r = mg::dyn::extentOf(rewriter, loc, rhs, 0);
+ Value eq = rewriter.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, l, r);
+ rewriter.create<cf::AssertOp>(loc, eq, rewriter.getStringAttr("mg.matmul: inner dimensions differ at runtime"));
+ }
+ // m = rows of lhs, n = columns of rhs, k = columns of lhs.
+ llvm::SmallVector<Value, 2> extents = {outType.isDynamicDim(0) ? mg::dyn::extentOf(rewriter, loc, lhs, 0) : Value(),
+ outType.isDynamicDim(1) ? mg::dyn::extentOf(rewriter, loc, rhs, 1) : Value()};
+ auto alloc = mg::dyn::allocFor(rewriter, loc, memrefType, extents);
+ Value zero = rewriter.create<arith::ConstantOp>(loc, rewriter.getF64FloatAttr(0.0));
+ mg::dyn::nest2D(rewriter, loc, outType, extents, [&](OpBuilder &b, Location l, ValueRange ivs) {
+ b.create<affine::AffineStoreOp>(l, zero, alloc, ivs);
+ });
+
+ auto bound = [&](int64_t staticSize, Value lhsOrRhs, int64_t dim, Value known) -> Value {
+ if (!ShapedType::isDynamic(staticSize))
+ return rewriter.create<arith::ConstantIndexOp>(loc, staticSize);
+ return known ? known : mg::dyn::extentOf(rewriter, loc, lhsOrRhs, dim);
+ };
+ Value m = bound(outType.getDimSize(0), lhs, 0, Value());
+ Value n = bound(outType.getDimSize(1), rhs, 1, Value());
+ int64_t kStatic = !lt.isDynamicDim(1) ? lt.getDimSize(1) : rt.getDimSize(0);
+ Value k = bound(kStatic, lhs, 1, Value());
+ Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+ affine::buildAffineLoopNest(rewriter, loc, ValueRange{c0, c0, c0}, ValueRange{m, n, k}, {1, 1, 1},
+ [&](OpBuilder &b, Location l, ValueRange ivs) {
+ Value i = ivs[0], j = ivs[1], kk = ivs[2];
+ Value acc = b.create<affine::AffineLoadOp>(l, alloc, ValueRange{i, j});
+ Value x = b.create<affine::AffineLoadOp>(l, lhs, ValueRange{i, kk});
+ Value y = b.create<affine::AffineLoadOp>(l, rhs, ValueRange{kk, j});
+ Value prod = b.create<arith::MulFOp>(l, x, y);
+ b.create<affine::AffineStoreOp>(l, b.create<arith::AddFOp>(l, acc, prod), alloc, ValueRange{i, j});
+ });
+ rewriter.replaceOp(op, alloc.getResult());
+ return success();
+ }
+};
+
/// mg.transpose lowers to a real `affine.for` loop nest over the *output*
/// shape, loading `input[j, i]` and storing to `result[i, j]` -- the
/// transpose is the real index swap between load and store, nothing more.
@@ -263,7 +396,11 @@
RewritePatternSet patterns(&getContext());
patterns.add<ConstantOpLowering, AddOpLowering, TransposeOpLowering,
- PrintOpLowering, TensorCastLowering>(typeConverter, &getContext());
+ PrintOpLowering, TensorCastLowering,
+ ElementwiseBinaryLowering<mg::SubOp, arith::SubFOp>,
+ ElementwiseBinaryLowering<mg::MulOp, arith::MulFOp>,
+ ElementwiseBinaryLowering<mg::DivOp, arith::DivFOp>,
+ ScalarOpLowering, NegOpLowering, MatmulOpLowering>(typeConverter, &getContext());
populateFunctionOpInterfaceTypeConversionPattern<func::FuncOp>(patterns,
typeConverter);
populateCallOpTypeConversionPattern(patterns, typeConverter);
--- ../../part20/code/mgfront.py 2026-10-02 02:58:36.640920215 +0000
+++ mgfront.py 2026-10-02 03:15:00.006673338 +0000
@@ -1,13 +1,16 @@
#!/usr/bin/env python3
"""mgfront: the Mountain Goat surface-syntax front end. Reads .mg source, writes MLIR (mg dialect) to stdout.
+Operators (Chapter 21): + - * / are elementwise on matrices ('*' is the Hadamard product) and also combine a matrix
+with a scalar; '@' is the matrix product; unary '-' negates. Scalars are numbers known at compile time.
Grammar (one statement per line, '#' starts a comment):
program := (def | stmt)*
def := 'def' NAME '(' [NAME ':' type (',' NAME ':' type)*] ')' '=' expr -- a function of tensors
stmt := 'let' NAME '=' expr | 'print' expr
type := 'tensor' '[' dim 'x' dim ']' dim := INT | '?'
- expr := term ('+' term)*
- term := NAME | matrix | 'transpose' '(' expr ')' | NAME '(' expr (',' expr)* ')' | '(' expr ')'
+ expr := product (('+'|'-') product)*
+ product := unary (('*'|'/'|'@') unary)* unary := '-' unary | term
+ term := NAME | NUMBER | matrix | 'transpose' '(' expr ')' | NAME '(' expr (',' expr)* ')' | '(' expr ')'
matrix := '[' row (',' row)* ']' row := '[' NUM (',' NUM)* ']'
Every value is a rank-2 f64 tensor. A dimension is an int or None (dynamic, written '?').
"""
@@ -16,7 +19,7 @@
class MgError(Exception):
def __init__(self, line, msg): self.line, self.msg = line, msg
-TOKEN = re.compile(r"\s*(?:(?P<num>-?\d+\.?\d*(?:[eE][-+]?\d+)?)|(?P<id>[A-Za-z_]\w*)|(?P<sym>[()\[\],+=:?]))")
+TOKEN = re.compile(r"\s*(?:(?P<num>\d+\.?\d*(?:[eE][-+]?\d+)?)|(?P<id>[A-Za-z_]\w*)|(?P<sym>[()\[\],+\-*/@=:?]))")
def tokenize(text, line):
toks, pos = [], 0
@@ -30,6 +33,13 @@
toks.append((kind, m.group(kind)))
return toks
+def mlir_float(text):
+ """MLIR's float literals need a decimal point before any exponent: '3e2' is rejected, '3.0e2' is accepted."""
+ t = repr(float(text)) # '300.0', '0.25', '1e+16', '1.5e-07'
+ if "e" in t and "." not in t:
+ mantissa, exp = t.split("e"); t = mantissa + ".0e" + exp
+ return t
+
def fmt(shape): # (2,None) -> 'tensor<2x?xf64>'
return "tensor<" + "x".join("?" if d is None else str(d) for d in shape) + "xf64>"
@@ -56,23 +66,56 @@
if want is not None and v != want:
raise MgError(self.line, f"expected {want!r} but found {v!r}")
self.i += 1; return k, v
- def expr(self):
- v, s = self.term()
- while self.peek()[1] == "+":
- self.take("+"); w, s2 = self.term()
- for a, b in zip(s, s2):
- if a is not None and b is not None and a != b:
- raise MgError(self.line, f"cannot add shapes {'x'.join(map(str,s))} and {'x'.join(map(str,s2))}")
- shape = tuple(a if a is not None else b for a, b in zip(s, s2))
+ # A value is (ssa_name, shape) for a matrix, or (None, float) for a compile-time scalar.
+ OPNAME = {"+": "add", "-": "sub", "*": "mul", "/": "div"}
+ VERB = {"+": "add", "-": "subtract", "*": "multiply", "/": "divide", "@": "multiply (matmul)"}
+ def expr(self): # additive level: + and -
+ x = self.product()
+ while self.peek()[1] in ("+", "-"):
+ op = self.take()[1]; y = self.product(); x = self.binary(op, x, y)
+ return x
+ def product(self): # multiplicative level: * (elementwise), / (elementwise), @ (matmul)
+ x = self.unary()
+ while self.peek()[1] in ("*", "/", "@"):
+ op = self.take()[1]; y = self.unary(); x = self.binary(op, x, y)
+ return x
+ def unary(self):
+ if self.peek()[1] == "-":
+ self.take("-"); v, s = self.unary()
+ if v is None: return None, -s
+ r = self.fn.tmp(); self.fn.emit(f"{r} = mg.neg {v} : {fmt(s)} -> {fmt(s)}"); return r, s
+ return self.term()
+ def binary(self, op, x, y):
+ (v, s), (w, s2) = x, y
+ if op == "@":
+ if v is None or w is None: raise MgError(self.line, "'@' (matrix product) needs two matrices, not a scalar")
+ if s[1] is not None and s2[0] is not None and s[1] != s2[0]:
+ raise MgError(self.line, f"cannot multiply (matmul) shapes {s[0]}x{s[1]} and {s2[0]}x{s2[1]}: inner dimensions differ")
+ shape = (s[0], s2[1]); r = self.fn.tmp()
+ self.fn.emit(f"{r} = mg.matmul {v}, {w} : {fmt(s)}, {fmt(s2)} -> {fmt(shape)}")
+ return r, shape
+ if v is None and w is None: # scalar op scalar: folded here
+ if op == "/" and s2 == 0: raise MgError(self.line, "division of a scalar by zero")
+ return None, {"+": s + s2, "-": s - s2, "*": s * s2, "/": (s / s2 if op == "/" else 0)}[op]
+ if v is None or w is None: # matrix with a scalar: mg.scalar, reversed when the scalar is on the left
+ rev = v is None
+ m, shp, c = (w, s2, s) if rev else (v, s, s2)
r = self.fn.tmp()
- self.fn.emit(f"{r} = mg.add {v}, {w} : {fmt(s)}, {fmt(s2)} -> {fmt(shape)}")
- v, s = r, shape
- return v, s
+ self.fn.emit(f'{r} = mg.scalar {m} {{op = "{self.OPNAME[op]}", value = {mlir_float(str(c))} : f64, reversed = {"true" if rev else "false"}}} : {fmt(shp)} -> {fmt(shp)}')
+ return r, shp
+ for a, b in zip(s, s2):
+ if a is not None and b is not None and a != b:
+ raise MgError(self.line, f"cannot {self.VERB[op]} shapes {'x'.join(map(str,s))} and {'x'.join(map(str,s2))}")
+ shape = tuple(a if a is not None else b for a, b in zip(s, s2)); r = self.fn.tmp()
+ self.fn.emit(f"{r} = mg.{self.OPNAME[op]} {v}, {w} : {fmt(s)}, {fmt(s2)} -> {fmt(shape)}")
+ return r, shape
def term(self):
k, v = self.peek()
+ if k is None: raise MgError(self.line, "expected an expression, found the end of the line")
if v == "(":
self.take("("); r = self.expr(); self.take(")"); return r
if v == "[": return self.matrix()
+ if k == "num": self.take(); return None, float(v)
if k == "id":
self.take()
if v == "transpose":
@@ -95,6 +138,7 @@
raise MgError(self.line, f"{name} takes {len(params)} arguments, got {len(args)}")
vals = []
for (v, s), p in zip(args, params):
+ if v is None: raise MgError(self.line, f"{name}: a scalar cannot be passed where a matrix is expected")
for a, b in zip(s, p):
if b is not None and a != b:
raise MgError(self.line, f"argument shape {'x'.join('?' if d is None else str(d) for d in s)} does not fit parameter {'x'.join('?' if d is None else str(d) for d in p)}")
@@ -105,12 +149,16 @@
sig = ", ".join(fmt(p) for p in params)
self.fn.emit(f"{r} = func.call @{name}({', '.join(vals)}) : ({sig}) -> {fmt(result)}")
return r, result
+ def entry(self):
+ neg = self.peek()[1] == "-"
+ if neg: self.take("-")
+ return ("-" if neg else "") + self.take()[1]
def matrix(self):
rows = []
self.take("[")
while True:
- self.take("["); row = [self.take()[1]]
- while self.peek()[1] == ",": self.take(","); row.append(self.take()[1])
+ self.take("["); row = [self.entry()]
+ while self.peek()[1] == ",": self.take(","); row.append(self.entry())
self.take("]"); rows.append(row)
if self.peek()[1] == ",": self.take(","); continue
break
@@ -118,7 +166,7 @@
if len({len(r) for r in rows}) != 1: raise MgError(self.line, "matrix rows have different lengths")
def f(x):
if not re.fullmatch(r"-?\d+\.?\d*(?:[eE][-+]?\d+)?", x): raise MgError(self.line, f"{x!r} is not a number")
- return x if ("." in x or "e" in x.lower()) else x + ".0"
+ return mlir_float(x)
body = "[" + ", ".join("[" + ", ".join(f(x) for x in r) + "]" for r in rows) + "]"
shape = (len(rows), len(rows[0])); r = self.fn.tmp()
self.fn.emit(f"{r} = mg.constant dense<{body}> : {fmt(shape)}")
@@ -131,8 +179,8 @@
for i in range(2):
v = p.take()[1]
if v == "?": dims.append(None)
- elif v.isdigit(): dims.append(int(v))
- else: raise MgError(p.line, f"bad dimension {v!r}")
+ elif v.isdigit() and int(v) > 0: dims.append(int(v))
+ else: raise MgError(p.line, f"bad dimension {v!r}: use a positive integer or '?'")
if i == 0 and p.take()[1] != "x": raise MgError(p.line, "expected 'x' between dimensions")
p.take("]"); return tuple(dims)
@@ -156,6 +204,7 @@
p.take(")"); p.take("=")
for i, (n, shp) in enumerate(zip(names, params)): fn.env[n] = (f"%{n}", shp)
v, s = p.expr()
+ if v is None: raise MgError(ln, f"def {name}: a function must return a matrix, not a scalar")
if p.i != len(toks): raise MgError(ln, f"unexpected {p.peek()[1]!r}")
em.funcs[name] = (params, s)
em.defs.append((name, names, params, s))
@@ -166,7 +215,9 @@
if kw == "let":
n = p.take()[1]; p.take("="); v, s = p.expr(); main.env[n] = (v, s)
elif kw == "print":
- v, s = p.expr(); main.emit(f"mg.print {v} : {fmt(s)}")
+ v, s = p.expr()
+ if v is None: raise MgError(ln, "print needs a matrix; a scalar has no shape to print")
+ main.emit(f"mg.print {v} : {fmt(s)}")
else:
raise MgError(ln, f"a statement must start with 'let', 'print' or 'def', not {kw!r}")
if p.i != len(toks): raise MgError(ln, f"unexpected {p.peek()[1]!r}")
--- ../../part20/code/mgc 2026-10-02 02:58:36.642607976 +0000
+++ mgc 2026-10-02 03:11:46.808300801 +0000
@@ -5,10 +5,10 @@
# mgc lib prog.mg [-o dir] compile the 'def's to prog.o plus a C++ header prog.h (no main), for linking into your own program
# mgc mlir prog.mg print the mg-dialect MLIR the front end produces
# mgc ptx prog.mg compile to PTX (no launch: this machine has no GPU) and print the PTX
-# Env: MG_OPT (default Chapter 20 build), MGC_KEEP=dir (keep every intermediate file there).
+# Env: MG_OPT (default Chapter 21 build), MGC_KEEP=dir (keep every intermediate file there).
set -e
HERE=$(cd "$(dirname "$0")" && pwd)
-MG_OPT=${MG_OPT:-$HERE/../../part20/code/build/mg-opt}
+MG_OPT=${MG_OPT:-$HERE/../../part21/code/build/mg-opt}
LIBDIR=${MLIR_LIBDIR:-/usr/lib/llvm-18/lib}
usage() { echo "usage: mgc run|build|lib|mlir|ptx prog.mg [-o exe]" >&2; exit 2; }
[ $# -ge 2 ] || usage
@@ -20,7 +20,7 @@
LIBFLAG=; [ "$cmd" = lib ] && LIBFLAG=--lib
python3 "$HERE/mgfront.py" $LIBFLAG "$src" > "$W/$base.mlir" # 1. surface syntax -> mg dialect (exit 1 + message on a syntax/shape error)
[ "$cmd" = mlir ] && { cat "$W/$base.mlir"; exit 0; }
-[ -x "$MG_OPT" ] || { echo "mgc: no mg-opt at $MG_OPT (run part20/code/build.sh)" >&2; exit 2; }
+[ -x "$MG_OPT" ] || { echo "mgc: no mg-opt at $MG_OPT (run part21/code/build.sh)" >&2; exit 2; }
case $cmd in
build|run|lib)
LOWER='--lower-affine --convert-scf-to-cf --mg-lower-assert-to-stderr --convert-arith-to-llvm --finalize-memref-to-llvm --convert-func-to-llvm --reconcile-unrealized-casts'
@@ -42,8 +42,8 @@
mlir-opt-18 "$W/$base.affine.mlir" \
--pass-pipeline='builtin.module(func.func(affine-parallelize,lower-affine,gpu-map-parallel-loops,convert-parallel-loops-to-gpu),gpu-kernel-outlining)' -o "$W/$base.gpu.mlir"
mlir-opt-18 "$W/$base.gpu.mlir" \
- --pass-pipeline='builtin.module(lower-affine,nvvm-attach-target{chip=sm_70 features=+ptx60},gpu.module(convert-gpu-to-nvvm{index-bitwidth=64},reconcile-unrealized-casts),gpu-module-to-binary{format=isa})' -o "$W/$base.bin.mlir"
- python3 "$HERE/../../part10/code/decode_ptx.py" "$W/$base.bin.mlir"
+ --pass-pipeline='builtin.module(lower-affine,nvvm-attach-target{chip=sm_70 features=+ptx60},gpu.module(convert-scf-to-cf,convert-gpu-to-nvvm{index-bitwidth=64},reconcile-unrealized-casts),gpu-module-to-binary{format=isa})' -o "$W/$base.bin.mlir"
+ python3 "$HERE/decode_ptx.py" "$W/$base.bin.mlir"
;;
*) usage;;
esac
--- ../../part10/code/decode_ptx.py 2026-10-01 21:25:13.721605362 +0000
+++ decode_ptx.py 2026-10-02 03:09:33.392749115 +0000
@@ -1,7 +1,8 @@
# --gpu-module-to-binary stores the PTX as a hex-escaped string attribute (\0A for newline).
+# Chapter 21: a program with several loop nests has several gpu.modules; this prints EVERY kernel's PTX (the Chapter 10 version stopped at the first).
# This pulls it back out as plain PTX text: python3 decode_ptx.py add_bin.mlir > out.ptx
import re, sys
s = open(sys.argv[1]).read()
-m = re.search(r'assembly\s*=\s*"((?:[^"\\]|\\.)*)"', s)
-t = re.sub(r'\\([0-9A-Fa-f]{2})', lambda x: chr(int(x.group(1), 16)), m.group(1))
-sys.stdout.write(t.replace('\x00', ''))
+for m in re.finditer(r'assembly\s*=\s*"((?:[^"\\]|\\.)*)"', s):
+ t = re.sub(r'\\([0-9A-Fa-f]{2})', lambda x: chr(int(x.group(1), 16)), m.group(1))
+ sys.stdout.write(t.replace('\x00', ''))
To rebuild mg-opt with everything above:
#!/bin/sh
# Rebuild mg-opt for Chapter 21: Chapter 20's build plus one new lowering (tensor.cast -> memref.cast). Produces ./build/mg-opt.
# Needs mlir-18-tools, libmlir-18-dev, llvm-18-dev, cmake, make.
set -e
HERE=$(cd "$(dirname "$0")" && pwd); D=$HERE/../..; T=$HERE/tree
mkdir -p $T/include/mg $T/lib $T/tools
cp $HERE/CMakeLists.txt $T/
cp $D/part3/code/MgDialect.td $HERE/MgOps.td $D/part3/code/MgDialect.h $D/part3/code/MgOps.h $T/include/mg/
cp $D/part14/code/DynamicShapes.h $T/include/mg/
cp $HERE/MgDialect.cpp $T/lib/
cp $D/part14/code/MgBufferizableOpInterfaceImpl.cpp $T/lib/
cp $D/part19/code/AssertToStderr.cpp $T/lib/
cp $HERE/LowerToAffine.cpp $T/lib/
cp $HERE/mg-opt.cpp $T/tools/
mkdir -p $HERE/build/include/mg && cd $HERE/build
cmake $T -DMLIR_DIR=/usr/lib/llvm-18/lib/cmake/mlir -DLLVM_DIR=/usr/lib/llvm-18/lib/cmake/llvm \
-DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=OFF > cmake_out.txt 2>&1
make -j8 2>&1 | tail -2
The front end: new operators¶
The Chapter 20 front end parsed + and transpose only. Operator parsing now has the three levels from the design section, one small method each:
def take(self, want=None):
k, v = self.peek()
if k is None: raise MgError(self.line, "unexpected end of line")
if want is not None and v != want:
raise MgError(self.line, f"expected {want!r} but found {v!r}")
self.i += 1; return k, v
# A value is (ssa_name, shape) for a matrix, or (None, float) for a compile-time scalar.
OPNAME = {"+": "add", "-": "sub", "*": "mul", "/": "div"}
VERB = {"+": "add", "-": "subtract", "*": "multiply", "/": "divide", "@": "multiply (matmul)"}
def expr(self): # additive level: + and -
x = self.product()
while self.peek()[1] in ("+", "-"):
op = self.take()[1]; y = self.product(); x = self.binary(op, x, y)
return x
def product(self): # multiplicative level: * (elementwise), / (elementwise), @ (matmul)
x = self.unary()
while self.peek()[1] in ("*", "/", "@"):
op = self.take()[1]; y = self.unary(); x = self.binary(op, x, y)
return x
def unary(self):
if self.peek()[1] == "-":
self.take("-"); v, s = self.unary()
if v is None: return None, -s
r = self.fn.tmp(); self.fn.emit(f"{r} = mg.neg {v} : {fmt(s)} -> {fmt(s)}"); return r, s
return self.term()
def binary(self, op, x, y):
(v, s), (w, s2) = x, y
if op == "@":
if v is None or w is None: raise MgError(self.line, "'@' (matrix product) needs two matrices, not a scalar")
if s[1] is not None and s2[0] is not None and s[1] != s2[0]:
raise MgError(self.line, f"cannot multiply (matmul) shapes {s[0]}x{s[1]} and {s2[0]}x{s2[1]}: inner dimensions differ")
shape = (s[0], s2[1]); r = self.fn.tmp()
self.fn.emit(f"{r} = mg.matmul {v}, {w} : {fmt(s)}, {fmt(s2)} -> {fmt(shape)}")
return r, shape
if v is None and w is None: # scalar op scalar: folded here
if op == "/" and s2 == 0: raise MgError(self.line, "division of a scalar by zero")
return None, {"+": s + s2, "-": s - s2, "*": s * s2, "/": (s / s2 if op == "/" else 0)}[op]
if v is None or w is None: # matrix with a scalar: mg.scalar, reversed when the scalar is on the left
rev = v is None
m, shp, c = (w, s2, s) if rev else (v, s, s2)
r = self.fn.tmp()
self.fn.emit(f'{r} = mg.scalar {m} {{op = "{self.OPNAME[op]}", value = {mlir_float(str(c))} : f64, reversed = {"true" if rev else "false"}}} : {fmt(shp)} -> {fmt(shp)}')
return r, shp
for a, b in zip(s, s2):
if a is not None and b is not None and a != b:
raise MgError(self.line, f"cannot {self.VERB[op]} shapes {'x'.join(map(str,s))} and {'x'.join(map(str,s2))}")
shape = tuple(a if a is not None else b for a, b in zip(s, s2)); r = self.fn.tmp()
self.fn.emit(f"{r} = mg.{self.OPNAME[op]} {v}, {w} : {fmt(s)}, {fmt(s2)} -> {fmt(shape)}")
binary is where each operator becomes an operation. Two matrices with + - * / become the elementwise operation (shapes checked, ? accepted); @ becomes mg.matmul with the result shape built from the two inputs; a matrix and a number become mg.scalar (reversed when the number is on the left); and two numbers are folded on the spot (dividing a scalar by a scalar zero is an error there, since the front end can see it; dividing a matrix by zero is not, because that is IEEE). A scalar passed where a matrix is expected, printed, or returned from a function is rejected with a message.
Fourteen examples¶
Every example lives in docs/part21/code/examples/. Each block below shows the program, the real output, and the exit status. Memref base addresses (0x…) vary and are masked. As in Chapter 20, the output format (Unranked Memref ... data = [[...]]) is MLIR's runtime printer, not Mountain Goat's.
1. Subtraction: order matters¶
################ ./mgc run examples/01_subtract.mg
| # Elementwise subtraction. Order matters: a - b is not b - a.
| let a = [[10, 20, 30], [40, 50, 60]]
| let b = [[1, 2, 3], [4, 5, 6]]
| print a - b
| print b - a
----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data =
[[9, 18, 27],
[36, 45, 54]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data =
[[-9, -18, -27],
[-36, -45, -54]]
----- exit status 0
a - b and b - a are negatives of each other, as arithmetic says.
2. The Hadamard product: elementwise¶
################ ./mgc run examples/02_hadamard.mg
| # '*' is the Hadamard (elementwise) product: each element times the matching element.
| # It is NOT the matrix product; that is '@' (example 4).
| let a = [[1, 2, 3], [4, 5, 6]]
| let b = [[2, 2, 2], [3, 3, 3]]
| print a * b
----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 3] strides = [3, 1] data =
[[2, 4, 6],
[12, 15, 18]]
----- exit status 0
Each element is multiplied by the matching element (3 * 2 = 6, 5 * 3 = 15). Compare with example 4, which uses the same kind of inputs but gives a different answer, because @ is a different operation.
3. Division, and what happens at zero¶
################ ./mgc run examples/03_divide.mg
| # Elementwise division, and IEEE behaviour: x/0 is inf, 0/0 is nan (no error, no abort).
| let a = [[10, 20], [0, -5]]
| let b = [[4, 8], [0, 0]]
| print a / b
----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[2.5, 2.5],
[-nan, -inf]]
----- exit status 0
10/4 = 2.5 and 20/8 = 2.5. The bottom row divides 0 by 0 and -5 by 0: the result is -nan and -inf, and the program exits normally with status 0. The sign shown on nan is whatever the hardware produced for 0/0 (x86 produces a negative quiet NaN); the book does not rely on it, and the test accepts either sign.
4. Matrix product¶
################ ./mgc run examples/04_matmul.mg
| # Matrix product: (2x3) @ (3x2) gives (2x2). Entry (i,j) = sum over k of a[i,k] * b[k,j].
| let a = [[1, 2, 3], [4, 5, 6]]
| let b = [[7, 8], [9, 10], [11, 12]]
| print a @ b
| print b @ a
----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[58, 64],
[139, 154]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [3, 3] strides = [3, 1] data =
[[39, 54, 69],
[49, 68, 87],
[59, 82, 105]]
----- exit status 0
A 2x3 matrix times a 3x2 matrix is 2x2. Check the top-left entry by hand: row [1, 2, 3] dotted with column [7, 9, 11] is 7 + 18 + 33 = 58. The second product, b @ a, is 3x3: reversing the operands changes the shape, not just the values.
5. Scalars, on either side¶
################ ./mgc run examples/05_scalar_ops.mg
| # Scalars combine with every element. The scalar can be on either side, which matters for - and /.
| let a = [[1, 2], [4, 8]]
| print a + 10
| print a - 1
| print 10 - a
| print a * 0.5
| print 16 / a
| print a / 4
----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[11, 12],
[14, 18]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[0, 1],
[3, 7]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[9, 8],
[6, 2]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[0.5, 1],
[2, 4]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[16, 8],
[4, 2]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[0.25, 0.5],
[1, 2]]
----- exit status 0
Six forms: a + 10, a - 1, 10 - a, a * 0.5, 16 / a, a / 4. Compare a - 1 ([[0, 1], [3, 7]]) with 10 - a ([[9, 8], [6, 2]]): the reversed flag is why they differ. 16 / a divides 16 by each element.
6. Precedence¶
################ ./mgc run examples/06_precedence.mg
| # Precedence: unary minus binds tightest, then * / @ (left to right), then + -.
| let a = [[1, 2], [3, 4]]
| let b = [[1, 0], [0, 1]]
| print a + a * a
| print (a + a) * a
| print a - b @ a
| print -a + 10
| let k = 2 + 3 * 4
| print a * k
----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[2, 6],
[12, 20]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[2, 8],
[18, 32]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[0, 0],
[0, 0]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[9, 8],
[7, 6]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[14, 28],
[42, 56]]
----- exit status 0
Five lines: a + a * a, (a + a) * a, a - b @ a (here b is the identity, so the answer is all zeros), -a + 10, and a * k where k = 2 + 3 * 4 = 14. The first two give different answers from the same symbols, which is what precedence means.
7. A linear layer¶
################ ./mgc run examples/07_linear_layer.mg
| # A tiny "linear layer": y = x @ w + bias, with x 2x3 (two samples), w 3x2, bias 2x2.
| def layer(x: tensor[2x3], w: tensor[3x2], bias: tensor[2x2]) = x @ w + bias
| print layer([[1, 2, 3], [4, 5, 6]], [[1, 0], [0, 1], [1, 1]], [[0.5, 0.5], [0.5, 0.5]])
----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[4.5, 5.5],
[10.5, 11.5]]
----- exit status 0
x @ w + bias is the core computation of one layer of a neural network: two samples (rows of x) are multiplied by a weight matrix and a bias is added. This is a toy: there is no learning and no activation function. The point is that matrix product and elementwise addition compose in one expression.
8. A Gram matrix¶
################ ./mgc run examples/08_gram.mg
| # a @ transpose(a) is always square and symmetric (a "Gram matrix"): entry (i,j) = row i . row j.
| let a = [[1, 2, 3], [4, 5, 6]]
| print a @ transpose(a)
| print transpose(a) @ a
----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[14, 32],
[32, 77]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [3, 3] strides = [3, 1] data =
[[17, 22, 27],
[22, 29, 36],
[27, 36, 45]]
----- exit status 0
a @ transpose(a) is always square and always symmetric (the entry at row i, column j equals the one at row j, column i): look at 32 and 32. Each entry is the dot product of two rows.
9. Identities you can check by eye¶
################ ./mgc run examples/09_identity.mg
| # Identities you can check by eye: a @ I = a, a - a = 0, a / a = 1 (where a is not 0), -(-a) = a.
| let a = [[3, 1], [4, 1]]
| let i = [[1, 0], [0, 1]]
| print a @ i
| print a - a
| print a / a
| print -(-a)
----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[3, 1],
[4, 1]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[0, 0],
[0, 0]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[1, 1],
[1, 1]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[3, 1],
[4, 1]]
----- exit status 0
a @ I = a, a - a = 0, a / a = 1 (wherever a is not zero) and -(-a) = a. When you add an operation to a compiler, checking such identities is a cheap, independent way to catch a lowering bug, because they hold regardless of what the numbers are. They have blind spots, though: -(-a) = a cannot tell a real negation from one that does nothing (two no-ops also give back a), so negation gets its own direct checks (16_negate.mg, below).
10. One compiled matmul, many shapes¶
################ ./mgc run examples/10_dynamic_matmul.mg
| # One compiled matmul for any compatible sizes. The inner dimensions are checked at run time.
| def mm(a: tensor[?x?], b: tensor[?x?]) = a @ b
| print mm([[1, 2, 3], [4, 5, 6]], [[1], [1], [1]])
| print mm([[1, 2]], [[3, 4], [5, 6]])
| print mm([[2]], [[21]])
----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 1] strides = [1, 1] data =
[[6],
[15]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 2] strides = [2, 1] data =
[[13, 16]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 1] strides = [1, 1] data =
[[42]]
----- exit status 0
mm takes tensor[?x?] arguments, so it is compiled once. It is called with 2x3 times 3x1, then 1x2 times 2x2, then 1x1 times 1x1, and the results have sizes 2x1, 1x2 and 1x1.
11. Scalar operations on dynamic shapes¶
################ ./mgc run examples/11_dynamic_scalar.mg
| # Scalar operations need no run-time shape check (there is no second matrix), so they work on any shape.
| def affine(a: tensor[?x?]) = a * 2 + 1
| print affine([[1, 2, 3]])
| print affine([[1], [2]])
| print affine([[0, 0], [0, 0], [0, 0]])
----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data =
[[3, 5, 7]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 1] strides = [1, 1] data =
[[3],
[5]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [3, 2] strides = [2, 1] data =
[[1, 1],
[1, 1],
[1, 1]]
----- exit status 0
affine computes a * 2 + 1 for any shape. No run-time shape check is needed here, because there is one matrix and a compile-time scalar.
12. When the inner dimensions disagree¶
################ ./mgc run examples/12_matmul_mismatch.mg
| # Inner dimensions 3 and 2 do not agree. The front end cannot see this (both are ?), so the program
| # compiles and the run-time check aborts with a message on stderr.
| def mm(a: tensor[?x?], b: tensor[?x?]) = a @ b
| print mm([[1, 2, 3], [4, 5, 6]], [[1, 2], [3, 4]])
----- output
mg.matmul: inner dimensions differ at runtime
Aborted
----- exit status 134
2x3 times 2x2 has inner dimensions 3 and 2. The front end cannot see this (both parameters are ?), so the program compiles and the run-time check aborts it: message on stderr, exit status 134, and nothing on stdout. That is Chapters 14 and 19 working for an operation they were not written for, because mg.matmul uses the same cf.assert that mg.add does.
13 and 14. The GPU path for the new operations¶
################ ./mgc ptx examples/13_gpu_matmul.mg
| # The GPU path for matmul (compile-only). Use: mgc ptx 13_gpu_matmul.mg
| def mm(a: tensor[4x4], b: tensor[4x4]) = a @ b
----- the PTX instructions that do the arithmetic, and the kernel entry points
.visible .entry mm_kernel(
st.global.f64 [%rd11], %fd1;
.visible .entry mm_kernel(
ld.global.f64 %fd1, [%rd15];
ld.global.f64 %fd2, [%rd17];
ld.global.f64 %fd3, [%rd16];
mul.rn.f64 %fd4, %fd2, %fd3;
add.rn.f64 %fd5, %fd1, %fd4;
st.global.f64 [%rd15], %fd5;
ld.global.f64 %fd6, [%rd17+8];
ld.global.f64 %fd7, [%rd16+32];
mul.rn.f64 %fd8, %fd6, %fd7;
add.rn.f64 %fd9, %fd5, %fd8;
st.global.f64 [%rd15], %fd9;
ld.global.f64 %fd10, [%rd17+16];
ld.global.f64 %fd11, [%rd16+64];
mul.rn.f64 %fd12, %fd10, %fd11;
add.rn.f64 %fd13, %fd9, %fd12;
st.global.f64 [%rd15], %fd13;
ld.global.f64 %fd14, [%rd17+24];
ld.global.f64 %fd15, [%rd16+96];
mul.rn.f64 %fd16, %fd14, %fd15;
add.rn.f64 %fd17, %fd13, %fd16;
st.global.f64 [%rd15], %fd17;
################ ./mgc ptx examples/14_gpu_ops.mg
| # The GPU path for the elementwise and scalar ops (compile-only). Use: mgc ptx 14_gpu_ops.mg
| def f(a: tensor[2x2], b: tensor[2x2]) = (a - b) * a / 2 + 1
----- the PTX instructions that do the arithmetic, and the kernel entry points
.visible .entry f_kernel(
ld.global.f64 %fd1, [%rd16];
ld.global.f64 %fd2, [%rd17];
sub.rn.f64 %fd3, %fd1, %fd2;
st.global.f64 [%rd18], %fd3;
.visible .entry f_kernel(
ld.global.f64 %fd1, [%rd16];
ld.global.f64 %fd2, [%rd17];
mul.rn.f64 %fd3, %fd1, %fd2;
st.global.f64 [%rd18], %fd3;
.visible .entry f_kernel(
ld.global.f64 %fd2, [%rd14];
div.rn.f64 %fd3, %fd2, %fd1;
st.global.f64 [%rd15], %fd3;
.visible .entry f_kernel(
ld.global.f64 %fd2, [%rd14];
add.rn.f64 %fd3, %fd2, %fd1;
st.global.f64 [%rd15], %fd3;
These are the instructions out of the real PTX that mgc ptx produced (compile-only: nothing was launched, and the PTX was never assembled with ptxas). Three things are worth reading:
matmulbecomes two kernels. The first, which stores a zero, is the zero-fill loop nest. The second does the accumulate. Chapter 10's recipe parallelizes the loops it can prove safe (theiandjloops, one thread per result element), and leaves thekloop, the reduction, as an ordinary sequential loop inside each thread.- The
kloop is unrolled. The fourmul.rn.f64andadd.rn.f64pairs, with load offsets+8,+16,+24(a row of the left operand) and+32,+64,+96(a column of the right one), are the four iterations of a 4x4 product, becausekis the constant 4. Each thread loads, multiplies, adds and stores the running sum back, four times. - Elementwise chains become several kernels.
(a - b) * a / 2 + 1is four loop nests, hence four kernels (sub.rn.f64,mul.rn.f64,div.rn.f64,add.rn.f64), each reading the previous one's buffer. No fusion happened (Chapter 18 explains why fusion is conservative).
Two honest corrections are needed here. First, the very first attempt on matmul failed with failed to legalize operation 'builtin.unrealized_conversion_cast' that was explicitly marked illegal. The cause: the sequential k loop lowers to an scf.for inside the GPU module, and Chapter 10's device pipeline had no pass to turn structured loops into branches before convert-gpu-to-nvvm. Adding convert-scf-to-cf inside gpu.module(...) fixed it; Chapter 10's own examples never had a loop in a kernel, which is why the gap never showed. Second, Chapter 10's decode_ptx.py printed only the first kernel's PTX, so the first look at example 14 appeared to contain one kernel when it had four (grep -c assembly on the binary showed 4). This chapter ships a corrected decode_ptx.py (it prints every kernel).
Whether the kernels compute the right answer on a GPU is not established. Nothing here was run on one. What was checked is that the CPU path (examples 1 to 12) gives correct results for the same operations, and that the GPU path produces PTX of the expected structure.
Errors¶
################ ./mgc run examples/errors/bad_matmul.mg
| # Rejected by the front end: (2x3) @ (2x3) has inner dimensions 3 and 2.
| let a = [[1, 2, 3], [4, 5, 6]]
| print a @ a
----- output
errors/bad_matmul.mg:3: error: cannot multiply (matmul) shapes 2x3 and 2x3: inner dimensions differ
----- exit status 1
################ ./mgc run examples/errors/bad_scalar.mg
| # Rejected by the front end: a scalar has no shape, so it cannot be matrix-multiplied.
| let a = [[1, 2], [3, 4]]
| let k = 3
| print a @ k
----- output
errors/bad_scalar.mg:4: error: '@' (matrix product) needs two matrices, not a scalar
----- exit status 1
Both are caught by the front end, with the line number, before any MLIR exists.
Using the new operations from C++¶
mgc lib generates a C++ header for any set of defs (Chapter 20). The Mountain Goat source for six functions:
# The Chapter 21 operations, as functions callable from C++.
def sub(a: tensor[?x?], b: tensor[?x?]) = a - b
def hadamard(a: tensor[?x?], b: tensor[?x?]) = a * b
def divide(a: tensor[?x?], b: tensor[?x?]) = a / b
def matmul(a: tensor[?x?], b: tensor[?x?]) = a @ b
def affine(a: tensor[?x?]) = a * 2 + 1
def gram(a: tensor[?x?]) = a @ transpose(a)
A C++ program calls them and, for matmul, compares the result against a plain C++ triple loop on two inputs: the 2x3 by 3x2 example, and a 7x5 by 5x9 pair of generated matrices (an independent check, since the reference is not compiled by Mountain Goat):
// A C++ program calling the Chapter 21 operations. The elementwise results are checked by the test against exact expected
// values; matmul is also compared against a plain C++ triple loop, including after the heap has been deliberately dirtied.
#include <cmath>
#include <cstdio>
#include <cstdlib>
#include "ops.h"
static bool same(const mg::Matrix &x, const mg::Matrix &y) {
if (x.rows != y.rows || x.cols != y.cols) return false;
for (size_t i = 0; i < x.data.size(); i++) if (std::fabs(x.data[i] - y.data[i]) > 1e-12) return false;
return true;
}
static mg::Matrix ref_matmul(const mg::Matrix &a, const mg::Matrix &b) {
mg::Matrix c(a.rows, b.cols, std::vector<double>(a.rows * b.cols, 0.0));
for (int64_t i = 0; i < a.rows; i++)
for (int64_t j = 0; j < b.cols; j++)
for (int64_t k = 0; k < a.cols; k++) c(i, j) += a(i, k) * b(k, j);
return c;
}
static void show(const char *name, const mg::Matrix &m) {
std::printf("%s (%ldx%ld):\n", name, (long)m.rows, (long)m.cols);
for (int64_t i = 0; i < m.rows; i++) { for (int64_t j = 0; j < m.cols; j++) std::printf(" %g", m(i, j)); std::printf("\n"); }
}
int main(int argc, char **argv) {
mg::Matrix a(2, 3, {1, 2, 3, 4, 5, 6}), b(3, 2, {7, 8, 9, 10, 11, 12}), c(2, 3, {6, 5, 4, 3, 2, 1});
show("sub(a, c)", mg::sub(a, c));
show("hadamard(a, c)", mg::hadamard(a, c));
show("divide(a, c)", mg::divide(a, c));
show("matmul(a, b)", mg::matmul(a, b));
show("gram(a)", mg::gram(a));
show("affine(a)", mg::affine(a));
std::printf("matmul matches a plain C++ triple loop: %s\n", same(mg::matmul(a, b), ref_matmul(a, b)) ? "yes" : "NO");
// Dirty the heap first: allocate a block, fill it with 1000s, free it. The compiled code never zeroes memory it did not
// fill itself, so a matmul that forgot to zero its accumulators would add into these 1000s and print 1015, 1022 instead of
// 15, 22. (Without this, a fresh heap is all zeros and the bug hides.)
{
double *junk = (double *)std::malloc(4 * sizeof(double));
for (int i = 0; i < 4; i++) junk[i] = 1000.0;
std::free(junk);
mg::Matrix sq(2, 2, {1, 2, 3, 4});
mg::Matrix r = mg::matmul(sq, sq);
std::printf("dirty-heap matmul: %g %g %g %g\n", r.data[0], r.data[1], r.data[2], r.data[3]);
}
// A bigger one, to compare against the reference on more than a hand-sized case.
std::vector<double> pd(35), qd(45);
for (size_t i = 0; i < 35; i++) pd[i] = (double)((i * 7) % 11) - 5;
for (size_t i = 0; i < 45; i++) qd[i] = (double)((i * 5) % 13) - 6;
mg::Matrix p(7, 5, pd), q(5, 9, qd);
std::printf("7x5 @ 5x9 matches the plain loop: %s\n", same(mg::matmul(p, q), ref_matmul(p, q)) ? "yes" : "NO");
if (argc > 1) {
std::printf("about to matmul 2x3 by 2x3...\n"); std::fflush(stdout);
mg::matmul(a, a);
std::printf("not reached\n");
}
return 0;
}
#!/bin/sh
# Compile ops.mg with mgc, then build and run the C++ program against it. Output: run_out.txt
HERE=$(cd "$(dirname "$0")" && pwd); W=$HERE/../work/cpp; mkdir -p $W
$HERE/../mgc lib $HERE/ops.mg -o $W
cp $W/ops.h $HERE/ops_generated.h
clang++-18 -std=c++17 -Wall -I$W $HERE/app.cpp $W/ops.o -o $W/app
echo "--- ./app"; $W/app
echo "--- ./app mismatch (exit status follows)"; $W/app mismatch; echo "exit=$?"
wrote work/cpp/ops.o and work/cpp/ops.h
--- ./app
sub(a, c) (2x3):
-5 -3 -1
1 3 5
hadamard(a, c) (2x3):
6 10 12
12 10 6
divide(a, c) (2x3):
0.166667 0.4 0.75
1.33333 2.5 6
matmul(a, b) (2x2):
58 64
139 154
gram(a) (2x2):
14 32
32 77
affine(a) (2x3):
3 5 7
9 11 13
matmul matches a plain C++ triple loop: yes
dirty-heap matmul: 7 10 15 22
7x5 @ 5x9 matches the plain loop: yes
--- ./app mismatch (exit status follows)
sub(a, c) (2x3):
-5 -3 -1
1 3 5
hadamard(a, c) (2x3):
6 10 12
12 10 6
divide(a, c) (2x3):
0.166667 0.4 0.75
1.33333 2.5 6
matmul(a, b) (2x2):
58 64
139 154
gram(a) (2x2):
14 32
32 77
affine(a) (2x3):
3 5 7
9 11 13
matmul matches a plain C++ triple loop: yes
dirty-heap matmul: 7 10 15 22
7x5 @ 5x9 matches the plain loop: yes
about to matmul 2x3 by 2x3...
mg.matmul: inner dimensions differ at runtime
Aborted
exit=134
Both comparisons print yes. The second run, ./app mismatch, calls matmul on a 2x3 and a 2x3, and the compiled code aborts with its message, as in example 12.
Tests¶
Seventeen test files in test/ops/ (fifteen written with the chapter and two added after the review below), and three more in test/tour/ for the language tour. Like Chapter 20's, the program tests run the real example files rather than copies.
| Test | What it checks |
|---|---|
matmul-inner-mismatch |
the verifier's exact message for a static inner-dimension mismatch |
matmul-bad-result-shape |
(2x3)@(3x4) declared as 2x3 is rejected |
scalar-bad-op |
mg.scalar with op = "pow" is rejected |
elementwise-static-mismatch |
sub, mul and div each reject a static mismatch (three split inputs) |
dynamic-accepted |
? dimensions verify and the new ops print and re-parse |
matmul-lowering-structure |
the static form has a zero-fill nest and a three-deep nest with mulf/addf and no assert; the dynamic form has the assert with its exact message |
values |
examples 1, 2, 4, 5: subtract, Hadamard, matmul, scalars, with exact values |
division-ieee |
2.5, nan, inf and no abort |
precedence-and-identities |
examples 6 and 9 |
composed-programs |
the linear layer and Gram matrices |
dynamic-ops |
examples 10 and 11 |
matmul-mismatch-aborts |
run-time abort, message on stderr, empty stdout |
front-end-errors |
both error files, exact text |
ptx-new-ops |
the GPU path: four kernels for the chain, two for matmul, with the arithmetic instructions |
cpp-interop-ops |
the C++ program builds with -Wall -Werror; every function's result matches exact hand-computed values; matmul matches a plain C++ loop, including on a deliberately dirtied heap; and a mismatch aborts with the message |
elementwise-mismatch-aborts |
added after the review: sub, mul and div abort on a run-time shape mismatch, naming the operation, with empty stdout |
negate-and-dynamic-divide |
added after the review: neg on static and dynamic shapes; dynamic division with order-sensitive numbers |
Do the tests fail when the code is wrong? prove_tests_can_fail.sh runs the new tests against Chapter 20's mg-opt (which has none of the new ops) and against six broken copies of the front end and driver:
#!/bin/sh
# Show that the Chapter 21 tests (test/ops) can fail. Output: prove_tests_can_fail_out.txt
# (1) against Chapter 20's mg-opt, which has none of the new ops;
# (2) against broken copies of the front end (no rebuild needed; restored on exit);
# (3) against rebuilt mg-opt copies with a bug injected into the new lowering (lowering_mutation.sh).
# READ THIS FIRST: this script runs the Chapter 21 tests against an older build and against deliberately broken copies of the front end and driver.
# FAIL lines (and non-zero exit statuses) in this script's output are the EXPECTED result: each broken copy SHOULD make at least one test fail (a FAIL line means a test caught the bug). The "baseline" sections must show all tests passing.
# The unmodified, current build is the baseline and must show no failures. A mutation or old build that makes NO test fail would be the problem.
echo "NOTE: this script deliberately runs broken or older code. FAIL lines below are EXPECTED: they show the tests can detect the problem. The baseline (unmodified current build) must show none."
HERE=$(cd "$(dirname "$0")" && pwd); LIT=$HERE/../../part15/code; W=$HERE/work/prove; rm -rf $W; mkdir -p $W
run() { label=$1; mgopt=$2
printf '%s\n' "### $label"
MG_OPT=$mgopt MG_TEST_TMP=$W/out sh $LIT/run_lit.sh --filter ':: ops/' 2>&1 | grep -E "^(FAIL)| Passed| Failed" | sed 's#IR Forge / Mountain Goat :: ##'
}
run "baseline: Chapter 21 build" $HERE/build/mg-opt
run "Chapter 20 build (none of the new ops exist)" $HERE/../../part20/code/build/mg-opt
cp $HERE/mgfront.py $W/mgfront.orig; cp $HERE/mgc $W/mgc.orig
restore() { cp $W/mgfront.orig $HERE/mgfront.py; cp $W/mgc.orig $HERE/mgc; }
trap restore EXIT
mutate() { # label file sed-expr
restore; cp $2 $W/before; sed -i "$3" $2
if cmp -s $2 $W/before; then echo "### $1: MUTATION DID NOT APPLY"; return; fi
run "front-end/driver mutation: $1" $HERE/build/mg-opt
}
mutate "'-' is emitted as mg.add" $HERE/mgfront.py 's/OPNAME = {"+": "add", "-": "sub"/OPNAME = {"+": "add", "-": "add"/'
mutate "'*' and '@' swapped (Hadamard vs matmul)" $HERE/mgfront.py 's/if op == "@":/if op == "*":/'
mutate "scalar on the left is not marked reversed" $HERE/mgfront.py 's/rev = v is None/rev = False/'
mutate "'*' and '+' on the same precedence level" $HERE/mgfront.py 's/while self.peek()\[1\] in ("\*", "\/", "@"):/while self.peek()[1] in ("*", "\/", "@", "+"):/'
mutate "matmul result shape is (m, k) instead of (m, n)" $HERE/mgfront.py 's/shape = (s\[0\], s2\[1\]); r = self.fn.tmp()/shape = (s[0], s[1]); r = self.fn.tmp()/'
mutate "driver forgets convert-scf-to-cf for the GPU path" $HERE/mgc 's/convert-scf-to-cf,convert-gpu-to-nvvm/convert-gpu-to-nvvm/'
restore; echo "### restored"; run "baseline again" $HERE/build/mg-opt
NOTE: this script deliberately runs broken or older code. FAIL lines below are EXPECTED: they show the tests can detect the problem. The baseline (unmodified current build) must show none.
### baseline: Chapter 21 build
Passed : 17 (18.89%)
### Chapter 20 build (none of the new ops exist)
FAIL: ops/dynamic-accepted.mlir (1 of 17)
FAIL: ops/division-ieee.mlir (2 of 17)
FAIL: ops/cpp-interop-ops.mlir (3 of 17)
FAIL: ops/composed-programs.mlir (4 of 17)
FAIL: ops/values.mlir (5 of 17)
FAIL: ops/elementwise-mismatch-aborts.mlir (6 of 17)
FAIL: ops/dynamic-ops.mlir (7 of 17)
FAIL: ops/precedence-and-identities.mlir (8 of 17)
FAIL: ops/negate-and-dynamic-divide.mlir (9 of 17)
FAIL: ops/ptx-new-ops.mlir (10 of 17)
FAIL: ops/matmul-mismatch-aborts.mlir (11 of 17)
FAIL: ops/elementwise-static-mismatch.mlir (13 of 17)
FAIL: ops/matmul-lowering-structure.mlir (14 of 17)
FAIL: ops/matmul-inner-mismatch.mlir (15 of 17)
FAIL: ops/matmul-bad-result-shape.mlir (16 of 17)
FAIL: ops/scalar-bad-op.mlir (17 of 17)
Passed : 1 (1.11%)
Failed : 16 (17.78%)
### front-end/driver mutation: '-' is emitted as mg.add
FAIL: ops/values.mlir (1 of 17)
FAIL: ops/precedence-and-identities.mlir (2 of 17)
FAIL: ops/ptx-new-ops.mlir (4 of 17)
FAIL: ops/elementwise-mismatch-aborts.mlir (7 of 17)
FAIL: ops/cpp-interop-ops.mlir (17 of 17)
Passed : 12 (13.33%)
Failed : 5 (5.56%)
### front-end/driver mutation: '*' and '@' swapped (Hadamard vs matmul)
FAIL: ops/precedence-and-identities.mlir (1 of 17)
FAIL: ops/cpp-interop-ops.mlir (2 of 17)
FAIL: ops/values.mlir (3 of 17)
FAIL: ops/ptx-new-ops.mlir (4 of 17)
FAIL: ops/composed-programs.mlir (5 of 17)
FAIL: ops/dynamic-ops.mlir (6 of 17)
FAIL: ops/matmul-mismatch-aborts.mlir (7 of 17)
FAIL: ops/front-end-errors.mlir (8 of 17)
FAIL: ops/elementwise-mismatch-aborts.mlir (17 of 17)
Passed : 8 (8.89%)
Failed : 9 (10.00%)
### front-end/driver mutation: scalar on the left is not marked reversed
FAIL: ops/values.mlir (3 of 17)
Passed : 16 (17.78%)
Failed : 1 (1.11%)
### front-end/driver mutation: '*' and '+' on the same precedence level
FAIL: ops/precedence-and-identities.mlir (1 of 17)
Passed : 16 (17.78%)
Failed : 1 (1.11%)
### front-end/driver mutation: matmul result shape is (m, k) instead of (m, n)
FAIL: ops/values.mlir (2 of 17)
FAIL: ops/composed-programs.mlir (3 of 17)
Passed : 15 (16.67%)
Failed : 2 (2.22%)
### front-end/driver mutation: driver forgets convert-scf-to-cf for the GPU path
FAIL: ops/ptx-new-ops.mlir (8 of 17)
Passed : 16 (17.78%)
Failed : 1 (1.11%)
### restored
### baseline again
Passed : 17 (18.89%)
Front-end bugs are cheap to inject. Bugs in the lowering need mg-opt rebuilt (a full rebuild each), so lowering_mutation.sh injects five, one at a time, into a copy of the source tree:
#!/bin/sh
# Rebuild mg-opt with one bug injected into LowerToAffine.cpp and run the ops tests against it.
# Each build is a full rebuild in work/mutN (about two minutes). Output: lowering_mutation_out.txt
# READ THIS FIRST: this script rebuilds mg-opt with one deliberate bug injected into the lowering at a time (a full rebuild each) and runs the tests.
# FAIL lines (and non-zero exit statuses) in this script's output are the EXPECTED result: each bug SHOULD make at least one test fail (a FAIL line means a test caught it). A section with NO FAIL line would be a test gap.
# The unmodified, current build is the baseline and must show no failures. A mutation or old build that makes NO test fail would be the problem.
echo "NOTE: this script deliberately runs broken or older code. FAIL lines below are EXPECTED: they show the tests can detect the problem. The baseline (unmodified current build) must show none."
HERE=$(cd "$(dirname "$0")" && pwd); LIT=$HERE/../../part15/code; M=$HERE/work/mut; rm -rf $M; mkdir -p $M
n=0
try() { # label sed-expr
n=$((n+1)); d=$M/$n; mkdir -p $d/build; cp -r $HERE/tree $d/tree; rm -rf $d/tree/build
cp $d/tree/lib/LowerToAffine.cpp $d/before.cpp; sed -i "$2" $d/tree/lib/LowerToAffine.cpp
if cmp -s $d/before.cpp $d/tree/lib/LowerToAffine.cpp; then echo "### $1: MUTATION DID NOT APPLY"; return; fi
( cd $d/build && cmake $d/tree -DMLIR_DIR=/usr/lib/llvm-18/lib/cmake/mlir -DLLVM_DIR=/usr/lib/llvm-18/lib/cmake/llvm -DCMAKE_BUILD_TYPE=Release -DLLVM_ENABLE_ASSERTIONS=OFF >/dev/null 2>&1 && mkdir -p include/mg && make -j8 >/dev/null 2>&1 )
[ -x $d/build/mg-opt ] || { echo "### $1: BUILD FAILED"; return; }
printf '%s\n' "### lowering mutation: $1"
MG_OPT=$d/build/mg-opt MG_TEST_TMP=$d/out sh $LIT/run_lit.sh --filter ':: ops/' 2>&1 | grep -E "^(FAIL)| Passed| Failed" | sed 's#IR Forge / Mountain Goat :: ##'
}
try "matmul accumulator not zero-filled" 's/b.create<affine::AffineStoreOp>(l, zero, alloc, ivs);/(void)zero;/'
try "scalar op ignores reversed (10 - x computed as x - 10)" 's/Value lhs = rev ? c : x, rhs = rev ? x : c, r;/Value lhs = x, rhs = c, r;/'
try "matmul skips the run-time inner-dimension check" 's/if (lt.isDynamicDim(1) || rt.isDynamicDim(0)) {/if (false) {/'
try "mg.div lowered as a multiply" 's/ElementwiseBinaryLowering<mg::DivOp, arith::DivFOp>/ElementwiseBinaryLowering<mg::DivOp, arith::MulFOp>/'
try "mg.neg lowered as identity (no negation)" 's/b.create<arith::NegFOp>(l, x)/x/'
# Added after an independent review (Chapter 21's "After the review" section):
try "elementwise operands swapped (computes rhs OP lhs)" 's/b.create<affine::AffineStoreOp>(l, b.create<ArithOp>(l, x, y), alloc, ivs);/b.create<affine::AffineStoreOp>(l, b.create<ArithOp>(l, y, x), alloc, ivs);/'
try "elementwise drops its run-time shape check" '/struct ElementwiseBinaryLowering/,/^};/ s/mg::dyn::assertSameShape(rewriter, loc, lhs, rhs,/if (false) mg::dyn::assertSameShape(rewriter, loc, lhs, rhs,/'
try "matmul takes n from the wrong operand (lhs dim 1)" 's/Value n = bound(outType.getDimSize(1), rhs, 1, Value());/Value n = bound(outType.getDimSize(1), lhs, 1, Value());/'
### lowering mutation: matmul accumulator not zero-filled
FAIL: ops/matmul-lowering-structure.mlir (9 of 17)
FAIL: ops/cpp-interop-ops.mlir (11 of 17)
Passed : 15 (16.67%)
Failed : 2 (2.22%)
### lowering mutation: scalar op ignores reversed (10 - x computed as x - 10)
FAIL: ops/values.mlir (17 of 17)
Passed : 16 (17.78%)
Failed : 1 (1.11%)
### lowering mutation: matmul skips the run-time inner-dimension check
FAIL: ops/matmul-lowering-structure.mlir (9 of 17)
FAIL: ops/cpp-interop-ops.mlir (10 of 17)
FAIL: ops/matmul-mismatch-aborts.mlir (11 of 17)
Passed : 14 (15.56%)
Failed : 3 (3.33%)
### lowering mutation: mg.div lowered as a multiply
FAIL: ops/division-ieee.mlir (2 of 17)
FAIL: ops/cpp-interop-ops.mlir (10 of 17)
FAIL: ops/negate-and-dynamic-divide.mlir (14 of 17)
FAIL: ops/precedence-and-identities.mlir (16 of 17)
Passed : 13 (14.44%)
Failed : 4 (4.44%)
### lowering mutation: mg.neg lowered as identity (no negation)
FAIL: ops/negate-and-dynamic-divide.mlir (12 of 17)
FAIL: ops/precedence-and-identities.mlir (15 of 17)
Passed : 15 (16.67%)
Failed : 2 (2.22%)
### lowering mutation: elementwise operands swapped (computes rhs OP lhs)
FAIL: ops/division-ieee.mlir (2 of 17)
FAIL: ops/cpp-interop-ops.mlir (11 of 17)
FAIL: ops/negate-and-dynamic-divide.mlir (14 of 17)
FAIL: ops/values.mlir (15 of 17)
Passed : 13 (14.44%)
Failed : 4 (4.44%)
### lowering mutation: elementwise drops its run-time shape check
FAIL: ops/elementwise-mismatch-aborts.mlir (4 of 17)
Passed : 16 (17.78%)
Failed : 1 (1.11%)
### lowering mutation: matmul takes n from the wrong operand (lhs dim 1)
FAIL: ops/dynamic-ops.mlir (3 of 17)
FAIL: ops/cpp-interop-ops.mlir (10 of 17)
Passed : 15 (16.67%)
Failed : 2 (2.22%)
Read the second list closely, and read the corrected account that follows it (an independent review found that my first explanation here was wrong): removing the zero-fill was caught by only two of the original fifteen tests, the lowering-structure test and, as it turned out, the C++ program. No .mg program that prints a result noticed.
The same list shows the other mutations doing what they should. Ignoring the reversed flag is caught by values; dropping the run-time inner-dimension check is caught by three tests; compiling div as a multiply, and neg as the identity, were each caught by a single check, which is thinner than it looks (the review below made them direct).
After an independent review¶
I asked a second model (Opus) to review these mutation results and the tests behind them, without editing anything. It found real problems, which I confirmed before acting, and which this section records. This is the kind of thing a mutation run is for, and the kind of thing I missed on my own.
1. My explanation of the zero-fill result was wrong. I had written that malloc returned "fresh, already-zero memory for a small buffer". The real reason is different and more important: the generated code never frees anything (LowerToAffine.cpp contains no dealloc), so in a program run with mgc run every allocation comes from heap memory nobody has used yet, which the operating system hands out zeroed, whatever its size. No program written in Mountain Goat can ever catch a missing zero-fill, because it cannot reuse freed memory. The C++ wrapper does free its results (take() calls std::free), so a C++ program can reuse memory, and that is why the C++ test noticed. But it noticed by accident of how the C library recycles blocks: against the broken build, the second 2x2 matmul in the C++ program printed NO, because it reused a small block freed earlier, while the large 7x5-by-5x9 comparison printed yes, because its block was fresh. I had credited the large comparison; it was the small one. A test that passes or fails depending on the allocator is not a test to rely on.
2. A deterministic test now exists. The C++ program allocates a small block, fills it with 1000s, frees it, and then calls matmul. With the zero-fill, the result is 7 10 15 22. Without it, the accumulators start from the leftover 1000s and the result is 7 10 1015 1022, every time. The check is dirty-heap matmul: 7 10 15 22 in cpp-interop-ops.
3. Results that were printed but never checked. The C++ program printed the results of sub, hadamard, divide, gram and affine, but the test only checked matmul; under the mutation that compiled div as a multiply, that test still passed. The test now checks every function's output against hand-computed values. (The program's own header comment said "compare every result against plain C++ loops"; only matmul was. The comment is corrected.)
4. The run-time shape check was only tested for +. sub, mul and div share a lowering template with their own copy of the check (assertSameShape), and nothing exercised it for them. New examples 15_elementwise_mismatch.mg, 15b_mul_mismatch.mg and 15c_div_mismatch.mg abort with mg.sub:, mg.mul: and mg.div: operand shapes differ at runtime in dimension 0, with an empty stdout (test elementwise-mismatch-aborts).
5. neg and dynamic div now have direct tests (16_negate.mg, 17_dynamic_divide.mg; test negate-and-dynamic-divide): negation on static and dynamic shapes with a negative input, and a dynamic division with numbers where order matters ([[1, 4], [9, 3]] / [[2, 8], [3, 4]] is [[0.5, 0.5], [3, 0.75]]). Negation had been caught by exactly one check line, and -(-a), which I had cited as an identity check, cannot see a negation that does nothing.
6. dynamic-ops now checks every value it prints, including the 2x1 and 3x2 cases that previously only had their sizes checked.
Three new mutations accompany these, and the script was rerun in full (eight rebuilds):
### lowering mutation: matmul accumulator not zero-filled
FAIL: ops/matmul-lowering-structure.mlir (9 of 17)
FAIL: ops/cpp-interop-ops.mlir (11 of 17)
Passed : 15 (16.67%)
Failed : 2 (2.22%)
### lowering mutation: scalar op ignores reversed (10 - x computed as x - 10)
FAIL: ops/values.mlir (17 of 17)
Passed : 16 (17.78%)
Failed : 1 (1.11%)
### lowering mutation: matmul skips the run-time inner-dimension check
FAIL: ops/matmul-lowering-structure.mlir (9 of 17)
FAIL: ops/cpp-interop-ops.mlir (10 of 17)
FAIL: ops/matmul-mismatch-aborts.mlir (11 of 17)
Passed : 14 (15.56%)
Failed : 3 (3.33%)
### lowering mutation: mg.div lowered as a multiply
FAIL: ops/division-ieee.mlir (2 of 17)
FAIL: ops/cpp-interop-ops.mlir (10 of 17)
FAIL: ops/negate-and-dynamic-divide.mlir (14 of 17)
FAIL: ops/precedence-and-identities.mlir (16 of 17)
Passed : 13 (14.44%)
Failed : 4 (4.44%)
### lowering mutation: mg.neg lowered as identity (no negation)
FAIL: ops/negate-and-dynamic-divide.mlir (12 of 17)
FAIL: ops/precedence-and-identities.mlir (15 of 17)
Passed : 15 (16.67%)
Failed : 2 (2.22%)
### lowering mutation: elementwise operands swapped (computes rhs OP lhs)
FAIL: ops/division-ieee.mlir (2 of 17)
FAIL: ops/cpp-interop-ops.mlir (11 of 17)
FAIL: ops/negate-and-dynamic-divide.mlir (14 of 17)
FAIL: ops/values.mlir (15 of 17)
Passed : 13 (14.44%)
Failed : 4 (4.44%)
### lowering mutation: elementwise drops its run-time shape check
FAIL: ops/elementwise-mismatch-aborts.mlir (4 of 17)
Passed : 16 (17.78%)
Failed : 1 (1.11%)
### lowering mutation: matmul takes n from the wrong operand (lhs dim 1)
FAIL: ops/dynamic-ops.mlir (3 of 17)
FAIL: ops/cpp-interop-ops.mlir (10 of 17)
Passed : 15 (16.67%)
Failed : 2 (2.22%)
Before and after, for the five original mutants, counting the test files that fail out of the seventeen:
| Mutation | Before the review | After | Who catches it now |
|---|---|---|---|
| matmul accumulator not zero-filled | 2 (one by allocator luck) | 2 | structure test, and the dirty-heap check (deterministic) |
scalar op ignores reversed |
1 | 1 | values: 10 - a and 16 / a have a direct, deterministic check |
| matmul skips the run-time inner-dimension check | 3 | 3 | structure, the abort test, the C++ abort run |
div computed as a multiply |
2 | 4 | division, identities, dynamic divide, C++ output |
neg computed as the identity |
1 | 2 | identities, and the new direct neg test |
The three new mutations are each caught as well: operands swapped (4 tests), the elementwise run-time shape check removed (1 test, the new elementwise-mismatch-aborts, which is the only thing that exercises that code), and the matmul output width read from the wrong operand (2 tests).
Two honest remarks. First, "caught by only one test" is not automatically a weakness: reversed and the shape check are each caught by a direct, deterministic check of exactly that behavior, which is what you want. What was a weakness was a catch that depended on chance (the allocator) or on a test that did not look at the thing it printed. Second, the zero-fill mutant is still caught by only two checks, and one of them (the structure test) looks at the implementation, not the behavior. That is the best available: the language cannot observe the bug, so only the C++ side, which can reuse memory, can.
The full suite after these additions:
The new examples' real output:
################ ./mgc run examples/15_elementwise_mismatch.mg
| # The run-time shape check is not only for '+': subtract, multiply and divide each abort on a dimension mismatch,
| # naming the operation. Here 1x3 minus 2x1: the 1x3 is a ?-typed parameter, so nothing broadcasts and the check runs.
| def f(a: tensor[?x?], b: tensor[?x?]) = a - b
| print f([[1, 2, 3]], [[1], [2]])
----- output
mg.sub: operand shapes differ at runtime in dimension 0
Aborted
----- exit status 134
################ ./mgc run examples/15b_mul_mismatch.mg
| def f(a: tensor[?x?], b: tensor[?x?]) = a * b
| print f([[1, 2, 3]], [[1], [2]])
----- output
mg.mul: operand shapes differ at runtime in dimension 0
Aborted
----- exit status 134
################ ./mgc run examples/15c_div_mismatch.mg
| def f(a: tensor[?x?], b: tensor[?x?]) = a / b
| print f([[1, 2, 3]], [[1], [2]])
----- output
mg.div: operand shapes differ at runtime in dimension 0
Aborted
----- exit status 134
################ ./mgc run examples/16_negate.mg
| # Negation, on its own: static and dynamic shapes. (-(-a) = a cannot tell a real negation from a no-op; these can.)
| def h(a: tensor[?x?]) = -a
| print h([[1, -2, 3]])
| print -[[0.5, -4]]
| print h([[7], [-8]])
----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 3] strides = [3, 1] data =
[[-1, 2, -3]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [1, 2] strides = [2, 1] data =
[[-0.5, 4]]
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 1] strides = [1, 1] data =
[[-7],
[8]]
----- exit status 0
################ ./mgc run examples/17_dynamic_divide.mg
| # Division on dynamic shapes, with numbers where the order matters (a / b is not b / a).
| def g(a: tensor[?x?], b: tensor[?x?]) = a / b
| print g([[1, 4], [9, 3]], [[2, 8], [3, 4]])
----- output
Unranked Memref base@ = 0x… rank = 2 offset = 0 sizes = [2, 2] strides = [2, 1] data =
[[0.5, 0.5],
[3, 0.75]]
----- exit status 0
The C++ program's output, with the new check:
wrote work/cpp/ops.o and work/cpp/ops.h
--- ./app
sub(a, c) (2x3):
-5 -3 -1
1 3 5
hadamard(a, c) (2x3):
6 10 12
12 10 6
divide(a, c) (2x3):
0.166667 0.4 0.75
1.33333 2.5 6
matmul(a, b) (2x2):
58 64
139 154
gram(a) (2x2):
14 32
32 77
affine(a) (2x3):
3 5 7
9 11 13
matmul matches a plain C++ triple loop: yes
dirty-heap matmul: 7 10 15 22
7x5 @ 5x9 matches the plain loop: yes
--- ./app mismatch (exit status follows)
sub(a, c) (2x3):
-5 -3 -1
1 3 5
hadamard(a, c) (2x3):
6 10 12
12 10 6
divide(a, c) (2x3):
0.166667 0.4 0.75
1.33333 2.5 6
matmul(a, b) (2x2):
58 64
139 154
gram(a) (2x2):
14 32
32 77
affine(a) (2x3):
3 5 7
9 11 13
matmul matches a plain C++ triple loop: yes
dirty-heap matmul: 7 10 15 22
7x5 @ 5x9 matches the plain loop: yes
about to matmul 2x3 by 2x3...
mg.matmul: inner dimensions differ at runtime
Aborted
exit=134
The full suite against the Chapter 21 build:
Bugs found while writing the language tour¶
Writing the tour page (docs/tour/language-tour.md) with real, runnable examples found two problems in this chapter's front end, both now fixed and covered by tests:
- Scientific notation broke the compiled program. The front end accepted
3e2but copied it into the MLIR, and MLIR only accepts a decimal point before an exponent (3.0e2). The error came from MLIR (expected ']'), pointing into a file the user never wrote. The front end now rewrites every number through one helper (mlir_float), andscientific-literalschecks the MLIR text and the compiled result (including a scalar with a large exponent,1e20). - Two confusing messages. A
defwith no expression (for example, putting aletin its body) producedunexpected None, andtensor[0x3]was accepted. They now sayexpected an expression, found the end of the lineandbad dimension '0': use a positive integer or '?'.
Limits and what is not established¶
- Rank 2,
f64only. No other ranks, no other element types, no integer arithmetic. - No broadcasting (added for static shapes in Chapter 22). In this chapter
a + bneeds equal shapes (or?resolved at run time). - No compile-time folding for the new operations (only
mg.addhas a folder).1 + 1-style scalar arithmetic is folded by the front end, butmg.subof two constant matrices is computed at run time. mg.scalaruses a string attribute for its operation name; an enum would be stricter.- Matrix product is the textbook triple loop. No tiling, no blocking, no vectorization, no library call. No performance claim is made at all, and nothing here measures speed.
- Floating-point caveats: the matrix product's sum order is the loop order, and floating-point addition is not associative, so a different order (say, a GPU running the reduction differently, or a library) can give slightly different results. Division by zero produces
inf/nansilently. - The bufferization path (Chapter 6) was not extended.
--one-shot-bufferizeknowsmg.addandmg.transposeonly. The new operations lower through the hand-written conversion, whichmgcuses. - GPU: PTX is produced and its structure is inspected. It was never assembled, loaded or launched, and its correctness and speed on a GPU are not established. The matmul kernels are one-thread-per-output-element with a sequential reduction, nothing cleverer.
- Platform: x86-64 Linux, LLVM 18.1.3. POSIX only (Chapter 19).
- The language itself is described in the language tour. This chapter only adds operations to it.
Reproducing¶
cd docs/part21/code
./build.sh && ./demo.sh > demo_out.txt && ./show.sh && cpp/run.sh
./prove_tests_can_fail.sh # front-end mutations (a few seconds)
./lowering_mutation.sh # lowering mutations (a full rebuild each: roughly 10 minutes)
cd ../../part15/code && ./run_lit.sh # 78 tests
Chapter summary¶
- Four new operations and one new operand kind:
mg.sub,mg.mul(Hadamard),mg.div,mg.neg,mg.matmul, andmg.scalar(add, sub, mul, div with a number, on either side). *is elementwise and@is the matrix product: one operator, one meaning.- Elementwise operations share one verifier and one lowering template;
matmulis a zero-fill plus a three-deep accumulate loop, with one run-time check of the inner dimensions;scalarneeds no check at all. - Division follows IEEE 754:
x/0isinfornan, not an error. - Fourteen worked examples ran for real, including a linear layer, a Gram matrix, a dynamic matmul and a run-time abort.
- The GPU path produced PTX for all of it, which exposed two real gaps (no
scf-to-cffor loops inside kernels; a decoder that printed only the first kernel), both fixed. - C++ calls the new operations and matches a plain C++ loop on a 7x5 by 5x9 product.
- Seventeen test files (plus three for the tour); each file was shown to fail under at least one injected bug, and an independent review then found and fixed the places where a single check was carrying too much.
Self-check questions¶
Each answer is collapsed; try the question first.
-
What are the three meanings of "multiply" for matrices, and which operator in Mountain Goat is each?
Answer
Elementwise (Hadamard): each element times the matching element, written
*, needs equal shapes. Matrix product: row-by-column dot products, written@, needs only the inner dimensions to match and gives (rows of the left) by (columns of the right). By a scalar: every element times one number, written*with a number on one side. The language disambiguates by operator and by operand type (matrix or number), never by guessing from shapes. -
Why does
a - 1give a different answer from10 - a, and how does the compiler represent the difference?Answer
Subtraction is not commutative:
a - 1subtracts 1 from each element, while10 - asubtracts each element from 10. Both aremg.scalarwithop = "sub"; thereversedattribute isfalsefor the first andtruefor the second. The lowering uses it to decide which side ofarith.subfthe constant goes on (subf %0, %cstversussubf %cst, %0). The same applies to division (16 / aversusa / 16). -
What is a reduction, and which loop in the lowered matmul is one?
Answer
A reduction combines many values into one, here by summing. In the lowered matrix product it is the innermost loop over
k: for each output position(i, j), it sumslhs[i][k] * rhs[k][j]over allkinto the same memory location. Theiandjloops are independent of each other (each result element is its own computation), which is why the GPU recipe makes them parallel and leaves thekloop sequential inside each thread. -
Why must the result of a matmul be zero-filled first?
Answer
The main loop accumulates: it reads the current value at
result[i][j], adds a product, and writes it back.memref.allocgives memory with unspecified contents, so without the fill every sum would start from garbage. The lowering therefore runs a separate loop nest that stores 0.0 everywhere before the accumulate nest. One of this chapter's mutations removes exactly this step. Only the lowering-structure test and one deliberately dirtied-heap test in the C++ program reliably notice (see "After an independent review" below); no program written in the language itself can, because the generated code never reuses memory. -
mg.matmul's verifier does not require the two operand shapes to be equal. What does it require, and what happens when a dimension is??Answer
It requires the inner dimensions to be compatible (the left operand's columns against the right operand's rows) and the result to be (left rows) by (right columns). "Compatible" is Chapter 13's rule: equal, or
?on either side. When the inner dimension is?on either side, the verifier accepts it, and the lowering emits onecf.assertthat compares the two extents at run time. A failed check aborts withmg.matmul: inner dimensions differ at runtimeon stderr. -
Dividing a matrix by zero does not stop the program, but dividing a scalar by a scalar zero in source is an error. Why the difference?
Answer
Matrix division is carried out at run time by the hardware's floating-point divide, which follows IEEE 754:
x/0givesinfor-infand0/0givesnan; there is no exception to raise. Scalar-by-scalar arithmetic is folded by the front end while compiling, where a literal division by zero is certainly a mistake it can see and report with a line number. The front end cannot tell what a matrix will contain, so it does not guess. -
Why did the first attempt to compile
matmulto PTX fail, and why had Chapter 10 not hit that?Answer
The
kloop stays sequential, so it ends up as anscf.forinside the GPU kernel. The device pipeline from Chapter 10 converted the GPU dialect to NVVM but never converted structured loops (scf) to branches (cf), so a leftover cast could not be legalized. Chapter 10's kernels had no loops inside them (every loop became a thread index), so the missing pass was never needed. Addingconvert-scf-to-cfinsidegpu.module(...)fixed it. -
Why does
(a - b) * a / 2 + 1produce four kernels, and what does that say about the generated code?Answer
Each operation lowers to its own loop nest over its own result buffer, and each nest becomes one kernel. So four operations give four kernels, each reading the previous one's output from memory and writing a new buffer. Nothing fused them into one pass over the data. That is correct but does more memory traffic than a fused version would; the book does not measure it and makes no speed claim.
-
The test
matmul-mismatch-abortschecks stdout as well as stderr. What does the stdout check establish?Answer
That the abort happens before anything is printed: the program printed nothing to stdout, and the message went only to stderr. Without that check, a bug that printed a partial or wrong result and then aborted would pass.
-
The C++ test compares Mountain Goat's matmul against a hand-written C++ triple loop. Why is that a stronger check than the values in example 4?
Answer
Example 4's expected values were computed by hand for one small pair of matrices; the C++ comparison checks 7x5 by 5x9 matrices with 63 outputs, against an implementation that Mountain Goat did not produce. A bug that coincidentally gives the right numbers on a tiny input (say, a transposed index that happens to be symmetric) is unlikely to survive a non-square, non-symmetric comparison.