[MLIR][WasmSSA] Add global_set and select ops (#196613)

Adds two missing ops to the WasmSSA dialect.

- `wasmssa.global_set`
- `wasmssa.select`

---------

Co-authored-by: Ferdinand Lemaire <flscminecraft@gmail.com>

mlir/include/mlir/Dialect/WasmSSA/IR/WasmSSA.h

@@ -38,6 +38,7 @@
 #include "mlir/Interfaces/CallInterfaces.h"
 #include "mlir/Interfaces/FunctionInterfaces.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
 
 //===----------------------------------------------------------------------===//
 // WebAssembly Constraints

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

@@ -16,6 +16,7 @@ include "mlir/Dialect/WasmSSA/IR/WasmSSAInterfaces.td"
 
 include "mlir/Interfaces/FunctionInterfaces.td"
 include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/IR/BuiltinAttributeInterfaces.td"
 include "mlir/IR/SymbolInterfaces.td"
 
@@ -381,6 +382,25 @@ def WasmSSA_GlobalGetOp : WasmSSA_Op<"global_get", [DeclareOpInterfaceMethods<Sy
   let assemblyFormat = "$global attr-dict `:` type($global_val)";
 }
 
+def WasmSSA_GlobalSetOp : WasmSSA_Op<"global_set", [
+    DeclareOpInterfaceMethods<SymbolUserOpInterface>]> {
+  let summary = "Sets the value of the mutable global passed as argument.";
+  let description = [{
+     Sets the value of the mutable global referenced by the symbol to the given
+     value.
+
+     Example:
+
+     ```mlir
+     // Sets the value of `@global_0` to `%v`.
+     wasmssa.global_set @global_0 to %v : i32
+     ```
+    }];
+  let arguments = (ins FlatSymbolRefAttr: $global,
+                       WasmSSA_ValType: $value);
+  let assemblyFormat = "$global `to` $value `:` type($value) attr-dict";
+}
+
 def WasmSSA_IfOp : WasmSSA_Op<"if", [Terminator,
     DeclareOpInterfaceMethods<LabelLevelOpInterface>]> {
   let summary = "Execute the if region if condition value is non-zero, the else region otherwise.";
@@ -609,6 +629,29 @@ def WasmSSA_ReturnOp : WasmSSA_Op<"return", [Terminator]> {
   ];
 }
 
+def WasmSSA_SelectOp : WasmSSA_Op<"select", [
+    AllTypesMatch<["true_value", "false_value", "result"]>,
+    Pure]> {
+  let summary = "Select between two values based on a condition.";
+  let description = [{
+     Returns `true_value` if `condition` is non-zero, `false_value` otherwise.
+     Both operands and the result must share the same Wasm value type.
+
+     Example:
+
+     ```mlir
+     // Returns %a if %cond is non-zero, otherwise %b.
+     %r = wasmssa.select %cond, %a, %b : i32
+     ```
+    }];
+  let arguments = (ins WasmSSA_ValType: $true_value,
+                       WasmSSA_ValType: $false_value,
+                       I32: $condition);
+  let results = (outs WasmSSA_ValType: $result);
+  let assemblyFormat =
+      "$condition `,` $true_value `,` $false_value attr-dict `:` type($result)";
+}
+
 // ---- Numeric ops
 
 class WasmSSA_BinaryNumericalOp<string mnemonic, string summaryStr, string descStr,

mlir/lib/Conversion/RaiseWasm/RaiseWasmMLIR.cpp

@@ -671,6 +671,23 @@ struct WasmGlobalWithGetGlobalInitConversion
   }
 };
 
+struct WasmGlobalSetOpConversion : OpConversionPattern<GlobalSetOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(GlobalSetOp globalSetOp, GlobalSetOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = globalSetOp.getLoc();
+    auto globalPtr = memref::GetGlobalOp::create(
+        rewriter, loc, MemRefType::get({1}, adaptor.getValue().getType()),
+        globalSetOp.getGlobal());
+    auto idx = arith::ConstantIndexOp::create(rewriter, loc, 0);
+    rewriter.replaceOpWithNewOp<memref::StoreOp>(
+        globalSetOp, adaptor.getValue(), globalPtr.getResult(),
+        ValueRange{idx.getResult()});
+    return success();
+  }
+};
+
 struct WasmMemoryOpConversion : OpConversionPattern<MemOp> {
   using OpConversionPattern::OpConversionPattern;
 
@@ -796,6 +813,24 @@ struct WasmReturnOpConversion : OpConversionPattern<ReturnOp> {
   }
 };
 
+struct WasmSelectOpConversion : OpConversionPattern<SelectOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(SelectOp selectOp, SelectOp::Adaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto loc = selectOp.getLoc();
+    auto zero =
+        arith::ConstantOp::create(rewriter, loc, rewriter.getI32IntegerAttr(0));
+    auto flag = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::ne,
+                                      adaptor.getCondition(), zero.getResult());
+    rewriter.replaceOpWithNewOp<arith::SelectOp>(selectOp, flag.getResult(),
+                                                 adaptor.getTrueValue(),
+                                                 adaptor.getFalseValue());
+    return success();
+  }
+};
+
 struct RaiseWasmMLIRPass : public impl::RaiseWasmMLIRBase<RaiseWasmMLIRPass> {
   void runOnOperation() override {
     ConversionTarget target{getContext()};
@@ -876,6 +911,7 @@ void mlir::populateRaiseWasmMLIRConversionPatterns(
            WasmGeSIOpConversion,
            WasmGeUIOpConversion,
            WasmGlobalImportOpConverter,
+           WasmGlobalSetOpConversion,
            WasmGlobalWithConstInitConversion,
            WasmGlobalWithGetGlobalInitConversion,
            WasmGtOpConversion,
@@ -906,6 +942,7 @@ void mlir::populateRaiseWasmMLIRConversionPatterns(
            WasmReturnOpConversion,
            WasmRotlOpConversion,
            WasmRotrOpConversion,
+           WasmSelectOpConversion,
            WasmShLOpConversion,
            WasmShRSOpConversion,
            WasmShRUOpConversion,

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

@@ -331,6 +331,44 @@ GlobalGetOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// GlobalSetOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult
+GlobalSetOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  Operation *symTabOp = SymbolTable::getNearestSymbolTable(*this);
+  StringRef referencedSymbol = getGlobal();
+  Operation *definitionOp = symbolTable.lookupSymbolIn(
+      symTabOp, StringAttr::get(this->getContext(), referencedSymbol));
+  if (!definitionOp)
+    return emitError() << "symbol @" << referencedSymbol << " is undefined";
+
+  Type globalType;
+  bool isMutable = false;
+  if (auto global = dyn_cast<GlobalOp>(definitionOp)) {
+    globalType = global.getType();
+    isMutable = global.getIsMutable();
+  } else if (auto globalImport = dyn_cast<GlobalImportOp>(definitionOp)) {
+    globalType = globalImport.getType();
+    isMutable = globalImport.getIsMutable();
+  } else {
+    return emitError() << "symbol @" << referencedSymbol
+                       << " is not a global symbol";
+  }
+
+  if (!isMutable)
+    return emitError("global.set target must be mutable");
+
+  Type valueType = getValue().getType();
+  if (globalType != valueType)
+    return emitError("global.set value type does not match target global "
+                     "type: expected ")
+           << globalType << " but got " << valueType;
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // GlobalImportOp
 //===----------------------------------------------------------------------===//

mlir/test/Conversion/RaiseWasm/wasm-global-set-to-memref.mlir

@@ -0,0 +1,35 @@
+// RUN: mlir-opt %s --raise-wasm-mlir | FileCheck %s
+
+wasmssa.global @global_i32 i32 mutable : {
+  %0 = wasmssa.const 10 : i32
+  wasmssa.return %0 : i32
+}
+
+wasmssa.global @global_f64 f64 mutable : {
+  %0 = wasmssa.const 3.14 : f64
+  wasmssa.return %0 : f64
+}
+
+// CHECK-LABEL:   func.func @set_global_i32() {
+wasmssa.func exported @set_global_i32() {
+// CHECK:           %[[VAL_0:.*]] = arith.constant 42 : i32
+  %0 = wasmssa.const 42 : i32
+// CHECK:           %[[VAL_1:.*]] = memref.get_global @global_i32 : memref<1xi32>
+// CHECK:           %[[VAL_2:.*]] = arith.constant 0 : index
+// CHECK:           memref.store %[[VAL_0]], %[[VAL_1]]{{\[}}%[[VAL_2]]] : memref<1xi32>
+  wasmssa.global_set @global_i32 to %0 : i32
+// CHECK:           return
+  wasmssa.return
+}
+
+// CHECK-LABEL:   func.func @set_global_f64() {
+wasmssa.func exported @set_global_f64() {
+// CHECK:           %[[VAL_0:.*]] = arith.constant 2.500000e-01 : f64
+  %0 = wasmssa.const 0.25 : f64
+// CHECK:           %[[VAL_1:.*]] = memref.get_global @global_f64 : memref<1xf64>
+// CHECK:           %[[VAL_2:.*]] = arith.constant 0 : index
+// CHECK:           memref.store %[[VAL_0]], %[[VAL_1]]{{\[}}%[[VAL_2]]] : memref<1xf64>
+  wasmssa.global_set @global_f64 to %0 : f64
+// CHECK:           return
+  wasmssa.return
+}

mlir/test/Conversion/RaiseWasm/wasm-select-to-arith-select.mlir

@@ -0,0 +1,65 @@
+// RUN: mlir-opt %s --raise-wasm-mlir | FileCheck %s
+
+// CHECK-LABEL:   func.func @select_i32() -> i32 {
+wasmssa.func @select_i32() -> i32 {
+// CHECK:           %[[COND:.*]] = arith.constant 1 : i32
+  %cond = wasmssa.const 1 : i32
+// CHECK:           %[[VAL_A:.*]] = arith.constant 12 : i32
+  %a = wasmssa.const 12 : i32
+// CHECK:           %[[VAL_B:.*]] = arith.constant 50 : i32
+  %b = wasmssa.const 50 : i32
+// CHECK:           %[[ZERO:.*]] = arith.constant 0 : i32
+// CHECK:           %[[FLAG:.*]] = arith.cmpi ne, %[[COND]], %[[ZERO]] : i32
+// CHECK:           %[[RES:.*]] = arith.select %[[FLAG]], %[[VAL_A]], %[[VAL_B]] : i32
+  %r = wasmssa.select %cond, %a, %b : i32
+// CHECK:           return %[[RES]] : i32
+  wasmssa.return %r : i32
+}
+
+// CHECK-LABEL:   func.func @select_i64() -> i64 {
+wasmssa.func @select_i64() -> i64 {
+// CHECK:           %[[COND:.*]] = arith.constant 0 : i32
+  %cond = wasmssa.const 0 : i32
+// CHECK:           %[[VAL_A:.*]] = arith.constant 12 : i64
+  %a = wasmssa.const 12 : i64
+// CHECK:           %[[VAL_B:.*]] = arith.constant 50 : i64
+  %b = wasmssa.const 50 : i64
+// CHECK:           %[[ZERO:.*]] = arith.constant 0 : i32
+// CHECK:           %[[FLAG:.*]] = arith.cmpi ne, %[[COND]], %[[ZERO]] : i32
+// CHECK:           %[[RES:.*]] = arith.select %[[FLAG]], %[[VAL_A]], %[[VAL_B]] : i64
+  %r = wasmssa.select %cond, %a, %b : i64
+// CHECK:           return %[[RES]] : i64
+  wasmssa.return %r : i64
+}
+
+// CHECK-LABEL:   func.func @select_f32() -> f32 {
+wasmssa.func @select_f32() -> f32 {
+// CHECK:           %[[COND:.*]] = arith.constant 1 : i32
+  %cond = wasmssa.const 1 : i32
+// CHECK:           %[[VAL_A:.*]] = arith.constant 1.250000e-01 : f32
+  %a = wasmssa.const 0.125 : f32
+// CHECK:           %[[VAL_B:.*]] = arith.constant 2.500000e-01 : f32
+  %b = wasmssa.const 0.25 : f32
+// CHECK:           %[[ZERO:.*]] = arith.constant 0 : i32
+// CHECK:           %[[FLAG:.*]] = arith.cmpi ne, %[[COND]], %[[ZERO]] : i32
+// CHECK:           %[[RES:.*]] = arith.select %[[FLAG]], %[[VAL_A]], %[[VAL_B]] : f32
+  %r = wasmssa.select %cond, %a, %b : f32
+// CHECK:           return %[[RES]] : f32
+  wasmssa.return %r : f32
+}
+
+// CHECK-LABEL:   func.func @select_f64() -> f64 {
+wasmssa.func @select_f64() -> f64 {
+// CHECK:           %[[COND:.*]] = arith.constant 0 : i32
+  %cond = wasmssa.const 0 : i32
+// CHECK:           %[[VAL_A:.*]] = arith.constant 3.140000e+00 : f64
+  %a = wasmssa.const 3.14 : f64
+// CHECK:           %[[VAL_B:.*]] = arith.constant 2.718000e+00 : f64
+  %b = wasmssa.const 2.718 : f64
+// CHECK:           %[[ZERO:.*]] = arith.constant 0 : i32
+// CHECK:           %[[FLAG:.*]] = arith.cmpi ne, %[[COND]], %[[ZERO]] : i32
+// CHECK:           %[[RES:.*]] = arith.select %[[FLAG]], %[[VAL_A]], %[[VAL_B]] : f64
+  %r = wasmssa.select %cond, %a, %b : f64
+// CHECK:           return %[[RES]] : f64
+  wasmssa.return %r : f64
+}

mlir/test/Dialect/WasmSSA/custom_parser/global.mlir

@@ -19,6 +19,12 @@ module {
     %0 = wasmssa.global_get @global_0 : i32
     wasmssa.return %0 : i32
   }
+
+  wasmssa.func @set_global() {
+    %0 = wasmssa.const 7 : i32
+    wasmssa.global_set @global_2 to %0 : i32
+    wasmssa.return
+  }
 }
 
 // CHECK-LABEL:   wasmssa.import_global "from_js" from "env" as @global_0 : i32
@@ -42,3 +48,9 @@ module {
 // CHECK:           %[[VAL_0:.*]] = wasmssa.global_get @global_0 : i32
 // CHECK:           wasmssa.return %[[VAL_0]] : i32
 // CHECK:         }
+
+// CHECK-LABEL:   wasmssa.func @set_global() {
+// CHECK:           %[[VAL_0:.*]] = wasmssa.const 7 : i32
+// CHECK:           wasmssa.global_set @global_2 to %[[VAL_0]] : i32
+// CHECK:           wasmssa.return
+// CHECK:         }

mlir/test/Dialect/WasmSSA/custom_parser/select.mlir

@@ -0,0 +1,37 @@
+// RUN: mlir-opt %s | FileCheck %s
+
+module {
+  wasmssa.func @select_i32(%cond: !wasmssa<local ref to i32>,
+                           %a: !wasmssa<local ref to i32>,
+                           %b: !wasmssa<local ref to i32>) -> i32 {
+    %0 = wasmssa.local_get %cond : ref to i32
+    %1 = wasmssa.local_get %a : ref to i32
+    %2 = wasmssa.local_get %b : ref to i32
+    %r = wasmssa.select %0, %1, %2 : i32
+    wasmssa.return %r : i32
+  }
+
+  wasmssa.func @select_f64(%cond: !wasmssa<local ref to i32>,
+                           %a: !wasmssa<local ref to f64>,
+                           %b: !wasmssa<local ref to f64>) -> f64 {
+    %0 = wasmssa.local_get %cond : ref to i32
+    %1 = wasmssa.local_get %a : ref to f64
+    %2 = wasmssa.local_get %b : ref to f64
+    %r = wasmssa.select %0, %1, %2 : f64
+    wasmssa.return %r : f64
+  }
+}
+
+// CHECK-LABEL:   wasmssa.func @select_i32(
+// CHECK:           %[[COND:.*]] = wasmssa.local_get
+// CHECK:           %[[A:.*]] = wasmssa.local_get
+// CHECK:           %[[B:.*]] = wasmssa.local_get
+// CHECK:           %[[R:.*]] = wasmssa.select %[[COND]], %[[A]], %[[B]] : i32
+// CHECK:           wasmssa.return %[[R]] : i32
+
+// CHECK-LABEL:   wasmssa.func @select_f64(
+// CHECK:           %[[COND:.*]] = wasmssa.local_get
+// CHECK:           %[[A:.*]] = wasmssa.local_get
+// CHECK:           %[[B:.*]] = wasmssa.local_get
+// CHECK:           %[[R:.*]] = wasmssa.select %[[COND]], %[[A]], %[[B]] : f64
+// CHECK:           wasmssa.return %[[R]] : f64

mlir/test/Dialect/WasmSSA/global-invalid.mlir

@@ -40,3 +40,58 @@ module {
     wasmssa.return %0 : i32
   }
 }
+
+// -----
+
+module {
+  wasmssa.func @f() {
+    %0 = wasmssa.const 1 : i32
+    // expected-error@+1 {{symbol @missing is undefined}}
+    wasmssa.global_set @missing to %0 : i32
+    wasmssa.return
+  }
+}
+
+// -----
+
+module {
+  wasmssa.global @global_0 i32 : {
+    %0 = wasmssa.const 0 : i32
+    wasmssa.return %0 : i32
+  }
+  wasmssa.func @f() {
+    %0 = wasmssa.const 1 : i32
+    // expected-error@+1 {{global.set target must be mutable}}
+    wasmssa.global_set @global_0 to %0 : i32
+    wasmssa.return
+  }
+}
+
+// -----
+
+module {
+  wasmssa.global @global_0 i64 mutable : {
+    %0 = wasmssa.const 0 : i64
+    wasmssa.return %0 : i64
+  }
+  wasmssa.func @f() {
+    %0 = wasmssa.const 1 : i32
+    // expected-error@+1 {{global.set value type does not match target global type: expected 'i64' but got 'i32'}}
+    wasmssa.global_set @global_0 to %0 : i32
+    wasmssa.return
+  }
+}
+
+// -----
+
+module {
+  wasmssa.func @not_global() {
+    wasmssa.return
+  }
+  wasmssa.func @f() {
+    %0 = wasmssa.const 1 : i32
+    // expected-error@+1 {{symbol @not_global is not a global symbol}}
+    wasmssa.global_set @not_global to %0 : i32
+    wasmssa.return
+  }
+}