[Compilation Hints] Start Compilation Hints by adding Call support in the text format (#7587)

Left as TODOs are other instructions, the binary format, and
function-level annotations.

Author: kripken

src/parser/contexts.h

@@ -1257,9 +1257,47 @@ struct ParseImplicitTypeDefsCtx : TypeParserCtx<ParseImplicitTypeDefsCtx> {
   }
 };
 
+struct AnnotationParserCtx {
+  // Return the inline hint for a call instruction, if there is one.
+  std::optional<std::uint8_t>
+  getInlineHint(const std::vector<Annotation>& annotations) {
+    // Find and apply (the last) inline hint.
+    const Annotation* hint = nullptr;
+    for (auto& a : annotations) {
+      if (a.kind == Annotations::InlineHint) {
+        hint = &a;
+      }
+    }
+    if (!hint) {
+      return std::nullopt;
+    }
+
+    Lexer lexer(hint->contents);
+    if (lexer.empty()) {
+      std::cerr << "warning: empty InlineHint\n";
+      return std::nullopt;
+    }
+
+    auto str = lexer.takeString();
+    if (!str || str->size() != 1) {
+      std::cerr << "warning: invalid InlineHint string\n";
+      return std::nullopt;
+    }
+
+    uint8_t value = (*str)[0];
+    if (value > 127) {
+      std::cerr << "warning: invalid InlineHint value\n";
+      return std::nullopt;
+    }
+
+    return value;
+  }
+};
+
 // Phase 4: Parse and set the types of module elements.
 struct ParseModuleTypesCtx : TypeParserCtx<ParseModuleTypesCtx>,
-                             NullInstrParserCtx {
+                             NullInstrParserCtx,
+                             AnnotationParserCtx {
   // In this phase we have constructed all the types, so we can materialize and
   // validate them when they are used.
 
@@ -1367,6 +1405,13 @@ struct ParseModuleTypesCtx : TypeParserCtx<ParseModuleTypesCtx>,
         Builder::addVar(f.get(), l.name, l.type);
       }
     }
+    // TODO: Add function-level annotations (stored using the nullptr key, as
+    // they are tied to no instruction in particular), but this should wait on
+    // figuring out
+    // https://github.com/WebAssembly/tool-conventions/issues/251
+    // if (auto inline_ = getInlineHint(annotations)) {
+    //   f->codeAnnotations[nullptr].inline_ = inline_;
+    // }
     return Ok{};
   }
 
@@ -1427,7 +1472,7 @@ struct ParseModuleTypesCtx : TypeParserCtx<ParseModuleTypesCtx>,
 };
 
 // Phase 5: Parse module element definitions, including instructions.
-struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
+struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx>, AnnotationParserCtx {
   using GlobalTypeT = Ok;
   using TableTypeT = Ok;
   using TypeUseT = HeapType;
@@ -2329,7 +2374,8 @@ struct ParseDefsCtx : TypeParserCtx<ParseDefsCtx> {
                     const std::vector<Annotation>& annotations,
                     Name func,
                     bool isReturn) {
-    return withLoc(pos, irBuilder.makeCall(func, isReturn));
+    auto inline_ = getInlineHint(annotations);
+    return withLoc(pos, irBuilder.makeCall(func, isReturn, inline_));
   }
 
   Result<> makeCallIndirect(Index pos,

src/passes/Print.cpp

@@ -19,6 +19,7 @@
 //
 
 #include <algorithm>
+#include <fstream>
 
 #include <ir/iteration.h>
 #include <ir/module-utils.h>
@@ -275,6 +276,10 @@ struct PrintSExpression : public UnifiedExpressionVisitor<PrintSExpression> {
   // Prints debug info and code annotations.
   void printMetadata(Expression* curr);
 
+  // Print code annotations for an expression. If the expression is nullptr,
+  // prints for the current function.
+  void printCodeAnnotations(Expression* curr);
+
   void printExpressionContents(Expression* curr);
 
   void visit(Expression* curr) {
@@ -2642,18 +2647,8 @@ void PrintSExpression::printMetadata(Expression* curr) {
       }
     }
 
-    // Show a code annotation, if there is one.
-    if (auto iter = currFunction->codeAnnotations.find(curr);
-        iter != currFunction->codeAnnotations.end()) {
-      auto& annotation = iter->second;
-      if (annotation.branchLikely) {
-        Colors::grey(o);
-        o << "(@" << Annotations::BranchHint << " \"\\0"
-          << (*annotation.branchLikely ? "1" : "0") << "\")\n";
-        restoreNormalColor(o);
-        doIndent(o, indent);
-      }
-    }
+    // Show code annotations.
+    printCodeAnnotations(curr);
   }
 }
 
@@ -2670,6 +2665,31 @@ void PrintSExpression::printDebugDelimiterLocation(Expression* curr, Index i) {
   }
 }
 
