[MLIR][WASM] Control flow, conversion and comparison in Wasm importer (#154674)

This is the following of PR #154452.

It extend Wasm binary to Wasm SSA importer with support of control flow
operations, comparison operations and conversion operations.

---------

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

Author: 4 people

mlir/include/mlir/Dialect/WasmSSA/IR/WasmSSAOps.td

@@ -43,9 +43,41 @@ class WasmSSA_BlockLikeOp<string mnemonic, string summaryStr> :
   let assemblyFormat = "(`(`$inputs^`)` `:` type($inputs))? attr-dict  `:` $body `>` $target";
 }
 
-def WasmSSA_BlockOp : WasmSSA_BlockLikeOp<"block", "Create a nesting level"> {}
+def WasmSSA_BlockOp : WasmSSA_BlockLikeOp<
+    "block",
+    "Create a nesting level with a label at its exit."> {
+  let description = [{
+  Defines a Wasm block, creating a new nested scope.
+  A block contains a body region and an optional list of input values.
+  Control can enter the block and later branch out to the block target.
+  Example:
+
+  ```mlir
+
+  wasmssa.block {
+
+    // instructions
+
+  } > ^successor
+  }];
+}
+
+def WasmSSA_LoopOp : WasmSSA_BlockLikeOp<
+    "loop",
+    "Create a nesting level that define its entry as jump target."> {
+  let description = [{
+  Represents a Wasm loop construct. This defines a nesting level with
+  a label at the entry of the region.
 
-def WasmSSA_LoopOp : WasmSSA_BlockLikeOp<"loop", "Create a nesting level similar to Block Op, except that it has itself as a successor."> {}
+  Example:
+
+  ```mlir
+
+  wasmssa.loop {
+
+  } > ^successor
+  }];
+}
 
 def WasmSSA_BlockReturnOp : WasmSSA_Op<"block_return", [Terminator,
     DeclareOpInterfaceMethods<LabelBranchingOpInterface>]> {
@@ -55,9 +87,16 @@ def WasmSSA_BlockReturnOp : WasmSSA_Op<"block_return", [Terminator,
     ::mlir::Block* getTarget();
   }];
   let description = [{
-     Marks a return from the current block.
+    Escape from the current nesting level and return the control flow to its successor.
+    Optionally, mark the arguments that should be transfered to the successor block.
 
-     Example:
+    This shouldn't be confused with branch operations that targets the label defined
+    by the nesting level operation.
+
+    For instance, a `wasmssa.block_return` in a loop will give back control to the
+    successor of the loop, where a `branch` targeting the loop will flow back to the entry block of the loop.
+
+    Example:
 
     ```mlir
       wasmssa.block_return
@@ -127,12 +166,18 @@ def WasmSSA_FuncOp : WasmSSA_Op<"func", [
     - Arguments of the entry block of type `!wasm<local T>`, with T the corresponding type
      in the function type.
 
+    By default, `wasmssa.func` have nested visibility. Functions exported by the module
+    are marked with the exported attribute. This gives them public visibility.
+
      Example:
 
      ```mlir
-     // A simple function with no arguments that returns a float32
+     // Internal function with no arguments that returns a float32
      wasmssa.func @my_f32_func() -> f32
 
+     // Exported function with no arguments that returns a float32
+     wasmssa.func exported @my_f32_func() -> f32
+
      // A function that takes a local ref argument
      wasmssa.func @i64_wrap(%a: !wasmssa<local ref to i64>) -> i32
      ```
@@ -141,7 +186,7 @@ def WasmSSA_FuncOp : WasmSSA_Op<"func", [
                      WasmSSA_FuncTypeAttr: $functionType,
                      OptionalAttr<DictArrayAttr>:$arg_attrs,
                      OptionalAttr<DictArrayAttr>:$res_attrs,
-                     DefaultValuedAttr<StrAttr, "\"nested\"">:$sym_visibility);
+                     UnitAttr: $exported);
   let regions = (region AnyRegion: $body);
   let extraClassDeclaration = [{
 
@@ -162,6 +207,12 @@ def WasmSSA_FuncOp : WasmSSA_Op<"func", [
 
     /// Returns the result types of this function.
     ArrayRef<Type> getResultTypes() { return getFunctionType().getResults(); }
+
+    ::mlir::SymbolTable::Visibility getVisibility() {
+      return getExported() ?
+        ::mlir::SymbolTable::Visibility::Public :
+        ::mlir::SymbolTable::Visibility::Nested;
+    };
   }];
 
   let builders = [
@@ -207,8 +258,7 @@ def WasmSSA_FuncImportOp : WasmSSA_Op<"import_func", [
                      StrAttr: $importName,
                      WasmSSA_FuncTypeAttr: $type,
                      OptionalAttr<DictArrayAttr>:$arg_attrs,
-                     OptionalAttr<DictArrayAttr>:$res_attrs,
-                     OptionalAttr<StrAttr>:$sym_visibility);
+                     OptionalAttr<DictArrayAttr>:$res_attrs);
   let extraClassDeclaration = [{
     bool isDeclaration() const { return true; }
 
@@ -221,6 +271,10 @@ def WasmSSA_FuncImportOp : WasmSSA_Op<"import_func", [
     ::llvm::ArrayRef<Type> getResultTypes() {
       return getType().getResults();
     }
+
+    ::mlir::SymbolTable::Visibility getVisibility() {
+      return ::mlir::SymbolTable::Visibility::Nested;
+    };
   }];
   let builders = [
     OpBuilder<(ins "StringRef":$symbol,
@@ -238,30 +292,41 @@ def WasmSSA_GlobalOp : WasmSSA_Op<"global", [
   let arguments = (ins SymbolNameAttr: $sym_name,
                      WasmSSA_ValTypeAttr: $type,
                      UnitAttr: $isMutable,
-                     OptionalAttr<StrAttr>:$sym_visibility);
+                     UnitAttr: $exported);
   let description = [{
     WebAssembly global variable.
     Body contains the initialization instructions for the variable value.
     The body must contain only instructions considered `const` in a webassembly context,
     such as `wasmssa.const` or `global.get`.
 
+    By default, `wasmssa.global` have nested visibility. Global exported by the module
+    are marked with the exported attribute. This gives them public visibility.
+
     Example:
 
     ```mlir
-    // Define a global_var, a mutable i32 global variable equal to 10.
-    wasmssa.global @global_var i32 mutable nested : {
+    // Define module_global_var, an internal mutable i32 global variable equal to 10.
+    wasmssa.global @module_global_var i32 mutable : {
           %[[VAL_0:.*]] = wasmssa.const 10 : i32
           wasmssa.return %[[VAL_0]] : i32
     }
+
+    // Define global_var, an exported constant i32 global variable equal to 42.
+    wasmssa.global @global_var i32 : {
+          %[[VAL_0:.*]] = wasmssa.const 42 : i32
+          wasmssa.return %[[VAL_0]] : i32
+    }
     ```
   }];
   let regions = (region AnyRegion: $initializer);
 
-  let builders = [
-    OpBuilder<(ins "StringRef":$symbol,
-                   "Type": $type,
-                   "bool": $isMutable)>
-  ];
+  let extraClassDeclaration = [{
+    ::mlir::SymbolTable::Visibility getVisibility() {
+      return getExported() ?
+        ::mlir::SymbolTable::Visibility::Public :
+        ::mlir::SymbolTable::Visibility::Nested;
+    };
+  }];
   let hasCustomAssemblyFormat = 1;
 }
 
@@ -283,18 +348,14 @@ def WasmSSA_GlobalImportOp : WasmSSA_Op<"import_global", [
                      StrAttr: $moduleName,
                      StrAttr: $importName,
                      WasmSSA_ValTypeAttr: $type,
-                     UnitAttr: $isMutable,
-                     OptionalAttr<StrAttr>:$sym_visibility);
+                     UnitAttr: $isMutable);
   let extraClassDeclaration = [{
     bool isDeclaration() const { return true; }
+
+    ::mlir::SymbolTable::Visibility getVisibility() {
+      return ::mlir::SymbolTable::Visibility::Nested;
+    };
   }];
-  let builders = [
-    OpBuilder<(ins "StringRef":$symbol,
-                   "StringRef":$moduleName,
-                   "StringRef":$importName,
-                   "Type": $type,
-                   "bool": $isMutable)>
-  ];
   let hasCustomAssemblyFormat = 1;
 }
 
@@ -442,23 +503,33 @@ def WasmSSA_MemOp : WasmSSA_Op<"memory", [Symbol]> {
     Define a memory to be used by the program.
     Multiple memories can be defined in the same module.
 
+    By default, `wasmssa.memory` have nested visibility. Memory exported by
+    the module are marked with the exported attribute. This gives them public
+    visibility.
+
      Example:
 
      ```mlir
-     // Define the `mem_0` memory with defined bounds of 0 -> 65536
+     // Define the `mem_0` (internal)  memory with defined size bounds of [0:65536]
      wasmssa.memory @mem_0 !wasmssa<limit[0:65536]>
+
+     // Define the `mem_1` exported  memory with minimal size of 512
+     wasmssa.memory exported @mem_1 !wasmssa<limit[512:]>
      ```
     }];
   let arguments = (ins SymbolNameAttr: $sym_name,
                      WasmSSA_LimitTypeAttr: $limits,
-                     OptionalAttr<StrAttr>:$sym_visibility);
-  let builders = [
-    OpBuilder<(ins
-    "::llvm::StringRef":$symbol,
-    "wasmssa::LimitType":$limit)>
-  ];
+                     UnitAttr: $exported);
 
-  let assemblyFormat = "$sym_name custom<WasmVisibility>($sym_visibility) $limits attr-dict";
+  let extraClassDeclaration = [{
+    ::mlir::SymbolTable::Visibility getVisibility() {
+      return getExported() ?
+        ::mlir::SymbolTable::Visibility::Public :
+        ::mlir::SymbolTable::Visibility::Nested;
+    };
+  }];
+
+  let assemblyFormat = "(`exported` $exported^)? $sym_name $limits attr-dict";
 }
 
 def WasmSSA_MemImportOp : WasmSSA_Op<"import_mem", [Symbol, ImportOpInterface]> {
@@ -476,28 +547,29 @@ def WasmSSA_MemImportOp : WasmSSA_Op<"import_mem", [Symbol, ImportOpInterface]>
   let arguments = (ins SymbolNameAttr: $sym_name,
                      StrAttr: $moduleName,
                      StrAttr: $importName,
-                     WasmSSA_LimitTypeAttr: $limits,
-                     OptionalAttr<StrAttr>:$sym_visibility);
+                     WasmSSA_LimitTypeAttr: $limits);
   let extraClassDeclaration = [{
-     bool isDeclaration() const { return true; }
+      bool isDeclaration() const { return true; }
+      ::mlir::SymbolTable::Visibility getVisibility() {
+       return ::mlir::SymbolTable::Visibility::Nested;
+      };
    }];
-  let builders = [OpBuilder<(ins
-    "::llvm::StringRef":$symbol,
-    "::llvm::StringRef":$moduleName,
-    "::llvm::StringRef":$importName,
-    "wasmssa::LimitType":$limits)>];
   let assemblyFormat = "$importName `from` $moduleName `as` $sym_name attr-dict";
 }
 
 def WasmSSA_TableOp : WasmSSA_Op<"table", [Symbol]> {
   let summary= "WebAssembly table value";
   let arguments = (ins SymbolNameAttr: $sym_name,
                      WasmSSA_TableTypeAttr: $type,
-                     OptionalAttr<StrAttr>:$sym_visibility);
-  let builders = [OpBuilder<(ins
-      "::llvm::StringRef":$symbol,
-      "wasmssa::TableType":$type)>];
-  let assemblyFormat = "$sym_name custom<WasmVisibility>($sym_visibility) $type attr-dict";
+                     UnitAttr: $exported);
+  let extraClassDeclaration = [{
+    ::mlir::SymbolTable::Visibility getVisibility() {
+      return getExported() ?
+        ::mlir::SymbolTable::Visibility::Public :
+        ::mlir::SymbolTable::Visibility::Nested;
+    };
+  }];
+  let assemblyFormat = "(`exported` $exported^)? $sym_name $type attr-dict";
 }
 
 def WasmSSA_TableImportOp : WasmSSA_Op<"import_table", [Symbol, ImportOpInterface]> {
@@ -515,17 +587,14 @@ def WasmSSA_TableImportOp : WasmSSA_Op<"import_table", [Symbol, ImportOpInterfac
   let arguments = (ins SymbolNameAttr: $sym_name,
                      StrAttr: $moduleName,
                      StrAttr: $importName,
-                     WasmSSA_TableTypeAttr: $type,
-                     OptionalAttr<StrAttr>:$sym_visibility);
+                     WasmSSA_TableTypeAttr: $type);
   let extraClassDeclaration = [{
     bool isDeclaration() const { return true; }
+    ::mlir::SymbolTable::Visibility getVisibility() {
+       return ::mlir::SymbolTable::Visibility::Nested;
+    };
   }];
   let assemblyFormat = "$importName `from` $moduleName `as` $sym_name attr-dict";
-  let builders = [OpBuilder<(ins
-      "::llvm::StringRef":$symbol,
-      "::llvm::StringRef":$moduleName,
-      "::llvm::StringRef":$importName,
-      "wasmssa::TableType":$type)>];
 }
 
 def WasmSSA_ReturnOp : WasmSSA_Op<"return", [Terminator]> {

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

@@ -19,6 +19,14 @@ namespace mlir {
 struct WasmBinaryEncoding {
   /// Byte encodings for Wasm instructions.
   struct OpCode {
+    // Control instructions.
+    static constexpr std::byte block{0x02};
+    static constexpr std::byte loop{0x03};
+    static constexpr std::byte ifOpCode{0x04};
+    static constexpr std::byte elseOpCode{0x05};
+    static constexpr std::byte branchIf{0x0D};
+    static constexpr std::byte call{0x10};
+
     // Locals, globals, constants.
     static constexpr std::byte localGet{0x20};
     static constexpr std::byte localSet{0x21};
@@ -29,6 +37,42 @@ struct WasmBinaryEncoding {
     static constexpr std::byte constFP32{0x43};
     static constexpr std::byte constFP64{0x44};
 
+    // Comparisons.
+    static constexpr std::byte eqzI32{0x45};
+    static constexpr std::byte eqI32{0x46};
+    static constexpr std::byte neI32{0x47};
+    static constexpr std::byte ltSI32{0x48};
+    static constexpr std::byte ltUI32{0x49};
+    static constexpr std::byte gtSI32{0x4A};
+    static constexpr std::byte gtUI32{0x4B};
+    static constexpr std::byte leSI32{0x4C};
+    static constexpr std::byte leUI32{0x4D};
+    static constexpr std::byte geSI32{0x4E};
+    static constexpr std::byte geUI32{0x4F};
+    static constexpr std::byte eqzI64{0x50};
+    static constexpr std::byte eqI64{0x51};
+    static constexpr std::byte neI64{0x52};
+    static constexpr std::byte ltSI64{0x53};
+    static constexpr std::byte ltUI64{0x54};
+    static constexpr std::byte gtSI64{0x55};
+    static constexpr std::byte gtUI64{0x56};
+    static constexpr std::byte leSI64{0x57};
+    static constexpr std::byte leUI64{0x58};
+    static constexpr std::byte geSI64{0x59};
+    static constexpr std::byte geUI64{0x5A};
+    static constexpr std::byte eqF32{0x5B};
+    static constexpr std::byte neF32{0x5C};
+    static constexpr std::byte ltF32{0x5D};
+    static constexpr std::byte gtF32{0x5E};
+    static constexpr std::byte leF32{0x5F};
+    static constexpr std::byte geF32{0x60};
+    static constexpr std::byte eqF64{0x61};
+    static constexpr std::byte neF64{0x62};
+    static constexpr std::byte ltF64{0x63};
+    static constexpr std::byte gtF64{0x64};
+    static constexpr std::byte leF64{0x65};
+    static constexpr std::byte geF64{0x66};
+
     // Numeric operations.
     static constexpr std::byte clzI32{0x67};
     static constexpr std::byte ctzI32{0x68};
@@ -93,6 +137,33 @@ struct WasmBinaryEncoding {
     static constexpr std::byte maxF64{0xA5};
     static constexpr std::byte copysignF64{0xA6};
     static constexpr std::byte wrap{0xA7};
+
+    // Conversion operations
+    static constexpr std::byte extendS{0xAC};
+    static constexpr std::byte extendU{0xAD};
+    static constexpr std::byte convertSI32F32{0xB2};
+    static constexpr std::byte convertUI32F32{0xB3};
+    static constexpr std::byte convertSI64F32{0xB4};
+    static constexpr std::byte convertUI64F32{0xB5};
+
+    static constexpr std::byte demoteF64ToF32{0xB6};
+
+    static constexpr std::byte convertSI32F64{0xB7};
+    static constexpr std::byte convertUI32F64{0xB8};
+    static constexpr std::byte convertSI64F64{0xB9};
+    static constexpr std::byte convertUI64F64{0xBA};
+
+    static constexpr std::byte promoteF32ToF64{0xBB};
+    static constexpr std::byte reinterpretF32AsI32{0xBC};
+    static constexpr std::byte reinterpretF64AsI64{0xBD};
+    static constexpr std::byte reinterpretI32AsF32{0xBE};
+    static constexpr std::byte reinterpretI64AsF64{0xBF};
+
+    static constexpr std::byte extendI328S{0xC0};
+    static constexpr std::byte extendI3216S{0xC1};
+    static constexpr std::byte extendI648S{0xC2};
+    static constexpr std::byte extendI6416S{0xC3};
+    static constexpr std::byte extendI6432S{0xC4};
   };
 
   /// Byte encodings of types in Wasm binaries