MLIR-Toy-实践-3-Dialect转换

/// ToyCombine.cpp 定义pattern：transpose(transpose(x)) -> x
struct SimplifyRedundantTranspose : public mlir::OpRewritePattern<TransposeOp> {
  // ....
  mlir::LogicalResult
  matchAndRewrite(TransposeOp op,
                  mlir::PatternRewriter &rewriter) const override {
    // Look through the input of the current transpose.
    mlir::Value transposeInput = op.getOperand();
    TransposeOp transposeInputOp = transposeInput.getDefiningOp<TransposeOp>();

    // Input defined by another transpose? If not, no match.
    if (!transposeInputOp)
      return failure();

    // Otherwise, we have a redundant transpose. Use the rewriter.
    rewriter.replaceOp(op, {transposeInputOp.getOperand()});
    return success();
  }
};

// toy.cpp 创建pass
  mlir::PassManager pm(module.getContext());
  pm.addNestedPass<mlir::FuncOp>(mlir::createCanonicalizerPass());

// ToyCombine.cpp 注册pattarn
void TransposeOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.add<SimplifyRedundantTranspose>(context);
}

// 在ops.td中声明允许标准化操作
let hasCanonicalizer = 1;

class Pattern<
    dag sourcePattern, list<dag> resultPatterns,
    list<dag> additionalConstraints = [],
    dag benefitsAdded = (addBenefit 0)>;

// Reshape(Reshape(x)) = Reshape(x)
def ReshapeReshapeOptPattern : Pat<(ReshapeOp(ReshapeOp $arg)),
                                   (ReshapeOp $arg)>;

struct MyConversionPattern : public ConversionPattern {
  /// The `matchAndRewrite` hooks on ConversionPatterns take an additional
  /// `operands` parameter, containing the remapped operands of the original
  /// operation.
  virtual LogicalResult
  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                  ConversionPatternRewriter &rewriter) const;
};

  ConversionTarget target(getContext());

  // We define the specific operations, or dialects, that are legal targets for
  // this lowering. In our case, we are lowering to a combination of the
  // `Affine`, `MemRef` and `Standard` dialects.
  target.addLegalDialect<AffineDialect, memref::MemRefDialect,
                         StandardOpsDialect>();

  // We also define the Toy dialect as Illegal so that the conversion will fail
  // if any of these operations are *not* converted. Given that we actually want
  // a partial lowering, we explicitly mark the Toy operations that don't want
  // to lower, `toy.print`, as `legal`.
  target.addIllegalDialect<toy::ToyDialect>();
  target.addLegalOp<toy::PrintOp>();

  // 转换pattern定义
  struct TransposeOpLowering : public ConversionPattern {
    // ...
    LogicalResult
    matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                    ConversionPatternRewriter &rewriter)
    // ...
  ｝
  // 添加pattern到context
  RewritePatternSet patterns(&getContext());
  patterns.add<AddOpLowering, ConstantOpLowering, MulOpLowering,
               ReturnOpLowering, TransposeOpLowering>(&getContext());
  
  // 应用转换过程
  applyPartialConversion(getFunction(), target, std::move(patterns))

// 对OrOp添加Float转Int逻辑，并映射到Standard::OrOp
template <typename BinaryOp, typename LoweredBinaryOp>
struct BinaryOpLowering : public ConversionPattern {
  BinaryOpLowering(MLIRContext *ctx)
      : ConversionPattern(BinaryOp::getOperationName(), 1, ctx) {}

  LogicalResult
  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                  ConversionPatternRewriter &rewriter) const final {
    auto loc = op->getLoc();
    lowerOpToLoops(
        op, operands, rewriter,
        [loc, op](OpBuilder &builder, ValueRange memRefOperands,
              ValueRange loopIvs) {
          // ...

          // 对toy.or添加float转int逻辑，利用standard中的FPToUIOp
          auto opname = op->getName();
          if (opname.getStringRef().str() == "toy.or") {
            auto castLhs = builder.create<FPToUIOp>(loc, builder.getI64Type(), loadedLhs);
            auto castRhs = builder.create<FPToUIOp>(loc, builder.getI64Type(), loadedRhs);
            
            return builder.create<LoweredBinaryOp>(loc, castLhs, castRhs);
          }
          // ....
        });
    return success();
  }
};
// OrOp pattern定义
using OrOpLowering = BinaryOpLowering<toy::OrOp, OrOp>;

//
static void lowerOpToLoops(Operation *op, ArrayRef<Value> operands,
                           PatternRewriter &rewriter,
                           LoopIterationFn processIteration) {
  // ...
  buildAffineLoopNest(
      rewriter, loc, lowerBounds, tensorType.getShape(), steps,
      [&](OpBuilder &nestedBuilder, Location loc, ValueRange ivs) {
        Value valueToStore = processIteration(nestedBuilder, operands, ivs);
        // 将"toy.or" Op的结果从Int转换为Float
        if(nestedBuilder.getI64Type() == valueToStore.getType()) {
          valueToStore = nestedBuilder.create<UIToFPOp>(loc, nestedBuilder.getF64Type(), valueToStore);
        }
        nestedBuilder.create<AffineStoreOp>(loc, valueToStore, alloc, ivs);
      });
  // ...
}