+void PrintSExpression::printCodeAnnotations(Expression* curr) {
+  if (auto iter = currFunction->codeAnnotations.find(curr);
+      iter != currFunction->codeAnnotations.end()) {
+    auto& annotation = iter->second;
+    if (annotation.branchLikely) {
+      Colors::grey(o);
+      o << "(@" << Annotations::BranchHint << " \"\\0"
+        << (*annotation.branchLikely ? "1" : "0") << "\")\n";
+      restoreNormalColor(o);
+      doIndent(o, indent);
+    }
+    if (annotation.inline_) {
+      Colors::grey(o);
+      std::ofstream saved;
+      saved.copyfmt(std::cout);
+      o << "(@" << Annotations::InlineHint << " \"\\" << std::hex
+        << std::setfill('0') << std::setw(2) << int(*annotation.inline_)
+        << "\")\n";
+      std::cout.copyfmt(saved);
+      restoreNormalColor(o);
+      doIndent(o, indent);
+    }
+  }
+}
+
 void PrintSExpression::printExpressionContents(Expression* curr) {
   PrintExpressionContents(*this).visit(curr);
 }
@@ -3125,6 +3145,9 @@ void PrintSExpression::visitDefinedFunction(Function* curr) {
   if (currFunction->prologLocation) {
     printDebugLocation(*currFunction->prologLocation);
   }
+  // TODO: print code annotations in the right place, depending on
+  // https://github.com/WebAssembly/tool-conventions/issues/251
+  // printCodeAnnotations(nullptr);
   handleSignature(curr, true);
   incIndent();
   for (size_t i = curr->getVarIndexBase(); i < curr->getNumLocals(); i++) {

src/wasm-annotations.h

@@ -26,6 +26,7 @@
 namespace wasm::Annotations {
 
 extern const Name BranchHint;
+extern const Name InlineHint;
 
 } // namespace wasm::Annotations
 

src/wasm-ir-builder.h

@@ -124,7 +124,9 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
                      std::optional<bool> likely = std::nullopt);
   Result<> makeSwitch(const std::vector<Index>& labels, Index defaultLabel);
   // Unlike Builder::makeCall, this assumes the function already exists.
-  Result<> makeCall(Name func, bool isReturn);
+  Result<> makeCall(Name func,
+                    bool isReturn,
+                    std::optional<std::uint8_t> inline_ = std::nullopt);
   Result<> makeCallIndirect(Name table, HeapType type, bool isReturn);
   Result<> makeLocalGet(Index local);
   Result<> makeLocalSet(Index local);
@@ -701,6 +703,9 @@ class IRBuilder : public UnifiedExpressionVisitor<IRBuilder, Result<>> {
   // Add a branch hint, if |likely| is present.
   void addBranchHint(Expression* expr, std::optional<bool> likely);
 
+  // Add an inlining hint, if |inline_| is present.
+  void addInlineHint(Expression* expr, std::optional<std::uint8_t> inline_);
+
   void dump();
 };
 

src/wasm.h

@@ -2191,14 +2191,21 @@ class Function : public Importable {
     delimiterLocations;
   BinaryLocations::FunctionLocations funcLocation;
 
-  // Code annotations. As with debug info, we do not store these on Expressions
-  // themselves, as we assume most instances are unannotated, and do not want to
+  // Code annotations for VMs. As with debug info, we do not store these on
+  // Expressions as we assume most instances are unannotated, and do not want to
   // add constant memory overhead.
   struct CodeAnnotation {
-    // Branch hinting proposal: Whether the branch is likely, or unlikely.
+    // Branch Hinting proposal: Whether the branch is likely, or unlikely.
     std::optional<bool> branchLikely;
+
+    // Compilation Hints proposal.
+    static const uint8_t NeverInline = 0;
+    static const uint8_t AlwaysInline = 127;
+    std::optional<uint8_t> inline_;
   };
 
+  // Function-level annotations are implemented with a key of nullptr, matching
+  // the 0 byte offset in the spec.
   std::unordered_map<Expression*, CodeAnnotation> codeAnnotations;
 
   // The effects for this function, if they have been computed. We use a shared
@@ -2208,8 +2215,8 @@ class Function : public Importable {
   // See addsEffects() in pass.h for more details.
   std::shared_ptr<EffectAnalyzer> effects;
 
-  // Inlining metadata: whether to disallow full and/or partial inlining (for
-  // details on what those mean, see Inlining.cpp).
+  // Inlining metadata: whether to disallow full and/or partial inlining. This
+  // is a toolchain-level hint. For more details, see Inlining.cpp.
   bool noFullInline = false;
   bool noPartialInline = false;
 

src/wasm/wasm-ir-builder.cpp

@@ -1441,13 +1441,18 @@ Result<> IRBuilder::makeSwitch(const std::vector<Index>& labels,
   return Ok{};
 }
 
-Result<> IRBuilder::makeCall(Name func, bool isReturn) {
+Result<> IRBuilder::makeCall(Name func,
+                             bool isReturn,
+                             std::optional<std::uint8_t> inline_) {
   auto sig = wasm.getFunction(func)->getSig();
   Call curr(wasm.allocator);
   curr.target = func;
   curr.operands.resize(sig.params.size());
   CHECK_ERR(visitCall(&curr));
-  push(builder.makeCall(curr.target, curr.operands, sig.results, isReturn));
+  auto* call =
+    builder.makeCall(curr.target, curr.operands, sig.results, isReturn);
+  push(call);
+  addInlineHint(call, inline_);
   return Ok{};
 }
 
@@ -2541,4 +2546,13 @@ void IRBuilder::addBranchHint(Expression* expr, std::optional<bool> likely) {
   }
 }
 
+void IRBuilder::addInlineHint(Expression* expr,
+                              std::optional<uint8_t> inline_) {
+  if (inline_) {
+    // Branches are only possible inside functions.
+    assert(func);
+    func->codeAnnotations[expr].inline_ = inline_;
+  }
+}
+
 } // namespace wasm

src/wasm/wasm.cpp

@@ -67,6 +67,7 @@ const char* CustomDescriptorsFeature = "custom-descriptors";
 namespace Annotations {
 
 const Name BranchHint = "metadata.code.branch_hint";
+const Name InlineHint = "metadata.code.inline";
 
 } // namespace Annotations
 

test/lit/wat-annotations.wast renamed to test/lit/branch-hinting.wast

@@ -0,0 +1,36 @@
+;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
+
+;; RUN: wasm-opt -all             %s -S -o - | filecheck %s
+
+;; TODO: wasm-opt -all --roundtrip %s -S -o - | filecheck %s --check-prefix=RTRIP
+
+(module
+  ;; CHECK:      (type $0 (func))
+
+  ;; CHECK:      (func $func (type $0)
+  ;; CHECK-NEXT:  (@metadata.code.inline "\00")
+  ;; CHECK-NEXT:  (call $func)
+  ;; CHECK-NEXT:  (@metadata.code.inline "\01")
+  ;; CHECK-NEXT:  (call $func)
+  ;; CHECK-NEXT:  (@metadata.code.inline "\7e")
+  ;; CHECK-NEXT:  (call $func)
+  ;; CHECK-NEXT:  (@metadata.code.inline "\7f")
+  ;; CHECK-NEXT:  (call $func)
+  ;; CHECK-NEXT:  (call $func)
+  ;; CHECK-NEXT: )
+  (func $func
+    (@metadata.code.inline "\00")
+    (call $func)
+    (@metadata.code.inline "\01")
+    (call $func)
+    (@metadata.code.inline "\7e")
+    (call $func)
+    (@metadata.code.inline "\7f")
+    (call $func)
+    ;; Unannotated
+    (call $func)
+  )
+
+  ;; TODO: test function annotations, after
+  ;; https://github.com/WebAssembly/tool-conventions/issues/251
+)