/// This class defines the interface for handling inlining with Toy operations.

/// We simplify inherit from the base interface class and override

/// the necessary methods.

struct ToyInlinerInterface : public DialectInlinerInterface {

  using DialectInlinerInterface::DialectInlinerInterface;

  /// This hook checks to see if the given callable operation is legal to inline

  /// into the given call. For Toy this hook can simply return true, as the Toy

  /// Call operation is always inlinable.

  bool isLegalToInline(Operation *call, Operation *callable,  // 当遇到一个调用的时候，检查是否允许内联

                       bool wouldBeCloned) const final {

    return true;  // 所有调用都支持内联

  }

  /// This hook checks to see if the given operation is legal to inline into the

  /// given region. For Toy this hook can simply return true, as all Toy

  /// operations are inlinable.

  bool isLegalToInline(Operation *, Region *, bool,  // 内联操作

                       IRMapping &) const final {

    return true;

  }

  /// This hook cheks if the given 'src' region can be inlined into the 'dest'

  /// region. The regions here are the bodies of the callable functions. For

  /// Toy, any function can be inlined, so we simply return true.

  bool isLegalToInline(Region *dest, Region *src, bool wouldBeCloned,  // 内联区域

                       IRMapping &valueMapping) const final {

    return true;

  }

  /// This hook is called when a terminator operation has been inlined. The only

  /// terminator that we have in the Toy dialect is the return

  /// operation(toy.return). We handle the return by replacing the values

  /// previously returned by the call operation with the operands of the

  /// return.

  void handleTerminator(Operation *op,

                        ValueRange valuesToRepl) const final {

    // Only "toy.return" needs to be handled here.

    auto returnOp = cast<ReturnOp>(op);

    // Replace the values directly with the return operands.

    assert(returnOp.getNumOperands() == valuesToRepl.size());

    for (const auto &it : llvm::enumerate(returnOp.getOperands()))

      valuesToRepl[it.index()].replaceAllUsesWith(it.value());

  }

};

// 被内联的函数

toy.func @foo(%arg: f64) -> f64 {

  %0 = toy.mul %arg, %arg : f64

  toy.return %0 : f64           ← 终结符

}

// 调用点

toy.func @main() {

  %result = toy.generic_call @foo(%arg) : ... -> f64

}

1. MLIR 将 @foo 的代码复制到 @main 中。

2. 被复制的代码中有 toy.return %0。

3. handleTerminator 被调用：

   - returnOp = %0

   - valuesToRepl = [%result]（调用点的结果值）

4. valuesToRepl[0].replaceAllUsesWith(%0)

   → 所有使用 %result 的地方都被替换为 %0

5. toy.return 被移除

/// Emit a new function and add it to the MLIR module.

mlir::toy::FuncOp mlirGen(FunctionAST &funcAST) {

  ...

  // If this function isn't main, then set the visibility to private.

  if (funcAST.getProto()->getName() != "main")

    function.setPrivate();

  return function;

}

void ToyDialect::initialize() {

  addInterfaces<ToyInlinerInterface>();

}

def FuncOp : Toy_Op<"func",

    [FunctionOpInterface, IsolatedFromAbove]> {  // 添加 FunctionOpInterface 特性

  ...

}

def GenericCallOp : Toy_Op<"generic_call",

    [DeclareOpInterfaceMethods<CallOpInterface>]> {  // 添加 DeclareOpInterfaceMethods 特性

  ...

}

def FuncOp : Toy_Op<"func", [FunctionOpInterface, IsolatedFromAbove]> {

  // ...

  let extraClassDeclaration = [{

    ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }

    ArrayRef<Type> getResultTypes() { return getFunctionType().getResults(); }

    Region *getCallableRegion() { return &getBody(); }

  }];

}

let arguments = (ins

  ...

  OptionalAttr<DictArrayAttr>:$arg_attrs,  // 参数属性

  OptionalAttr<DictArrayAttr>:$res_attrs   // 结果属性

);

pm.addPass(mlir::createInlinerPass());

toy.func @multiply_transpose(%arg0: tensor<*xf64>, %arg1: tensor<*xf64>) -> tensor<*xf64> {

  %0 = toy.transpose(%arg0 : tensor<*xf64>) to tensor<*xf64>

  %1 = toy.transpose(%arg1 : tensor<*xf64>) to tensor<*xf64>

  %2 = toy.mul %0, %1 : tensor<*xf64>

  toy.return %2 : tensor<*xf64>

}

toy.func @main() {

  %0 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>

  %1 = toy.reshape(%0 : tensor<2x3xf64>) to tensor<2x3xf64>

  %2 = toy.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00, 5.000000e+00, 6.000000e+00]> : tensor<6xf64>

  %3 = toy.reshape(%2 : tensor<6xf64>) to tensor<2x3xf64>

  %4 = toy.generic_call @multiply_transpose(%1, %3) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64>

  %5 = toy.generic_call @multiply_transpose(%3, %1) : (tensor<2x3xf64>, tensor<2x3xf64>) -> tensor<*xf64>

  toy.print %5 : tensor<*xf64>

  toy.return

}

