MLIRの勉強 (4. 自作言語のMLIR変換とダンプ)

MLIRについて勉強している。

独自言語を作成し、その中間表現をMLIRを使って表現してみることに挑戦する。

ここで想定する流れは以下の通りだ。

1. 独自言語(MYSV : System Verilogオリジナル改造バージョン) を読み込んで、独自形式でASTを作成する
2. MYSVのASTをMLIRのツリーに変換する
3. MLIRのツリーをダンプする

次は、最適化の実装をするためにいくつか言語拡張をしていく。

Toy言語の場合は、Transpose(Transpose(A))の場合にはTransposeしない、という最適化を例に紹介がされている。

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

ただし自分の作っているMYSVはそういう構文を持っていないため、違う最適化の余地を実験してみたい。例えばこんな感じだ。

assign A = 0;
assign B = A;

この場合、最適化を行って、

assign B = 0;

という結論を導き出してくれれば嬉しい。こういうのはできないものだろうか。

このためには、まず自分のMYSVを拡張する必要がある。拡張するポイントは以下のようになるだろう。

• assign文のオペランドが変数を受け取るようにする(現在は定数のみ)

このためにASTの拡張を行おう。まず、式Exprが定数も変数も受け取れるようにする。

+++ b/mlir/examples/mysv/AST.cpp
@@ -26,6 +26,7 @@ class ASTDumper {
   void dump(AssignExprAST *assignExpr);
   void dump(ExprAST *expr);
   void dump(NumberExprAST *num);
+  void dump(VarExprAST *name);

Parserは、Assign文の=の次が何であるかに従ってParseの切り替えを行う。

--- a/mlir/examples/mysv/include/Parser.h
+++ b/mlir/examples/mysv/include/Parser.h
@@ -66,6 +66,40 @@ private:
     return result;
   }
 
+  /// identifierexpr
+  ///   ::= identifier
+  std::unique_ptr<ExprAST> parseIdentifierExpr() {
+    std::string name(lexer.getId());
+
+    auto loc = lexer.getLastLocation();
+    lexer.getNextToken(); // eat identifier.
+
+    return std::make_unique<VarExprAST>(std::move(loc), name);
+  }
+
+  /// primary
+  ///   ::= identifierexpr
+  ///   ::= numberexpr
+  std::unique_ptr<ExprAST> parsePrimary() {
+    switch (lexer.getCurToken()) {
+      default:
+        llvm::errs() << "unknown token '" << lexer.getCurToken()
+                     << "' when expecting an expression\n";
+        return nullptr;
+      case tok_identifier:
+        return parseIdentifierExpr();
+      case tok_number:
+        return parseNumberExpr();
+    }
+  }
+
+  std::unique_ptr<ExprAST> parseExpr() {
+    auto lhs = parsePrimary();
+    if (!lhs)
+      return nullptr;
+    return lhs;
+  }
+
   /// Parse a variable declaration, it starts with a `var` keyword followed by
   /// and identifier and an optional type (shape specification) before the
   /// initializer.
@@ -84,7 +118,7 @@ private:
 
     lexer.consume(Token('='));
 
-    auto expr = parseNumberExpr();
+    auto expr = parseExpr();
 
     lexer.consume(Token(';'));
     return std::make_unique<AssignExprAST>(std::move(loc), std::move(id), std::move(expr));

VarExprASTクラスを新たに作成する。ここではclassofメソッドを追加することでLLVMのTypeSwitchを使えることができるようになる。

+/// Expression class for variable literals.
+class VarExprAST : public ExprAST {
+  std::string name;
+
+ public:
+  VarExprAST(Location loc, llvm::StringRef name)
+      : ExprAST(Expr_Var, std::move(loc)), name(name) {}
+
+  llvm::StringRef getName() { return name; }
+
+  /// LLVM style RTTI
+  static bool classof(const ExprAST *c) { return c->getKind() == Expr_Var; }
+};
+

ExprAST用のmlirGenを新たに用意する。ASTの形式に従ってVarExprASTをベースにMLIRを生成するか、NumberExprASTをベースにMLIRを生成するか切り分ける。

+  /// Dispatch codegen for the right expression subclass using RTTI.
+  mlir::Value mlirGen(ExprAST &expr) {
+    switch (expr.getKind()) {
+      case mysv::ExprAST::Expr_Var:
+        return mlirGen(cast<VarExprAST>(expr));
+      case mysv::ExprAST::Expr_Num:
+        return mlirGen(cast<NumberExprAST>(expr));
+      default:
+        emitError(loc(expr.loc()))
+            << "MLIR codegen encountered an unhandled expr kind '"
+            << Twine(expr.getKind()) << "'";
+        return nullptr;
+    }
+  }
+
 
   /// Emit a literal/constant array. It will be emitted as a flattened array of
   /// data in an Attribute attached to a `mysv.constant` operation.
@@ -121,6 +136,19 @@ class MLIRGenImpl {
     return builder.create<ConstantOp>(loc(lit.loc()), elementType, lit.getValue());
   }
 
+
+  /// This is a reference to a variable in an expression. The variable is
+  /// expected to have been declared and so should have a value in the symbol
+  /// table, otherwise emit an error and return nullptr.
+  mlir::Value mlirGen(VarExprAST &expr) {
+    if (auto variable = symbolTable.lookup(expr.getName()))
+      return variable;
+
+    emitError(loc(expr.loc()), "error: unknown variable '")
+        << expr.getName() << "'";
+    return nullptr;
+  }
+

まずはASTでダンプしてみる。

 $ ./bin/mysv --emit=ast ../mlir/examples/mysv/test/assign_var.sv
  Module:
    assign A @../mlir/examples/mysv/test/assign_var.sv:1:1
      0 @../mlir/examples/mysv/test/assign_var.sv:1:12
    assign Hoge @../mlir/examples/mysv/test/assign_var.sv:2:1
      A @../mlir/examples/mysv/test/assign_var.sv:2:15

つぎはMLIRでダンプしてみる。

./bin/mysv --emit=mlir ../mlir/examples/mysv/test/assign_var.sv
loc("../mlir/examples/mysv/test/assign_var.sv":2:15): error: error: unknown variable 'A'
module {
  %0 = "mysv.constant"() {value = 3 : si64} : () -> i64
}

あれ？1文しか出力されないし、変数Aが無かったことになっている。これはなんだろう。