[MLIR][WASM] Custom assembly format for if memory and table ops (#152668)

Co-authored-by: Luc Forget <[email protected]>

Author: lforg37

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

@@ -329,11 +329,11 @@ def WasmSSA_IfOp : WasmSSA_Op<"if", [Terminator,
 
      ```mlir
      // Runs the if clause is %a is non-zero
-     "wasmssa.if"(%a)[^bb1] ({
+     wasmssa.if %a {
         // Execute if %a is non-zero
-     },{
+     } else {
         // else clause
-     }) : (i32) -> ()
+     }
      ```
     }];
   let arguments = (ins I32:$condition, Variadic<WasmSSA_ValType>: $inputs);
@@ -359,6 +359,7 @@ def WasmSSA_IfOp : WasmSSA_Op<"if", [Terminator,
       return createBlock(getElse());
     }
   }];
+  let assemblyFormat = "$condition (`(`$inputs^`)` `:` type($inputs))? attr-dict  `:` $if custom<ElseRegion>($else) `>` $target";
 }
 
 def WasmSSA_LocalOp : WasmSSA_Op<"local", [
@@ -445,7 +446,7 @@ def WasmSSA_MemOp : WasmSSA_Op<"memory", [Symbol]> {
 
      ```mlir
      // Define the `mem_0` memory with defined bounds of 0 -> 65536
-     "wasmssa.memory"() <{limits = !wasmssa<limit[0:65536]>, sym_name = "mem_0"}> : () -> ()
+     wasmssa.memory @mem_0 !wasmssa<limit[0:65536]>
      ```
     }];
   let arguments = (ins SymbolNameAttr: $sym_name,
@@ -456,6 +457,8 @@ def WasmSSA_MemOp : WasmSSA_Op<"memory", [Symbol]> {
     "::llvm::StringRef":$symbol,
     "wasmssa::LimitType":$limit)>
   ];