def CastOp : Toy_Op<"cast", [

    DeclareOpInterfaceMethods<CastOpInterface>,  // 增加一个 CastOpInterface 特性

    Pure,

    SameOperandsAndResultShape]

  > {

  let summary = "shape cast operation";

  let description = [{

    The "cast" operation converts a tensor from one type to an equivalent type

    without changing any data elements. The source and destination types

    must both be tensor types with the same element type. If both are ranked,

    then shape is required to match. The operation is invalid if converting

    to a mismatching constant dimension.

  }];

  let arguments = (ins F64Tensor:$input);

  let results = (outs F64Tensor:$output);

  let assemblyFormat = "$input attr-dict `:` type($input) `to` type($output)";

}

/// Returns true if the given set of input and result types are compatible with

/// this cast operation. This is required by the `CastOpInterface` to verify

/// this operation and provide other additional utilities.

bool CastOp::areCastCompatible(TypeRange inputs, TypeRange outputs) {

  if (inputs.size() != 1 || outputs.size() != 1)  // 只接受一个输入和一个输出，否则不合法

    return false;

  // The inputs must be Tensors with the same element type.

  TensorType input = llvm::dyn_cast<TensorType>(inputs.front());  // 强制转换为 TensorType 类型

  TensorType output = llvm::dyn_cast<TensorType>(outputs.front());

  if (!input || !output || input.getElementType() != output.getElementType())

    return false;  //input 或 output 不是 tensor 类型，不合法；两者类型不同，不合法

  // The shape is required to match if both types are ranked.

  return !input.hasRank() || !output.hasRank() || input == output; // 有一个是未定型，合法；都是具体类型，且形状相同，合法

}

struct ToyInlinerInterface : public DialectInlinerInterface {

  ...

  /// Attempts to materialize a conversion for a type mismatch between a call

  /// from this dialect, and a callable region. This method should generate an

  /// operation that takes 'input' as the only operand, and produces a single

  /// result of 'resultType'. If a conversion can not be generated, nullptr

  /// should be returned.

  Operation *materializeCallConversion(OpBuilder &builder, Value input,

                                       Type resultType,

                                       Location conversionLoc) const final {

    return CastOp::create(builder, conversionLoc, resultType, input);

  }

};

toy.func @main() {

  %0 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>

  %1 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>

  %2 = toy.cast %1 : tensor<2x3xf64> to tensor<*xf64>

  %3 = toy.cast %0 : tensor<2x3xf64> to tensor<*xf64>

  %4 = toy.transpose(%2 : tensor<*xf64>) to tensor<*xf64>

  %5 = toy.transpose(%3 : tensor<*xf64>) to tensor<*xf64>

  %6 = toy.mul %4, %5 : tensor<*xf64>

  toy.print %6 : tensor<*xf64>

  toy.return

}

def ShapeInferenceOpInterface : OpInterface<"ShapeInference"> {  // 定义一个接口，继承 OpInterface 类

  let description = [{  // 函数描述

    Interface to access a registered method to infer the return types for an

    operation that can be used during type inference.

  }];

  let methods = [  // 定义继承接口之后需要实现的方法

    InterfaceMethod<"Infer and set the output shape for the current operation.",

                    "void", "inferShapes">  // <"方法描述", "返回值类型", "方法名称">

  ];

}

class ShapeInference : public OpInterface<ShapeInference> {

public:

  virtual void inferShapes() = 0;  // 纯虚函数

};

def MulOp : Toy_Op<"mul",

    [..., DeclareOpInterfaceMethods<ShapeInferenceOpInterface>]> {

  ...

}

/// Infer the output shape of the MulOp, this is required by the shape inference

/// interface.

void MulOp::inferShapes() { getResult().setType(getLhs().getType()); }

class ShapeInferencePass

    : public mlir::PassWrapper<ShapeInferencePass, OperationPass<FuncOp>> {

  void runOnOperation() override {

    FuncOp function = getOperation();

    ...

  }

};

std::unique_ptr<mlir::Pass> mlir::toy::createShapeInferencePass() {

  return std::make_unique<ShapeInferencePass>();

}

pm.addPass(mlir::createShapeInferencePass());

toy.func @main() {

  %0 = toy.constant dense<[[1.000000e+00, 2.000000e+00, 3.000000e+00], [4.000000e+00, 5.000000e+00, 6.000000e+00]]> : tensor<2x3xf64>

  %1 = toy.transpose(%0 : tensor<2x3xf64>) to tensor<3x2xf64>

  %2 = toy.mul %1, %1 : tensor<3x2xf64>

  toy.print %2 : tensor<3x2xf64>

  toy.return

}