Reference: https://mlir.llvm.org/docs/Tutorials/Toy/Ch-2/

Note: Check Setup the Environment of MLIR for the environment setup.

1. Run Example

Define a new env var:

export TOY_CH2_HOME="$MLIR_HOME/examples/toy/Ch2"

Create a file $TOY_CH2_HOME/input.toy ; Add the following content to the file:

# User defined generic function that operates on unknown shaped arguments.
def multiply_transpose(a, b) {
  return transpose(a) * transpose(b);
}

def main() {
  var a<2, 3> = [[1, 2, 3], [4, 5, 6]];
  var b<2, 3> = [1, 2, 3, 4, 5, 6];
  var c = multiply_transpose(a, b);
  var d = multiply_transpose(b, a);
  print(d);
}

Emit to AST (Abstract Syntax Tree):

toy-ch2 $TOY_CH2_HOME/input.toy -emit=ast

Emit to MLIR (Multi-Level Intermediate Representation):

toyc-ch2 $TOY_CH2_HOME/input.toy -emit=mlir

2. Add an Operator

2.1. Define the Operation

Add following code to $TOY_CH2_HOME/include/toy/Ops.td:

// SubtractOp

def SubtractOp : Toy_Op<"subtract"> {
  let summary = "element-wise subtraction operation";
  let description = [{
    The "subtract" operation performs element-wise subtraction between two
    tensors. The shapes of the tensor operands are expected to match.
  }];

  let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);
  let results = (outs F64Tensor);

  // Indicate that the operation has a custom parser and printer method.
  let hasCustomAssemblyFormat = 1;

  // Allow building an AddOp with from the two input operands.
  let builders = [
    OpBuilder<(ins "Value":$lhs, "Value":$rhs)>
  ];
}

Build MLIR again to imply the modifications:

bash $LLVM_PROJ_HOME/scripts/build-mlir.sh

Build errors pop out, because:

  • hasCustomAssemblyFormat is assigned with 1, but the custom parser and printer method is not implemented.
  • OpBuilder is not implemented.

These errors will be handled later.

Note that the C++ implementation of class SubtractOp has been generated in $LLVM_PROJ_HOME/build/tools/mlir/examples/toy/Ch2/include/toy/Ops.h.inc, and as a result, you are now able to use class SubtractOp with code completion and syntax highlighting of clangd.

2.2. Implement the Operations

To implement custom parser and printer methods as well as OpBuilder, add the following code to $TOY_CH2_HOME/mlir/Dialect.cpp:

// SubtractOp

void SubtractOp::build(mlir::OpBuilder &builder, mlir::OperationState &state, mlir::Value lhs, mlir::Value rhs) {
  state.addTypes(UnrankedTensorType::get(builder.getF64Type()));
  state.addOperands({lhs, rhs});
}

mlir::ParseResult SubtractOp::parse(mlir::OpAsmParser &parser, mlir::OperationState &result) {
  return parseBinaryOp(parser, result);
}

void SubtractOp::print(mlir::OpAsmPrinter &p) { printBinaryOp(p, *this); }

2.3. Emit - Operator

Go to $TOY_CH2_HOME/mlir/MLIRGen.cpp, locate function mlirGen and add the specific case for -, as shown below:

mlir::Value mlirGen(BinaryExprAST &binop) {
    // ...
    switch (binop.getOp()) {
    case '+':
      return builder.create<AddOp>(location, lhs, rhs);
    case '*':
      return builder.create<MulOp>(location, lhs, rhs);
    case '-':
      return builder.create<SubtractOp>(location, lhs, rhs);
    }
    // ...
}

Rebuild the MLIR:

$LLVM_PROJ_HOME/scripts/build-mlir.sh

2.4. Test the - Operator

Change the content of $MLIR_HOME/input.toy to:

# User defined generic function that operates on unknown shaped arguments.
def multiply_transpose(a, b) {
  return transpose(a) * transpose(b);
}

def main() {
  var a<2, 3> = [[1, 2, 3], [4, 5, 6]];
  var b<2, 3> = [1, 2, 3, 4, 5, 6];
  var c = multiply_transpose(a, b);
  var d = multiply_transpose(b, a);
  var e = a - b;
  print(e);
}

Generate MLIR:

toyc-ch2 $TOY_CH2_HOME/input.toy -emit=mlir