+
+  let assemblyFormat = "$sym_name custom<WasmVisibility>($sym_visibility) $limits attr-dict";
 }
 
 def WasmSSA_MemImportOp : WasmSSA_Op<"import_mem", [Symbol, ImportOpInterface]> {
@@ -494,6 +497,7 @@ def WasmSSA_TableOp : WasmSSA_Op<"table", [Symbol]> {
   let builders = [OpBuilder<(ins
       "::llvm::StringRef":$symbol,
       "wasmssa::TableType":$type)>];
+  let assemblyFormat = "$sym_name custom<WasmVisibility>($sym_visibility) $type attr-dict";
 }
 
 def WasmSSA_TableImportOp : WasmSSA_Op<"import_table", [Symbol, ImportOpInterface]> {

mlir/lib/Dialect/WasmSSA/IR/WasmSSAOps.cpp

@@ -22,14 +22,54 @@
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
+using namespace mlir;
+namespace {
+ParseResult parseElseRegion(OpAsmParser &opParser, Region &elseRegion) {
+  std::string keyword;
+  std::ignore = opParser.parseOptionalKeywordOrString(&keyword);
+  if (keyword == "else")
+    return opParser.parseRegion(elseRegion);
+  return ParseResult::success();
+}
+
+void printElseRegion(OpAsmPrinter &opPrinter, Operation *op,
+                     Region &elseRegion) {
+  if (elseRegion.empty())
+    return;
+  opPrinter.printKeywordOrString("else ");
+  opPrinter.printRegion(elseRegion);
+}
+
+ParseResult parseWasmVisibility(OpAsmParser &opParser, StringAttr &visibility) {
+  std::string keyword;
+  auto initLocation = opParser.getCurrentLocation();
+  std::ignore = opParser.parseOptionalKeywordOrString(&keyword);
+  if (keyword == "nested" or keyword == "") {
+    visibility = StringAttr::get(opParser.getContext(), "nested");
+    return ParseResult::success();
+  }
+
+  if (keyword == "public" || keyword == "private") {
+    visibility = StringAttr::get(opParser.getContext(), keyword);
+    return ParseResult::success();
+  }
+  opParser.emitError(initLocation, "expecting symbol visibility");
+  return ParseResult::failure();
+}
+
+void printWasmVisibility(OpAsmPrinter &opPrinter, Operation *op,
+                         Attribute visibility) {
+  opPrinter.printKeywordOrString(cast<StringAttr>(visibility).strref());
+}
+} // namespace
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/WasmSSA/IR/WasmSSAOps.cpp.inc"
 
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/Types.h"
 #include "llvm/Support/LogicalResult.h"
 
-using namespace mlir;
 using namespace wasmssa;
 
 namespace {

mlir/test/Dialect/WasmSSA/custom_parser/if.mlir

@@ -0,0 +1,53 @@
+// RUN: mlir-opt %s | FileCheck %s
+
+// CHECK-LABEL:   wasmssa.func nested @func_0(
+// CHECK-SAME:      %[[ARG0:.*]]: !wasmssa<local ref to i32>) -> i32 {
+// CHECK:           %[[VAL_0:.*]] = wasmssa.local_get %[[ARG0]] :  ref to i32
+// CHECK:           wasmssa.if %[[VAL_0]] : {
+// CHECK:             %[[VAL_1:.*]] = wasmssa.const 5.000000e-01 : f32
+// CHECK:             wasmssa.block_return %[[VAL_1]] : f32
+// CHECK:           } "else "{
+// CHECK:             %[[VAL_2:.*]] = wasmssa.const 2.500000e-01 : f32
+// CHECK:             wasmssa.block_return %[[VAL_2]] : f32
+// CHECK:           }> ^bb1
+// CHECK:         ^bb1(%[[VAL_3:.*]]: f32):
+// CHECK:           wasmssa.return %[[VAL_3]] : f32
+wasmssa.func nested @func_0(%arg0 : !wasmssa<local ref to i32>) -> i32 {
+  %cond = wasmssa.local_get %arg0 : ref to i32
+  wasmssa.if %cond : {
+    %c0 = wasmssa.const 0.5 : f32
+    wasmssa.block_return %c0 : f32
+  } else {
+   %c1 = wasmssa.const 0.25 : f32
+   wasmssa.block_return %c1 : f32
+  } >^bb1
+ ^bb1(%retVal: f32):
+  wasmssa.return %retVal : f32
+}
+
+// CHECK-LABEL:   wasmssa.func nested @func_1(
+// CHECK-SAME:      %[[ARG0:.*]]: !wasmssa<local ref to i32>) -> i32 {
+// CHECK:           %[[VAL_0:.*]] = wasmssa.local_get %[[ARG0]] :  ref to i32
+// CHECK:           %[[VAL_1:.*]] = wasmssa.local of type i32
+// CHECK:           %[[VAL_2:.*]] = wasmssa.const 0 : i64
+// CHECK:           wasmssa.if %[[VAL_0]] : {
+// CHECK:             %[[VAL_3:.*]] = wasmssa.const 1 : i32
+// CHECK:             wasmssa.local_set %[[VAL_1]] :  ref to i32 to %[[VAL_3]] : i32
+// CHECK:             wasmssa.block_return
+// CHECK:           } > ^bb1
+// CHECK:         ^bb1:
+// CHECK:           %[[VAL_4:.*]] = wasmssa.local_get %[[VAL_1]] :  ref to i32
+// CHECK:           wasmssa.return %[[VAL_4]] : i32
+wasmssa.func nested @func_1(%arg0 : !wasmssa<local ref to i32>) -> i32 {
+  %cond = wasmssa.local_get %arg0 : ref to i32
+  %var = wasmssa.local of type i32
+  %zero = wasmssa.const 0
+  wasmssa.if %cond : {
+    %c1 = wasmssa.const 1 : i32
+    wasmssa.local_set %var : ref to i32 to %c1 : i32
+    wasmssa.block_return
+  } >^bb1
+ ^bb1:
+  %res = wasmssa.local_get %var : ref to i32
+  wasmssa.return %res : i32
+}

mlir/test/Dialect/WasmSSA/custom_parser/memory.mlir

@@ -0,0 +1,7 @@
+// RUN: mlir-opt %s | FileCheck %s
+
+// CHECK:   wasmssa.memory @mem0 public !wasmssa<limit[0: 65536]>
+wasmssa.memory @mem0 public !wasmssa<limit[0:65536]>
+
+// CHECK:   wasmssa.memory @mem1 nested !wasmssa<limit[512:]>
+wasmssa.memory @mem1 !wasmssa<limit[512:]>

mlir/test/Dialect/WasmSSA/custom_parser/table.mlir

@@ -0,0 +1,7 @@
+// RUN: mlir-opt %s | FileCheck %s
+
+// CHECK:   wasmssa.table @tab0 public !wasmssa<tabletype !wasmssa.externref [0: 65536]>
+wasmssa.table @tab0 public !wasmssa<tabletype !wasmssa.externref [0:65536]>
+
+// CHECK:   wasmssa.table @tab1 nested !wasmssa<tabletype !wasmssa.funcref [348:]>
+wasmssa.table @tab1 !wasmssa<tabletype !wasmssa.funcref [348:]>