Add a python API for reference interpreter (#2062)

```py
ASM = """
func.func @test(%arg0: tensor<2x2xf16>) -> tensor<2x2xf16> {
  %0 = stablehlo.add %arg0, %arg0 : (tensor<2x2xf16>, tensor<2x2xf16>) -> tensor<2x2xf16>
  func.return %0 : tensor<2x2xf16>
}
"""

m = ir.Module.parse(ASM)
arg = np.asarray([1, 2, 3, 4], np.float16).reshape(2,2)
args = [ir.DenseIntElementsAttr.get(arg)]
res = stablehlo.eval_module(m, args)
```

- Added python API for reference interpreter, and test cases for
int/float datatypes
- Added reference API which takes DenseElementsAttrs to be used by
python API
- Added method to convert between `stablehlo::Tensor` and
`mlir::DenseElementsAttr` more easily (only supports int/float types)
- Added more error checking at interpreter entry point to improve UX
- Changed error reporting at API boundary from `llvm::Error` to
`mlir::LogicalResult` to be more consistent for external users, also
improves python error experience

Author: GleasonK

BUILD.bazel

@@ -431,6 +431,7 @@ cc_library(
         ":reference_ops",
         ":reference_process",
         ":reference_scope",
+        ":reference_tensor",
         ":reference_value",
         ":register",
         "@llvm-project//llvm:Support",

stablehlo/integrations/python/CMakeLists.txt

@@ -85,6 +85,7 @@ declare_mlir_python_extension(StablehloPythonExtensions.Main
     StablehloCAPI
   PRIVATE_LINK_LIBS
     StablehloPortableApi
+    StablehloReferenceApi
     StablehloSerialization
     LLVMSupport
 )

stablehlo/integrations/python/StablehloModule.cpp

@@ -11,14 +11,18 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <vector>
+
 #include "mlir-c/IR.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 #include "mlir/CAPI/IR.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "stablehlo/dialect/Serialization.h"
 #include "stablehlo/integrations/c/StablehloAttributes.h"
 #include "stablehlo/integrations/c/StablehloDialect.h"
 #include "stablehlo/integrations/c/StablehloTypes.h"
 #include "stablehlo/integrations/python/PortableApi.h"
+#include "stablehlo/reference/Api.h"
 
 namespace py = pybind11;
 
@@ -483,6 +487,38 @@ PYBIND11_MODULE(_stablehlo, m) {
   //
   mlir::stablehlo::AddPortableApi(m);
 
+  //
+  // Reference APIs
+  //
+  m.def(
+      "eval_module",
+      [](MlirModule module,
+         std::vector<MlirAttribute> &args) -> std::vector<MlirAttribute> {
+        std::vector<mlir::DenseElementsAttr> inputs;
+        for (auto arg : args) {
+          auto attr = unwrap(arg).dyn_cast<mlir::DenseElementsAttr>();
+          if (!attr) {
+            PyErr_SetString(PyExc_ValueError,
+                            "input args must be DenseElementsAttr");
+            return {};
+          }
+          inputs.push_back(attr);
+        }
+
+        mlir::stablehlo::InterpreterConfiguration config;
+        auto results =
+            mlir::stablehlo::evalModule(unwrap(module), inputs, config);
+        if (failed(results)) {
+          PyErr_SetString(PyExc_ValueError, "interpreter failed");
+          return {};
+        }
+
+        std::vector<MlirAttribute> pyResults;
+        for (auto res : *results) pyResults.push_back(wrap(res));
+        return pyResults;
+      },
+      py::arg("module"), py::arg("args"));
+
   //
   // Serialization APIs.
   //

stablehlo/integrations/python/tests/stablehlo.py

@@ -21,6 +21,7 @@
 import io
 from mlir import ir
 from mlir.dialects import stablehlo
+import numpy as np
 
 
 def run(f):
@@ -227,22 +228,51 @@ def test_minimum_version():
   assert is_semver_format(curr_version)
 
 
-ASM = """
-func.func @test(%arg0: tensor<2xf32>) -> tensor<2xf32> {
-  %0 = stablehlo.add %arg0, %arg0 : (tensor<2xf32>, tensor<2xf32>) -> tensor<2xf32>
-  func.return %0 : tensor<2xf32>
-}
+ASM_FORMAT = """
+func.func @test(%arg0: tensor<{0}>) -> tensor<{0}> {{
+  %0 = stablehlo.add %arg0, %arg0 : (tensor<{0}>, tensor<{0}>) -> tensor<{0}>
+  func.return %0 : tensor<{0}>
+}}
 """
 
+
+@run
+def test_reference_api():
+  # Formatted as (tensor_type, np_value)
+  # Program runs arg + arg, which is used for expected value
+  tests = [
+    # No numpy types for f8 - skipping fp8 tests
+    ("f16", np.asarray(1, np.float16)),
+    ("f32", np.asarray(2, np.float32)),
+    ("f64", np.asarray(3, np.double)),
+    ("1xi8", np.asarray([4], np.int8)),
+    ("1xi16", np.asarray([5], np.int16)),
+    ("1xi32", np.asarray([-6], np.int32)),
+    # Numpy's uint treated as int by DenseElementsAttr, skipping np.uint tests
+    ("2x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2,2)),
+    ("2x1x2xf16", np.asarray([1, 2, 3, 4], np.float16).reshape(2,1,2)),
+  ]
+  for test in tests:
+    tensor_type, arg = test
+    with ir.Context() as context:
+      stablehlo.register_dialect(context)
+      m = ir.Module.parse(ASM_FORMAT.format(tensor_type))
+      args = [ir.DenseIntElementsAttr.get(arg)]
+
+    actual = np.array(stablehlo.eval_module(m, args)[0])
+    expected = arg + arg
+    assert (actual == expected).all()
+
+
 @run
 def test_serialization_apis():
   curr_version = stablehlo.get_current_version()
 
   with ir.Context() as context:
     stablehlo.register_dialect(context)
-    m = ir.Module.parse(ASM)
-    module_str = str(m)
+    m = ir.Module.parse(ASM_FORMAT.format("2xf32"))
     assert m is not None
+    module_str = str(m)
     serialized = stablehlo.serialize_portable_artifact(m, curr_version)
     deserialized = stablehlo.deserialize_portable_artifact(context, serialized)
     assert module_str == str(deserialized)
@@ -258,9 +288,9 @@ def module_to_bytecode(module: ir.Module) -> bytes:
 
   with ir.Context() as context:
     stablehlo.register_dialect(context)
-    m = ir.Module.parse(ASM)
-    module_str = str(m)
+    m = ir.Module.parse(ASM_FORMAT.format("2xf32"))
     assert m is not None
+    module_str = str(m)
     bytecode = module_to_bytecode(m)
     serialized = stablehlo.serialize_portable_artifact(bytecode, curr_version)
     deserialized = stablehlo.deserialize_portable_artifact(serialized)

stablehlo/reference/Api.cpp

@@ -14,15 +14,23 @@ limitations under the License.
 ==============================================================================*/
 #include "stablehlo/reference/Api.h"
 
+#include <cstdint>
+
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/SourceMgr.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectRegistry.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/TypeRange.h"
 #include "mlir/Parser/Parser.h"
+#include "mlir/Support/LogicalResult.h"
 #include "stablehlo/dialect/Register.h"
 #include "stablehlo/reference/Configuration.h"
 #include "stablehlo/reference/Errors.h"
@@ -31,11 +39,13 @@ limitations under the License.
 #include "stablehlo/reference/Ops.h"
 #include "stablehlo/reference/Process.h"
 #include "stablehlo/reference/Scope.h"
+#include "stablehlo/reference/Tensor.h"
+#include "stablehlo/reference/Value.h"
 
 namespace mlir {
 namespace stablehlo {
 namespace {
-func::FuncOp getMainFunction(ModuleOp module, StringRef mainName) {
+FailureOr<func::FuncOp> getMainFunction(ModuleOp module, StringRef mainName) {
   auto functions = module.getOps<func::FuncOp>();
 
   for (auto funcOp : functions)
@@ -46,7 +56,8 @@ func::FuncOp getMainFunction(ModuleOp module, StringRef mainName) {
   bool isDefaultLookup = mainName == "main";
   if (isSingleFunction && isDefaultLookup) return *functions.begin();
 
-  return {};
+  return module.emitError()
+         << "module must have entry func with name " << mainName;
 }
 
 // DefaultInterpreterFallback is an implementation detail of run module. It
@@ -106,33 +117,77 @@ class DefaultInterpreterFallback : public InterpreterFallback {
   int64_t serializedProbeFileId = 0;
 };
 
+LogicalResult validateEntrySignature(func::FuncOp func,
+                                     ArrayRef<InterpreterValue> inputs) {
+  if (func.getNumArguments() != inputs.size())
+    return func->emitError()
+           << "incorrect number of arguments specified, provided "
+           << inputs.size() << " inputs but function expected"
+           << func.getNumArguments();
+
+  TypeRange signature = func.getArgumentTypes();
+  for (int64_t i = 0; i < func.getNumArguments(); ++i) {
+    Type sigType = signature[i];
+    Type argType = inputs[i].getType();
+    if (sigType != argType)
+      return func.emitError() << "invalid input argument type at index " << i
+                              << ", input type was " << argType
+                              << " but entry function expected " << sigType;
+  }
+  return success();
+}
+
 }  // namespace
 
-llvm::ErrorOr<SmallVector<InterpreterValue>> evalModule(
+FailureOr<SmallVector<InterpreterValue>> evalModule(
     ModuleOp module, ArrayRef<InterpreterValue> inputs,
     const InterpreterConfiguration &config) {
+  // Additional error checking at main function boundary.
+  // This is most likely user error, where future errors during interpreting are
+  // more likely invalid IR or interpreter bugs.
   if (module.getOps<func::FuncOp>().empty())
     return SmallVector<InterpreterValue>();
 
   auto mainFunc = getMainFunction(module, config.mainFunction);
-  if (!mainFunc) llvm::report_fatal_error("Requested main function not found.");
+  if (failed(mainFunc) || failed(validateEntrySignature(*mainFunc, inputs)))
+    return failure();
 
   if (!config.probeInstrumentationDir.empty()) {
     llvm::SmallString<128> instrumentationMetadataFile(
         config.probeInstrumentationDir);
     llvm::sys::path::append(instrumentationMetadataFile,
                             stablehlo::numpy::kInstrumentationMetadataFilename);
     if (llvm::sys::fs::remove(instrumentationMetadataFile))
-      llvm::report_fatal_error(
+      return emitError(
+          UnknownLoc::get(module.getContext()),
           "Failed to remove existing instrumentation metadata file.");
   }
 
   DefaultInterpreterFallback fallback(config);
-  return stablehlo::eval(mainFunc.getBody(), inputs, &fallback);
+  return stablehlo::eval(mainFunc->getBody(), inputs, &fallback);
+}
+
+FailureOr<SmallVector<DenseElementsAttr>> evalModule(
+    ModuleOp module, ArrayRef<DenseElementsAttr> inputs,
+    const InterpreterConfiguration &config) {
+  SmallVector<InterpreterValue> valueInputs = llvm::to_vector(
+      llvm::map_range(inputs, [](DenseElementsAttr attr) -> InterpreterValue {
+        return InterpreterValue(makeTensor(attr));
+      }));
+
+  auto values = evalModule(module, valueInputs, config);
+  if (failed(values)) return failure();
+
+  SmallVector<DenseElementsAttr> results = llvm::to_vector(llvm::map_range(
+      values.value(), [](InterpreterValue val) -> DenseElementsAttr {
+        return makeDenseElementsAttr(val.getTensor());
+      }));
+
+  return results;
 }
 
-llvm::ErrorOr<OwningOpRef<ModuleOp>> parseStablehloModule(
-    const std::string &mlir, MLIRContext &context) {
+FailureOr<OwningOpRef<ModuleOp>> parseStablehloModule(const std::string &mlir,
+                                                      MLIRContext &context) {
   llvm::SourceMgr source_mgr;
   source_mgr.AddNewSourceBuffer(llvm::MemoryBuffer::getMemBuffer(mlir),
                                 llvm::SMLoc());
@@ -145,7 +200,8 @@ llvm::ErrorOr<OwningOpRef<ModuleOp>> parseStablehloModule(
   mlir::OwningOpRef<mlir::ModuleOp> module(
       mlir::parseSourceFile<mlir::ModuleOp>(source_mgr, &context));
 
-  if (!module) return llvm::errc::invalid_argument;
+  if (!module)
+    return emitError(UnknownLoc::get(&context), "unable to parse module");
 
   return module;
 }

stablehlo/reference/Api.h

@@ -19,6 +19,7 @@ limitations under the License.
 #include <string>
 
 #include "llvm/Support/ErrorOr.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/OwningOpRef.h"
@@ -33,13 +34,19 @@ namespace stablehlo {
 /// module input and provided inputs. Returns a list of interpreter outputs.
 /// Can optionally pass a fallback interpreter callback which executes when no
 /// builtin kernels are matched.
-llvm::ErrorOr<SmallVector<InterpreterValue>> evalModule(
+FailureOr<SmallVector<InterpreterValue>> evalModule(
     ModuleOp module, ArrayRef<InterpreterValue> inputs,
     const InterpreterConfiguration &config);
 
+/// This wrapper is intended to be easily used by the StableHLO Python bindings.
+// It wraps the InterpreterValue API.
+FailureOr<SmallVector<DenseElementsAttr>> evalModule(
+    ModuleOp module, ArrayRef<DenseElementsAttr> inputs,
+    const InterpreterConfiguration &config);
+
 /// Parses a StableHLO MLIR text program into a ModuleOp.
-llvm::ErrorOr<OwningOpRef<ModuleOp>> parseStablehloModule(
-    const std::string &mlir, MLIRContext &context);
+FailureOr<OwningOpRef<ModuleOp>> parseStablehloModule(const std::string &mlir,
+                                                      MLIRContext &context);
 
 }  // namespace stablehlo
 }  // namespace mlir

stablehlo/reference/CMakeLists.txt

@@ -26,6 +26,7 @@ add_mlir_library(StablehloReferenceApi
   StablehloReferenceOps
   StablehloReferenceProcess
   StablehloReferenceScope
+  StablehloReferenceTensor
   StablehloReferenceValue
   StablehloRegister
 )