Skip to content

Instantly share code, notes, and snippets.

@airMeng
Last active December 25, 2023 06:57
Show Gist options
  • Save airMeng/f0a77592f508ac62a2f592c279d57a38 to your computer and use it in GitHub Desktop.
Save airMeng/f0a77592f508ac62a2f592c279d57a38 to your computer and use it in GitHub Desktop.
MLIR Hello.md

This is just my personal learning note in MLIR, just recording my questions, will or will not overlap with existing tutorial,

Frontend(parser)

8.30 update, parser is not necessary to understand MLIR and even makes it harder to understand MLIR itself.

I am not interested in any parser or frontend since it heavily depends on the source you choose and not general enough. However from my try, at least one thing about the parse is important. You have too options for get attributes(for example LHS of AddOp) of your customize Ops, one is like MLIR toy tutorial, you can define this method in your original IR(or called AST) and pass your method to parser then MLIR

167 /// Expression class for a binary operator.
168 class BinaryExprAST : public ExprAST {
...
171
172 public:
173   char getOp() { return op; }
174   ExprAST *getLHS() { return lhs.get(); }
175   ExprAST *getRHS() { return rhs.get(); }
...
184 };

Or you can get the attribute just from the header file as xxx.td

 845 def ONNXConvOp:ONNX_Op<"Conv",
 846   [NoSideEffect, DeclareOpInterfaceMethods<ShapeInferenceOpInterface>]> {
 847   let summary = "ONNX Conv operation";
 848   let description = [{
 849   "The convolution operator consumes an input tensor and a filter, and"
 850   "computes the output."
 851   }];
 852   let arguments = (ins AnyTypeOf<[TensorOf<[F16]>, TensorOf<[F32]>, TensorOf<[F64]>]>:$X,
 853     AnyTypeOf<[TensorOf<[F16]>, TensorOf<[F32]>, TensorOf<[F64]>]>:$W,
 854     AnyTypeOf<[TensorOf<[F16]>, TensorOf<[F32]>, TensorOf<[F64]>, NoneType]>:$B,
 ...
 874 }

then you can get the X directly.

auto inputOperand = ONNXConvOpAdaptor.X();

Define Op

Currently across the industry most operator in MLIR is generated by TableGen, so we don't care about traditional ways. A typical Op like this

 40 def AddOp : Hello_Op<"add", [NoSideEffect]> {
 41   let summary = "element-wise addition operation";
 42   let description = [{
 43     The "add" operation performs element-wise addition between two tensors.
 44     The shapes of the tensor operands are expected to match.
 45   }];
 46
 47   let arguments = (ins F64Tensor:$lhs, F64Tensor:$rhs);                                     // inputs
 48   let results = (outs F64Tensor);                                                           // outputs
 49
 50   // Indicate that the operation has a custom parser and printer method.
 51   let hasCustomAssemblyFormat = 1;                                                          // just as the comments, we need to define unique parser and printer.
 52
 53   // Allow building an AddOp with from the two input operands.
 54   let builders = [
 55     OpBuilder<(ins "mlir::Value":$lhs, "mlir::Value":$rhs)>                                 // declare builder here, so need to implement later.
 56   ];
 57 }

So PrintOp doesn't need build function but this is rare.

111 void hello::AddOp::build(mlir::OpBuilder &builder, mlir::OperationState &state,
112                   ::mlir::Value lhs, ::mlir::Value rhs) {
113   state.addTypes(UnrankedTensorType::get(builder.getF64Type()));
114   state.addOperands({lhs, rhs});
115 }

Conversion

A typical conversion between customized dialects and Affine is like

 67 static void lowerOpToLoops(Operation *op, ValueRange operands,
 68                            PatternRewriter &rewriter,
 69                            LoopIterationFn processIteration) {
 70   auto tensorType = (*op->result_type_begin()).cast<TensorType>();
 71   auto loc = op->getLoc();
 72
 73   // Insert an allocation and deallocation for the result of this operation.
 74   auto memRefType = convertTensorToMemRef(tensorType);
 75   auto alloc = insertAllocAndDealloc(memRefType, loc, rewriter);
 76
 77   // Create a nest of affine loops, with one loop per dimension of the shape.
 78   // The buildAffineLoopNest function takes a callback that is used to construct
 79   // the body of the innermost loop given a builder, a location and a range of
 80   // loop induction variables.
 81   SmallVector<int64_t, 4> lowerBounds(tensorType.getRank(), /*Value=*/0);
 82   SmallVector<int64_t, 4> steps(tensorType.getRank(), /*Value=*/1);
 83   buildAffineLoopNest(
 84       rewriter, loc, lowerBounds, tensorType.getShape(), steps,
 85       [&](OpBuilder &nestedBuilder, Location loc, ValueRange ivs) {
 86         // Call the processing function with the rewriter, the memref operands,
 87         // and the loop induction variables. This function will return the value
 88         // to store at the current index.
 89         Value valueToStore = processIteration(nestedBuilder, operands, ivs);
 90         nestedBuilder.create<AffineStoreOp>(loc, valueToStore, alloc, ivs);
 91       });
 92
 93   // Replace this operation with the generated alloc.
 94   rewriter.replaceOp(op, alloc);
 95 }
 ...
101 template <typename BinaryOp, typename LoweredBinaryOp>
102 struct BinaryOpLowering : public mlir::ConversionPattern {
103   BinaryOpLowering(MLIRContext *ctx)
104       : mlir::ConversionPattern(BinaryOp::getOperationName(), 1, ctx) {}
105
106   mlir::LogicalResult
107   matchAndRewrite(Operation *op, mlir::ArrayRef<mlir::Value> operands,
108                   mlir::ConversionPatternRewriter &rewriter) const final {
109     auto loc = op->getLoc();
110     lowerOpToLoops(op, operands, rewriter,
111                    [loc](OpBuilder &builder, ValueRange memRefOperands,
112                          ValueRange loopIvs) {
113                      // Generate an adaptor for the remapped operands of the
114                      // BinaryOp. This allows for using the nice named accessors
115                      // that are generated by the ODS.
116                      typename BinaryOp::Adaptor binaryAdaptor(memRefOperands);
117
118                      // Generate loads for the element of 'lhs' and 'rhs' at the
119                      // inner loop.
120                      auto loadedLhs = builder.create<AffineLoadOp>(
121                          loc, binaryAdaptor.lhs(), loopIvs);
122                      auto loadedRhs = builder.create<AffineLoadOp>(
123                          loc, binaryAdaptor.rhs(), loopIvs);
124
125                      // Create the binary operation performed on the loaded
126                      // values.
127                      return builder.create<LoweredBinaryOp>(loc, loadedLhs,
128                                                             loadedRhs);
129                    });
130     return mlir::success();
131   }
132 };

The key of conversion is the MatchAndRewrite, just as the name says, this function found your Op and create another Op to replace. This is basically a parse and construction case, except that construction should refer to complex MLIR API. Note lowerOpToLoops is the main construction case but you can easily find construction case for all MLIR Ops from official website.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment