[mlir][wasm] Expression parsing mechanism for Wasm importer

---------

Co-authored-by: Ferdinand Lemaire <[email protected]>
Co-authored-by: Jessica Paquette <[email protected]>

Author: 3 people

mlir/include/mlir/Target/Wasm/WasmBinaryEncoding.h

@@ -16,6 +16,15 @@
 #include <cstddef>
 namespace mlir {
 struct WasmBinaryEncoding {
+  /// Byte encodings for WASM instructions.
+  struct OpCode {
+    // Locals, globals, constants.
+    static constexpr std::byte constI32{0x41};
+    static constexpr std::byte constI64{0x42};
+    static constexpr std::byte constFP32{0x43};
+    static constexpr std::byte constFP64{0x44};
+  };
+
   /// Byte encodings of types in WASM binaries
   struct Type {
     static constexpr std::byte emptyBlockType{0x40};

mlir/lib/Target/Wasm/TranslateFromWasm.cpp

@@ -142,6 +142,8 @@ struct FunctionSymbolRefContainer : SymbolRefContainer {
 
 using ImportDesc = std::variant<TypeIdxRecord, TableType, LimitType, GlobalTypeRecord>;
 
+using parsed_inst_t = llvm::FailureOr<llvm::SmallVector<Value>>;
+
 struct WasmModuleSymbolTables {
   llvm::SmallVector<FunctionSymbolRefContainer> funcSymbols;
   llvm::SmallVector<GlobalSymbolRefContainer> globalSymbols;
@@ -173,6 +175,134 @@ struct WasmModuleSymbolTables {
     return getNewSymbolName("table_", id);
   }
 };
+
+class ParserHead;
+
+/// Wrapper around SmallVector to only allow access as push and pop on the
+/// stack. Makes sure that there are no "free accesses" on the stack to preserve
+/// its state.
+class ValueStack {
+private:
+  struct LabelLevel {
+    size_t stackIdx;
+    LabelLevelOpInterface levelOp;
+  };
+public:
+  bool empty() const { return values.empty(); }
+
+  size_t size() const { return values.size(); }
+
+  /// Pops values from the stack because they are being used in an operation.
+  /// @param operandTypes The list of expected types of the operation, used
+  ///   to know how many values to pop and check if the types match the
+  ///   expectation.
+  /// @param opLoc Location of the caller, used to report accurately the
+  /// location
+  ///   if an error occurs.
+  /// @return Failure or the vector of popped values.
+  llvm::FailureOr<llvm::SmallVector<Value>> popOperands(TypeRange operandTypes,
+                                                        Location *opLoc);
+
+  /// Push the results of an operation to the stack so they can be used in a
+  /// following operation.
+  /// @param results The list of results of the operation
+  /// @param opLoc Location of the caller, used to report accurately the
+  /// location
+  ///   if an error occurs.
+  LogicalResult pushResults(ValueRange results, Location *opLoc);
+
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+  /// A simple dump function for debugging.
+  /// Writes output to llvm::dbgs().
+  LLVM_DUMP_METHOD void dump() const;
+#endif
+
+private:
+  llvm::SmallVector<Value> values;
+};
+
+using local_val_t = TypedValue<wasmssa::LocalRefType>;
+
+class ExpressionParser {
+public:
+  using locals_t = llvm::SmallVector<local_val_t>;
+  ExpressionParser(ParserHead &parser, WasmModuleSymbolTables const &symbols,
+                   llvm::ArrayRef<local_val_t> initLocal)
+      : parser{parser}, symbols{symbols}, locals{initLocal} {}
+
+private:
+  template <std::byte opCode>
+  inline parsed_inst_t parseSpecificInstruction(OpBuilder &builder);
+
+  template <typename valueT>
+  parsed_inst_t
+  parseConstInst(OpBuilder &builder,
+                 std::enable_if_t<std::is_arithmetic_v<valueT>> * = nullptr);
+
+
+  /// This function generates a dispatch tree to associate an opcode with a
+  /// parser. Parsers are registered by specialising the
+  /// `parseSpecificInstruction` function for the op code to handle.
+  ///
+  /// The dispatcher is generated by recursively creating all possible patterns
+  /// for an opcode and calling the relevant parser on the leaf.
+  ///
+  /// @tparam patternBitSize is the first bit for which the pattern is not fixed
+  ///
+  /// @tparam highBitPattern is the fixed pattern that this instance handles for
+  /// the 8-patternBitSize bits
+  template <size_t patternBitSize = 0, std::byte highBitPattern = std::byte{0}>
+  inline parsed_inst_t dispatchToInstParser(std::byte opCode,
+                                            OpBuilder &builder) {
+    static_assert(patternBitSize <= 8,
+                  "PatternBitSize is outside of range of opcode space! "
+                  "(expected at most 8 bits)");
+    if constexpr (patternBitSize < 8) {
+      constexpr std::byte bitSelect{1 << (7 - patternBitSize)};
+      constexpr std::byte nextHighBitPatternStem = highBitPattern << 1;
+      constexpr size_t nextPatternBitSize = patternBitSize + 1;
+      if ((opCode & bitSelect) != std::byte{0})
+        return dispatchToInstParser < nextPatternBitSize,
+               nextHighBitPatternStem | std::byte{1} > (opCode, builder);
+      return dispatchToInstParser<nextPatternBitSize, nextHighBitPatternStem>(
+          opCode, builder);
+    } else {
+      return parseSpecificInstruction<highBitPattern>(builder);
+    }
+  }
+
+  struct ParseResultWithInfo {
+    llvm::SmallVector<Value> opResults;
+    std::byte endingByte;
+  };
+
+public:
+  template<std::byte ParseEndByte = WasmBinaryEncoding::endByte>
+  parsed_inst_t parse(OpBuilder &builder,
+                      UniqueByte<ParseEndByte> = {});
+
+  template <std::byte... ExpressionParseEnd>
+  llvm::FailureOr<ParseResultWithInfo>
+  parse(OpBuilder &builder,
+        ByteSequence<ExpressionParseEnd...> parsingEndFilters);
+
+  llvm::FailureOr<llvm::SmallVector<Value>>
+  popOperands(TypeRange operandTypes) {
+    return valueStack.popOperands(operandTypes, &currentOpLoc.value());
+  }
+
+  LogicalResult pushResults(ValueRange results) {
+    return valueStack.pushResults(results, &currentOpLoc.value());
+  }
+private:
+  std::optional<Location> currentOpLoc;
+  ParserHead &parser;
+  WasmModuleSymbolTables const &symbols;
+  locals_t locals;
+  ValueStack valueStack;
+  };
+
 class ParserHead {
 public:
   ParserHead(llvm::StringRef src, StringAttr name) : head{src}, locName{name} {}
@@ -382,6 +512,14 @@ class ParserHead {
              << static_cast<int>(*importType);
     }
   }
+
+  parsed_inst_t parseExpression(OpBuilder &builder,
+                                WasmModuleSymbolTables const &symbols,
+                                llvm::ArrayRef<local_val_t> locals = {}) {
+    auto eParser = ExpressionParser{*this, symbols, locals};
+    return eParser.parse(builder);
+  }
+
   bool end() const { return curHead().empty(); }
 
   ParserHead copy() const {
@@ -491,6 +629,201 @@ inline llvm::FailureOr<int64_t> ParserHead::parseI64() {
   return parseLiteral<int64_t>();
 }
 
+template <std::byte opCode>
+inline parsed_inst_t ExpressionParser::parseSpecificInstruction(OpBuilder &) {
+  return emitError(*currentOpLoc, "Unknown instruction opcode: ")
+         << static_cast<int>(opCode);
+}
+
+#if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
+void ValueStack::dump() const {
+  llvm::dbgs() << "================= Wasm ValueStack =======================\n";
+  llvm::dbgs() << "size: " << size() << "\n";
+  llvm::dbgs() << "<Top>"
+               << "\n";
+  // Stack is pushed to via push_back. Therefore the top of the stack is the
+  // end of the vector. Iterate in reverse so that the first thing we print
+  // is the top of the stack.
+  size_t stackSize = size();
+  for (size_t idx = 0 ; idx < stackSize ;) {
+    size_t actualIdx = stackSize - 1 - idx;
+    llvm::dbgs() << "  ";
+    values[actualIdx].dump();
+  }
+  llvm::dbgs() << "<Bottom>"
+               << "\n";
+  llvm::dbgs() << "=========================================================\n";
+}
+#endif
+
+parsed_inst_t ValueStack::popOperands(TypeRange operandTypes, Location *opLoc) {
+  LLVM_DEBUG(llvm::dbgs() << "Popping from ValueStack\n");
+  LLVM_DEBUG(llvm::dbgs() << "  Elements(s) to pop: " << operandTypes.size()
+                          << "\n");
+  LLVM_DEBUG(llvm::dbgs() << "  Current stack size: " << values.size() << "\n");
+  if (operandTypes.size() > values.size())
+    return emitError(*opLoc,
+                     "Stack doesn't contain enough values. Trying to get ")
+           << operandTypes.size() << " operands on a stack containing only "
+           << values.size() << " values.";
+  size_t stackIdxOffset = values.size() - operandTypes.size();
+  llvm::SmallVector<Value> res{};
+  res.reserve(operandTypes.size());
+  for (size_t i{0}; i < operandTypes.size(); ++i) {
+    Value operand = values[i + stackIdxOffset];
+    Type stackType = operand.getType();
+    if (stackType != operandTypes[i])
+      return emitError(*opLoc,
+                       "Invalid operand type on stack. Expecting ")
+             << operandTypes[i] << ", value on stack is of type " << stackType
+             << ".";
+    LLVM_DEBUG(llvm::dbgs() << "    POP: " << operand << "\n");
+    res.push_back(operand);
+  }
+  values.resize(values.size() - operandTypes.size());
+  LLVM_DEBUG(llvm::dbgs() << "  Updated stack size: " << values.size() << "\n");
+  return res;
+}
+
+LogicalResult ValueStack::pushResults(ValueRange results, Location *opLoc) {
+  LLVM_DEBUG(llvm::dbgs() << "Pushing to ValueStack\n");
+  LLVM_DEBUG(llvm::dbgs() << "  Elements(s) to push: " << results.size()
+                          << "\n");
+  LLVM_DEBUG(llvm::dbgs() << "  Current stack size: " << values.size() << "\n");
+  for (auto val : results) {
+    if (!isWasmValueType(val.getType()))
+      return emitError(*opLoc, "Invalid value type on stack: ")
+             << val.getType();
+    LLVM_DEBUG(llvm::dbgs() << "    PUSH: " << val << "\n");
+    values.push_back(val);
+  }
+
+  LLVM_DEBUG(llvm::dbgs() << "  Updated stack size: " << values.size() << "\n");
+  return success();
+}
+
+template<std::byte EndParseByte>
+parsed_inst_t ExpressionParser::parse(OpBuilder &builder, UniqueByte<EndParseByte> endByte) {
+  auto res = parse(builder, ByteSequence<EndParseByte>{});
+  if (failed(res))
+    return failure();
+  return res->opResults;
+}
+
+template <std::byte... ExpressionParseEnd>
+llvm::FailureOr<ExpressionParser::ParseResultWithInfo>
+ExpressionParser::parse(OpBuilder &builder,
+                        ByteSequence<ExpressionParseEnd...> parsingEndFilters) {
+  llvm::SmallVector<Value> res;
+  for (;;) {
+    currentOpLoc = parser.getLocation();
+    auto opCode = parser.consumeByte();
+    if (failed(opCode))
+      return failure();
+    if (isValueOneOf(*opCode, parsingEndFilters))
+      return {{res, *opCode}};
+    parsed_inst_t resParsed;
+    resParsed = dispatchToInstParser(*opCode, builder);
+    if (failed(resParsed))
+      return failure();
+    std::swap(res, *resParsed);
+    if (failed(pushResults(res)))
+      return failure();
+  }
+}
+
+
+template <typename T>
+inline Type buildLiteralType(OpBuilder &);
+
+template <>
+inline Type buildLiteralType<int32_t>(OpBuilder &builder) {
+  return builder.getI32Type();
+}
+
+template <>
+inline Type buildLiteralType<int64_t>(OpBuilder &builder) {
+  return builder.getI64Type();
+}
+
+template <>
+inline Type buildLiteralType<uint32_t>(OpBuilder &builder) {
+  return builder.getI32Type();
+}
+
+template <>
+inline Type buildLiteralType<uint64_t>(OpBuilder &builder) {
+  return builder.getI64Type();
+}
+
+template <>
+inline Type buildLiteralType<float>(OpBuilder &builder) {
+  return builder.getF32Type();
+}
+
+template <>
+inline Type buildLiteralType<double>(OpBuilder &builder) {
+  return builder.getF64Type();
+}
+
+template<typename ValT, typename E = std::enable_if_t<std::is_arithmetic_v<ValT>>>
+struct AttrHolder;
+
+template <typename ValT>
+struct AttrHolder<ValT, std::enable_if_t<std::is_integral_v<ValT>>> {
+  using type = IntegerAttr;
+};
+
+template <typename ValT>
+struct AttrHolder<ValT, std::enable_if_t<std::is_floating_point_v<ValT>>> {
+  using type = FloatAttr;
+};
+
+template<typename ValT>
+using attr_holder_t = typename AttrHolder<ValT>::type;
+
+template <typename ValT,
+          typename EnableT = std::enable_if_t<std::is_arithmetic_v<ValT>>>
+attr_holder_t<ValT> buildLiteralAttr(OpBuilder &builder, ValT val) {
+  return attr_holder_t<ValT>::get(buildLiteralType<ValT>(builder), val);
+}
+
+template <typename valueT>
+parsed_inst_t ExpressionParser::parseConstInst(
+    OpBuilder &builder, std::enable_if_t<std::is_arithmetic_v<valueT>> *) {
+  auto parsedConstant = parser.parseLiteral<valueT>();
+  if (failed(parsedConstant))
+    return failure();
+  auto constOp = builder.create<ConstOp>(
+      *currentOpLoc, buildLiteralAttr<valueT>(builder, *parsedConstant));
+  return {{constOp.getResult()}};
+}
+
+template <>
+inline parsed_inst_t ExpressionParser::parseSpecificInstruction<
+    WasmBinaryEncoding::OpCode::constI32>(OpBuilder &builder) {
+  return parseConstInst<int32_t>(builder);
+}
+
+template <>
+inline parsed_inst_t ExpressionParser::parseSpecificInstruction<
+    WasmBinaryEncoding::OpCode::constI64>(OpBuilder &builder) {
+  return parseConstInst<int64_t>(builder);
+}
+
+template <>
+inline parsed_inst_t ExpressionParser::parseSpecificInstruction<
+    WasmBinaryEncoding::OpCode::constFP32>(OpBuilder &builder) {
+  return parseConstInst<float>(builder);
+}
+
+template <>
+inline parsed_inst_t ExpressionParser::parseSpecificInstruction<
+    WasmBinaryEncoding::OpCode::constFP64>(OpBuilder &builder) {
+  return parseConstInst<double>(builder);
+}
+
+
 class WasmBinaryParser {
 private:
   struct SectionRegistry {