Add StableHLO shape assertion check pass and related tests (#2869)

This PR adds the CheckShapeAssertion from XLA into StableHLO:


https://github.com/openxla/xla/blob/949bb09eb7d0bb563f5b564fd7d0782e0687796e/xla/python/refine_polymorphic_shapes.cc#L72-L98

This idea was discussed in the discord channel:
https://discord.com/channels/999073994483433573/999074539138990131/1409946189721374803

This allows taking programs exported from JAX and using just the
satblehlo binary to produce a program that can be compile and executed
with PJRT plugins.

----


A usage example is for instance, export a program such as:

```
import jax
from jax import export
import jax.numpy as jnp
import numpy as np

def f1(x, y): # x: f32[a, 1], y : f32[a, 4]
 return x + y

# Assuming you have some actual args with concrete shapes
x = np.ones((3, 1), dtype=np.int32)
y = np.ones((3, 4), dtype=np.int32)
args_specs = export.symbolic_args_specs((x, y), 'a, ...')
exp = export.export(jax.jit(f1))(* args_specs)

code = exp.mlir_module()
print(code)
```

This gives us stableHLO like:

```
#loc = loc(unknown)
#loc2 = loc("x")
#loc3 = loc("y")
module @jit_f1 attributes {jax.uses_shape_polymorphism = true, mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<?x1xi32> loc(unknown), %arg1: tensor<?x4xi32> loc(unknown)) -> (tensor<?x4xi32> {jax.result_info = "result"}) {
    %c = stablehlo.constant dense<1> : tensor<i32> loc(#loc)
    %0 = stablehlo.get_dimension_size %arg0, dim = 0 : (tensor<?x1xi32>) -> tensor<i32> loc(#loc7)
    %1 = stablehlo.get_dimension_size %arg1, dim = 0 : (tensor<?x4xi32>) -> tensor<i32> loc(#loc7)
    %2 = stablehlo.compare  GE, %0, %c,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc8)
    stablehlo.custom_call @shape_assertion(%2, %0) {api_version = 2 : i32, error_message = "Input shapes do not match the polymorphic shapes specification. Expected value >= 1 for dimension variable 'a'. Using the following polymorphic shapes specifications: args[0].shape = (a, 1),args[1].shape = (a, 4). Obtained dimension variables: 'a' = {0} from specification 'a' for dimension args[0].shape[0] (= {0}), . Please see https://docs.jax.dev/en/latest/export/shape_poly.html#shape-assertion-errors for more details.", has_side_effect = true} : (tensor<i1>, tensor<i32>) -> () loc(#loc9)
    %3 = stablehlo.compare  EQ, %1, %0,  SIGNED : (tensor<i32>, tensor<i32>) -> tensor<i1> loc(#loc10)
    stablehlo.custom_call @shape_assertion(%3, %1, %0) {api_version = 2 : i32, error_message = "Input shapes do not match the polymorphic shapes specification. Found inconsistency between dimension size args[1].shape[0] (= {0}) and the specification 'a' (= {1}). Using the following polymorphic shapes specifications: args[0].shape = (a, 1),args[1].shape = (a, 4). Obtained dimension variables: 'a' = {1} from specification 'a' for dimension args[0].shape[0] (= {1}), . Please see https://docs.jax.dev/en/latest/export/shape_poly.html#shape-assertion-errors for more details.", has_side_effect = true} : (tensor<i1>, tensor<i32>, tensor<i32>) -> () loc(#loc9)
    %4 = call @_wrapped_jax_export_main(%0, %arg0, %arg1) : (tensor<i32>, tensor<?x1xi32>, tensor<?x4xi32>) -> tensor<?x4xi32> loc(#loc)
    return %4 : tensor<?x4xi32> loc(#loc)
  } loc(#loc)
  func.func private @_wrapped_jax_export_main(%arg0: tensor<i32> {jax.global_constant = "a"} loc(unknown), %arg1: tensor<?x1xi32> loc("x"), %arg2: tensor<?x4xi32> loc("y")) -> (tensor<?x4xi32> {jax.result_info = "result"}) {
    %c = stablehlo.constant dense<4> : tensor<1xi32> loc(#loc12)
    %0 = stablehlo.reshape %arg0 : (tensor<i32>) -> tensor<1xi32> loc(#loc12)
    %1 = stablehlo.concatenate %0, %c, dim = 0 : (tensor<1xi32>, tensor<1xi32>) -> tensor<2xi32> loc(#loc12)
    %2 = stablehlo.dynamic_broadcast_in_dim %arg1, %1, dims = [0, 1] : (tensor<?x1xi32>, tensor<2xi32>) -> tensor<?x4xi32> loc(#loc12)
    %3 = stablehlo.add %2, %arg2 : tensor<?x4xi32> loc(#loc12)
    return %3 : tensor<?x4xi32> loc(#loc)
  } loc(#loc)
} loc(#loc)
#loc1 = loc("<string>":13:6 to :46)
#loc4 = loc("<string>":7:8 to :13)
#loc5 = loc("<module>"(#loc1))
#loc6 = loc("f1"(#loc4))
#loc7 = loc("dimension_size"(#loc5))
#loc8 = loc("ge"(#loc5))
#loc9 = loc("shape_assertion"(#loc5))
#loc10 = loc("eq"(#loc5))
#loc11 = loc(callsite(#loc6 at #loc5))
#loc12 = loc("jit(f1)/add"(#loc11))
```

Before executing with PJRT one needs to refine shapes and remove the
shape assertions, with eg:

```
stablehlo-opt hello.mlir -stablehlo-refine-arguments="types='tensor<10x1xi32>, tensor<10x4xi32>'" -stablehlo-refine-shapes --stablehlo-check-shape-assertions
```

Which gives us:

```
module @jit_f1 attributes {jax.uses_shape_polymorphism = true, mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<10x1xi32>, %arg1: tensor<10x4xi32>) -> (tensor<10x4xi32> {jax.result_info = "result"}) {
    %c = stablehlo.constant dense<10> : tensor<i32>
    %c_0 = stablehlo.constant dense<true> : tensor<i1>
    %0 = call @_wrapped_jax_export_main(%arg0, %arg1) : (tensor<10x1xi32>, tensor<10x4xi32>) -> tensor<10x4xi32>
    return %0 : tensor<10x4xi32>
  }
  func.func private @_wrapped_jax_export_main(%arg0: tensor<10x1xi32>, %arg1: tensor<10x4xi32>) -> (tensor<10x4xi32> {jax.result_info = "result"}) {
    %c = stablehlo.constant dense<[10, 4]> : tensor<2xi32>
    %0 = stablehlo.dynamic_broadcast_in_dim %arg0, %c, dims = [0, 1] : (tensor<10x1xi32>, tensor<2xi32>) -> tensor<10x4xi32>
    %1 = stablehlo.add %0, %arg1 : tensor<10x4xi32>
    return %1 : tensor<10x4xi32>
  }
}
```

Author: dfalbel

BUILD.bazel

@@ -1135,6 +1135,7 @@ cc_library(
         "stablehlo/transforms/PassPipelines.cpp",
         "stablehlo/transforms/ShapeLegalizeToStablehlo.cpp",
         "stablehlo/transforms/StablehloCanonicalizeDynamism.cpp",
+        "stablehlo/transforms/StablehloCheckShapeAssertions.cpp",
         "stablehlo/transforms/StablehloCompatibilityExpander.cpp",
         "stablehlo/transforms/StablehloComplexMathExpander.cpp",
         "stablehlo/transforms/StablehloConvertToSignless.cpp",

docs/dynamism.md

@@ -105,13 +105,16 @@ Individually, the passes that tend to be useful for shape refinement are:
   shape information throughout the entire program.
 - [`stablehlo-canonicalize-dynamism`][canonicalize-dynamism] to replace dynamic
   ops with their static variants.
+- [`stablehlo-check-shape-assertions`][check-shape-assertions] to check and
+  remove shape assertions custom calls.
 
 See linked documentation for up-to-date information and examples.
 
 [remove-dynamism]:https://github.com/openxla/stablehlo/blob/ff13c96e56b73c62dcbb5b34b69f5ece9e71322f/stablehlo/transforms/Passes.h#L134
 [canonicalize-dynamism]:https://openxla.org/stablehlo/generated/stablehlo_passes#-stablehlo-canonicalize-dynamism
 [refine-arguments]:https://openxla.org/stablehlo/generated/stablehlo_passes#-stablehlo-refine-arguments
 [refine-shapes]:https://openxla.org/stablehlo/generated/stablehlo_passes#-stablehlo-refine-shapes
+[check-shape-assertions]:https://openxla.org/stablehlo/generated/stablehlo_passes#-stablehlo-check-shape-assertions
 
 ## Example: How is dynamism useful, and how can I use it?
 

docs/generated/stablehlo_passes.md

@@ -32,6 +32,32 @@ these ops are actually constants.
   %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<16xf32>
 ```
 
+### `-stablehlo-check-shape-assertions`
+
+_Check stablehlo.custom_call @shape_assertion ops._
+
+Validate shape_assertion custom calls.
+
+Shape assertions validate constraints on dynamic dimensions in StableHLO.
+For example if a framework needed to enforce a constraint of `DimA < 2`,
+the following IR could be emitted:
+
+```mlir
+%dimA = <get_dimension_size or input arg> : tensor<i32>
+%c2 = stablehlo.constant dense<2> : tensor<i32>
+%is_lt = stablehlo.compare LT %dimA, %c2 : tensor<i1>
+stablehlo.custom_call @shape_assertion(%is_lt) { error_message = "DimA must be less than 2" }
+```
+
+After the pass, if the shapes are correct, the `stablehlo.custom_call`
+will be removed.
+
+#### Options
+
+```
+-enable-shape-assertions : Whether shape assertions may generate errors.
+```
+
 ### `-stablehlo-compatibility-expander`
 
 _Compatibility expander for StableHLO operations._

stablehlo/tests/transforms/stablehlo_check_shape_assertions.mlir

@@ -0,0 +1,44 @@
+// RUN: stablehlo-opt --stablehlo-check-shape-assertions --split-input-file --verify-diagnostics %s | FileCheck %s --check-prefixes=CHECK
+
+// CHECK-LABEL: func.func @assertion_succeeds
+// CHECK-NOT: stablehlo.custom_call @shape_assertion
+// CHECK: return
+func.func @assertion_succeeds() {
+  %c1 = stablehlo.constant dense<true> : tensor<i1>
+  %c0 = stablehlo.constant dense<0> : tensor<i32>
+  stablehlo.custom_call @shape_assertion(%c1, %c0) {
+    api_version = 2 : i32,
+    error_message = "should not fire",
+    has_side_effect = true
+  } : (tensor<i1>, tensor<i32>) -> ()
+  return
+}
+
+// -----
+
+// ERR-LABEL: func.func @assertion_fails
+func.func @assertion_fails() {
+  %c1 = stablehlo.constant dense<false> : tensor<i1>
+  %c0 = stablehlo.constant dense<7> : tensor<i32>
+  // expected-error@+1 {{should fire}}
+  stablehlo.custom_call @shape_assertion(%c1, %c0) {
+    api_version = 2 : i32,
+    error_message = "should fire",
+    has_side_effect = true
+  } : (tensor<i1>, tensor<i32>) -> ()
+  return
+}
+
+// -----
+
+// ERR-LABEL: func.func @assertion_fails_not_constant
+func.func @assertion_fails_not_constant(%arg0 : tensor<i1>) {
+  %c0 = stablehlo.constant dense<7> : tensor<i32>
+  // expected-error@+1 {{expects static assert_what (operand #0)}}
+  stablehlo.custom_call @shape_assertion(%arg0, %c0) {
+    api_version = 2 : i32,
+    error_message = "not firing",
+    has_side_effect = true
+  } : (tensor<i1>, tensor<i32>) -> ()
+  return
+}

stablehlo/transforms/CMakeLists.txt

@@ -61,6 +61,7 @@ add_mlir_dialect_library(StablehloPasses
   PassPipelines.cpp
   ShapeLegalizeToStablehlo.cpp
   StablehloCanonicalizeDynamism.cpp
+  StablehloCheckShapeAssertions.cpp
   StablehloConvertToSignless.cpp
   StablehloCompatibilityExpander.cpp
   StablehloComplexMathExpander.cpp

stablehlo/transforms/Passes.td

@@ -55,6 +55,35 @@ def StablehloCanonicalizeDynamismPass : Pass<"stablehlo-canonicalize-dynamism",
   }];
 }
 
+def StablehloCheckShapeAssertionsPass
+    : Pass<"stablehlo-check-shape-assertions", "func::FuncOp"> {
+  let summary = "Check stablehlo.custom_call @shape_assertion ops.";
+
+  let description = [{
+    Validate shape_assertion custom calls.
+
+    Shape assertions validate constraints on dynamic dimensions in StableHLO.
+    For example if a framework needed to enforce a constraint of `DimA < 2`,
+    the following IR could be emitted:
+
+    ```mlir
+    %dimA = <get_dimension_size or input arg> : tensor<i32>
+    %c2 = stablehlo.constant dense<2> : tensor<i32>
+    %is_lt = stablehlo.compare LT %dimA, %c2 : tensor<i1>
+    stablehlo.custom_call @shape_assertion(%is_lt) { error_message = "DimA must be less than 2" }
+    ```
+
+    After the pass, if the shapes are correct, the `stablehlo.custom_call`
+    will be removed.
+  }];
+
+  let options = [
+    Option<"enable_shape_assertions", "enable-shape-assertions", "bool",
+           /*default=*/"true",
+           "Whether shape assertions may generate errors.">
+  ];
+}
+
 def StablehloCompatibilityExpanderPass : Pass<"stablehlo-compatibility-expander", "mlir::ModuleOp"> {
   let summary = "Compatibility expander for StableHLO operations.";
 

stablehlo/transforms/StablehloCheckShapeAssertions.cpp

@@ -0,0 +1,204 @@
+/* Copyright 2023 The JAX Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/LogicalResult.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Support/LLVM.h"
+#include "stablehlo/dialect/StablehloOps.h"
+#include "stablehlo/dialect/TypeInference.h"
+#include "stablehlo/transforms/Passes.h"
+
+namespace mlir {
+namespace stablehlo {
+
+#define GEN_PASS_DEF_STABLEHLOCHECKSHAPEASSERTIONSPASS
+#include "stablehlo/transforms/Passes.h.inc"
+
+namespace {
+
+constexpr llvm::StringRef shapeAssertionName = "shape_assertion";
+constexpr llvm::StringRef errorMessageAttrName = "error_message";
+// We bound the number of error_message_inputs for using llvm::formatv
+constexpr int maxErrorMessageInputs = 32;  // TODO(necula): Remove this bound
+
+// This pass is needed when we have shape assertions. A shape assertion is
+// represented via the `stablehlo.custom_call @shape_assertion`
+// custom call, and represents an assertion that the first operand
+// (`assert_what`) evaluates to `true`. The custom call also has an
+// `error_message` string attribute, and a variadic number
+// of integer scalar operands that may be used to format the error message.
+// The `error_message` may contain format specifiers `{0}`, `{1}`, ..., that
+// are replaced with the values of the error message inputs. The formatting is
+// done with the `llvm::formatv` function
+// (https://llvm.org/docs/ProgrammersManual.html#formatting-strings-the-formatv-function).
+//
+struct CheckShapeAssertionsPass
+    : public impl::StablehloCheckShapeAssertionsPassBase<
+          CheckShapeAssertionsPass> {
+  using StablehloCheckShapeAssertionsPassBase::
+      StablehloCheckShapeAssertionsPassBase;
+
+  void runOnOperation() override {
+    func::FuncOp funcOp = getOperation();
+    funcOp.walk([this](CustomCallOp op) {
+      if (op.getCallTargetName() != shapeAssertionName) return;
+      if (!enable_shape_assertions) {
+        op.erase();
+        return;
+      }
+      // Check first for ill-formed assertions, rather than silently fail.
+      if (failed(verifyShapeAssertion(op))) {
+        signalPassFailure();
+        return;
+      }
+      OperandRange inputs = op.getInputs();
+      SmallVector<int64_t> assertWhat;
+      if (failed(hlo::matchInts(inputs[0], assertWhat))) {
+        op.emitError() << "expects static assert_what (operand #0)";
+        signalPassFailure();
+        return;
+      }
+      if (assertWhat[0] != 0) {
+        op.erase();
+        return;
+      }
+      StringRef errorMessage = getErrorMessage(op);
+      SmallVector<int64_t> errorMessageInputs;
+      for (size_t i = 1; i < inputs.size(); ++i) {
+        SmallVector<int64_t> input;
+        if (failed(hlo::matchInts(inputs[i], input))) {
+          op.emitError() << "expects static error_message_input (operand #" << i
+                         << ")";
+          signalPassFailure();
+          return;
+        }
+        errorMessageInputs.push_back(input[0]);
+      }
+      op.emitError(formatErrorMessage(errorMessage, errorMessageInputs));
+      signalPassFailure();
+    });
+  }
+
+ private:
+  LogicalResult verifyShapeAssertion(CustomCallOp op) {
+    if (!(1 <= op->getNumOperands() &&
+          op->getNumOperands() <= 1 + maxErrorMessageInputs))
+      return op.emitError() << "expects 1 <= size(operands) <= "
+                            << (1 + maxErrorMessageInputs);
+    int nrErrorMessageInputs = op.getNumOperands() - 1;
+    if (op->getNumResults() != 0)
+      return op.emitError("expects size(results) = 0");
+    for (const auto& attr : op->getAttrs()) {
+      if (attr.getName() != "api_version" &&
+          attr.getName() != "backend_config" &&
+          attr.getName() != "call_target_name" &&
+          attr.getName() != "error_message" &&
+          attr.getName() != "has_side_effect")
+        return op.emitError()
+               << attr.getName() << " is not a supported attribute";
+    }
+    if (!op.hasEmptyBackendConfig())
+      return op.emitError() << "expects an empty backend_config";
+    if (op.getCallTargetName() != shapeAssertionName)
+      return op.emitError() << "expects @shape_assertion";
+
+    // input[0] (assert_what) : tensor<i1>
+    auto assertWhatType = dyn_cast<ShapedType>(op.getInputs()[0].getType());
+    if (!assertWhatType || !assertWhatType.hasRank() ||
+        assertWhatType.getRank() != 0 ||
+        !assertWhatType.getElementType().isSignlessInteger() ||
+        assertWhatType.getElementTypeBitWidth() != 1)
+      return op.emitError() << "expects assert_what (operand #0) "
+                            << "to be a constant of type tensor<i1>";
+
+    // input[1:] (error_message_inputs) : tensor<i32> or tensor<i64>
+    for (int i = 0; i < nrErrorMessageInputs; ++i) {
+      auto errorMessageInputType =
+          dyn_cast<ShapedType>(op.getInputs()[i + 1].getType());
+      if (!errorMessageInputType || !errorMessageInputType.hasRank() ||
+          errorMessageInputType.getRank() != 0 ||
+          !errorMessageInputType.getElementType().isSignlessInteger() ||
+          (errorMessageInputType.getElementTypeBitWidth() != 32 &&
+           errorMessageInputType.getElementTypeBitWidth() != 64))
+        return op.emitError()
+               << "expects error_message_input (operand #" << (i + 1) << ") "
+               << "to be a constant of type tensor<i32> or tensor<i64>";
+    }
+
+    if (!op->hasAttr(errorMessageAttrName))
+      return op.emitError() << "expects an error_message attribute";
+
+    // error_message contains valid format specifiers.
+    StringRef errorMessage = getErrorMessage(op);
+
+    // format specs: "{" index ["," layout] [":" format] "}"
+    size_t spec_begin = errorMessage.find_first_of('{');
+    size_t spec_end = errorMessage.find_first_of(",:}", spec_begin);
+
+    // Check that all specs reference valid input indices.
+    while (spec_begin != StringRef::npos && spec_end != StringRef::npos) {
+      StringRef index_str =
+          errorMessage.substr(spec_begin + 1, spec_end - spec_begin - 1);
+
+      int32_t index;
+      if (!index_str.getAsInteger(10, index) &&
+          !(0 <= index && index < nrErrorMessageInputs)) {
+        return op.emitError()
+               << "expects error_message to contain format specifiers with "
+               << "error_message_input index less than " << nrErrorMessageInputs
+               << ". Found specifier "
+               << errorMessage.substr(spec_begin, spec_end - spec_begin + 1);
+      }
+
+      spec_begin = errorMessage.find_first_of('{', spec_begin + 1);
+      spec_end = errorMessage.find_first_of(",:}", spec_begin);
+    }
+    return success();
+  }
+
+  StringRef getErrorMessage(CustomCallOp op) const {
+    return cast<StringAttr>(op->getAttr(errorMessageAttrName)).getValue();
+  }
+
+  std::string formatErrorMessage(
+      StringRef errorMessage,
+      const SmallVector<int64_t>& errorMessageInputs) const {
+    int nrErrorMessageInputs = errorMessageInputs.size();
+    auto errorMessageFormat = errorMessage.data();
+    if (nrErrorMessageInputs == 0) return errorMessageFormat;
+    auto errInput = [nrErrorMessageInputs, &errorMessageInputs](int idx) {
+      return (idx < nrErrorMessageInputs ? errorMessageInputs[idx] : -1);
+    };
+    return llvm::formatv(
+        false, errorMessageFormat, errInput(0), errInput(1), errInput(2),
+        errInput(3), errInput(4), errInput(5), errInput(6), errInput(7),
+        errInput(8), errInput(9), errInput(10), errInput(11), errInput(12),
+        errInput(13), errInput(14), errInput(15), errInput(16), errInput(17),
+        errInput(18), errInput(19), errInput(20), errInput(21), errInput(22),
+        errInput(23), errInput(24), errInput(25), errInput(26), errInput(27),
+        errInput(28), errInput(29), errInput(30), errInput(31));
+  }
+};
+
+}  // namespace
+
+}  // namespace stablehlo
+}  // namespace mlir