Integrate StableHLO at openxla/stablehlo@f4459e76

PiperOrigin-RevId: 619669095
tensorflow · Mar 27, 2024 · cae5b75 · cae5b75
1 parent 0b73f94
commit cae5b75
Show file tree

Hide file tree

Showing 31 changed files with 536 additions and 369 deletions.
diff --git a/stablehlo/BUILD b/stablehlo/BUILD
@@ -338,11 +338,12 @@ gentbl_cc_library(
 )
 
 cc_test(
-    name = "examples-add",
+    name = "example_add",
     srcs = [
-        "examples/c++/stablehlo_add.cpp",
+        "examples/c++/ExampleAdd.cpp",
     ],
     deps = [
+        ":reference_api",
         ":stablehlo_ops",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",

diff --git a/stablehlo/BUILD.bazel b/stablehlo/BUILD.bazel
@@ -275,6 +275,7 @@ cc_library(
         "@llvm-project//mlir:Dialect",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:InferTypeOpInterface",
+        "@llvm-project//mlir:InliningUtils",
         "@llvm-project//mlir:QuantOps",
         "@llvm-project//mlir:TransformUtils",
     ],
@@ -422,6 +423,7 @@ cc_library(
     hdrs = [
         "stablehlo/reference/Api.h",
     ],
+    strip_include_prefix = ".",
     deps = [
         ":interpreter_ops",
         ":reference_configuration",
@@ -1070,6 +1072,7 @@ cc_library(
         "@llvm-project//mlir:FunctionInterfaces",
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:InferTypeOpInterface",
+        "@llvm-project//mlir:InliningUtils",
         "@llvm-project//mlir:QuantOps",
         "@llvm-project//mlir:ShapeDialect",
         "@llvm-project//mlir:SparseTensorDialect",

diff --git a/stablehlo/CMakePresets.json b/stablehlo/CMakePresets.json
@@ -19,6 +19,27 @@
             "CMAKE_EXPORT_COMPILE_COMMANDS": "ON",
             "MLIR_DIR": "${sourceDir}/llvm-build/lib/cmake/mlir"
           }
+        },
+        {
+          "name": "debug-python",
+          "displayName": "Debug w/ python bindings",
+          "inherits": "debug",
+          "cacheVariables": {
+            "STABLEHLO_ENABLE_BINDINGS_PYTHON" : "ON",
+            "STABLEHLO_ENABLE_SANITIZER": "OFF"
+          }
+        }
+      ],
+      "buildPresets": [
+        {
+          "name": "debug",
+          "displayName": "Build Debug",
+          "configurePreset": "debug"
+        },
+        {
+          "name": "debug-python",
+          "displayName": "Build Debug w/ python bindings",
+          "configurePreset": "debug-python"
         }
     ]
 }
diff --git a/stablehlo/WORKSPACE.bazel b/stablehlo/WORKSPACE.bazel
@@ -17,9 +17,9 @@ workspace(name = "stablehlo")
 
 load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 
-LLVM_COMMIT = "a4ca07f13b560b4f6fa5459eef7159e4f9ee9a6b"
+LLVM_COMMIT = "7ac7d418ac2b16fd44789dcf48e2b5d73de3e715"
 
-LLVM_SHA256 = "fb936389d46b3ce7ee423c0d788e5359da8ce41cfe8996847719920c6f60b044"
+LLVM_SHA256 = "8b99a146881fbb2a2d8e812724550b2c88fed4403dfb4e133ee8b7107a6a9348"
 
 http_archive(
     name = "llvm-raw",

diff --git a/stablehlo/docs/spec.md b/stablehlo/docs/spec.md
@@ -39,7 +39,7 @@ func.func @main(
   %0 = "stablehlo.reshape"(%image) : (tensor<28x28xf32>) -> tensor<1x784xf32>
   %1 = "stablehlo.dot"(%0, %weights) : (tensor<1x784xf32>, tensor<784x10xf32>) -> tensor<1x10xf32>
   %2 = "stablehlo.add"(%1, %bias) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
-  %3 = "stablehlo.constant"() { value = dense<0.0> : tensor<1x10xf32> } : () -> tensor<1x10xf32>
+  %3 = "stablehlo.constant"() {value = dense<0.0> : tensor<1x10xf32>} : () -> tensor<1x10xf32>
   %4 = "stablehlo.maximum"(%2, %3) : (tensor<1x10xf32>, tensor<1x10xf32>) -> tensor<1x10xf32>
   "func.return"(%4): (tensor<1x10xf32>) -> ()
 }
@@ -1965,8 +1965,10 @@ semantics change.
 ```mlir
 %results = "stablehlo.composite"(%input0, %input1) {
   name = "my_namespace.my_op",
+  composite_attributes = {
+    my_attribute = "my_value"
+  },
   decomposition = @my_op,
-  composite_attributes = { my_attribute = "my_value" },
   version = 1 : i32
 } : (tensor<f32>, tensor<f32>) -> tensor<f32>
 ```
@@ -3871,7 +3873,7 @@ Semantics of `outfeed_config` is implementation-defined.
 #### Examples
 
 ```mlir
-%result = "stablehlo.outfeed"(%inputs0, %token) {
+%result = "stablehlo.outfeed"(%input0, %token) {
   outfeed_config = ""
 } : (tensor<2x2x2xi64>, !stablehlo.token) -> !stablehlo.token
 ```
@@ -6236,6 +6238,19 @@ At the moment, StableHLO does not provide guarantees about numerical accuracy,
 but this may change in the future
 ([#1156](https://github.com/openxla/stablehlo/issues/1156)).
 
+### Execution semantics of quantized operation
+
+The interpretation of quantized StableHLO operations may vary depending on the
+hardware requirements and capabilities. For instance, some hardware may opt to
+interpret quantized operations using a "dequantize, perform floating-point
+operation, and finally quantize" strategy. Others may perform the entire
+computation with integer arithmetic. Consequently, the interpretation of
+quantized StableHLO operations is exclusively determined by the specific
+implementation. The interpretation of hybrid quantization
+([#1575](https://github.com/openxla/stablehlo/issues/1575)) should be based on
+the it's semantics as prescribed in the specification (via
+[1792](https://github.com/openxla/stablehlo/pull/1792)).
+
 ### Errors
 
 StableHLO programs are validated through an extensive set of constraints for

diff --git a/stablehlo/examples/CMakeLists.txt b/stablehlo/examples/CMakeLists.txt
@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2024 The StableHLO Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,3 +13,7 @@
 # limitations under the License.
 
 add_subdirectory(c++)
+
+if(STABLEHLO_ENABLE_BINDINGS_PYTHON)
+  add_subdirectory(python)
+endif()
diff --git a/stablehlo/examples/c++/BUILD.bazel b/stablehlo/examples/c++/BUILD.bazel
@@ -19,11 +19,12 @@ package(
 )
 
 cc_binary(
-    name = "stablehlo-add",
+    name = "example-add",
     srcs = [
-        "stablehlo_add.cpp",
+        "ExampleAdd.cpp",
     ],
     deps = [
+        "//:reference_api",
         "//:stablehlo_ops",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",
@@ -32,11 +33,12 @@ cc_binary(
 )
 
 cc_test(
-    name = "stablehlo-add_test",
+    name = "example-add-test",
     srcs = [
-        "stablehlo_add.cpp",
+        "ExampleAdd.cpp",
     ],
     deps = [
+        "//:reference_api",
         "//:stablehlo_ops",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",

diff --git a/stablehlo/examples/c++/CMakeLists.txt b/stablehlo/examples/c++/CMakeLists.txt
@@ -14,16 +14,16 @@
 
 # TODO(fzakaria): As we think about adding more tests we should consider
 # making this a function to be DRY.
-add_executable(stablehlo-add stablehlo_add.cpp)
-llvm_update_compile_flags(stablehlo-add)
-target_link_libraries(stablehlo-add PRIVATE StablehloOps)
+add_executable(example-add ExampleAdd.cpp)
+llvm_update_compile_flags(example-add)
+target_link_libraries(example-add PRIVATE StablehloOps StablehloReferenceApi)
 
-mlir_check_all_link_libraries(stablehlo-add)
-add_custom_target(stablehlo-add-test
-  COMMAND stablehlo-add
+mlir_check_all_link_libraries(example-add)
+add_custom_target(example-add-test
+  COMMAND example-add
   WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
-  COMMENT "Executing stablehlo-add to validate it works."
+  COMMENT "Executing example-add to validate it works."
   DEPENDS
-    stablehlo-add
+    example-add
 )
-add_dependencies(check-stablehlo-quick stablehlo-add-test)
+add_dependencies(check-stablehlo-quick example-add-test)
diff --git a/stablehlo/examples/c++/stablehlo_add.cpp → stablehlo/examples/c++/ExampleAdd.cpp b/stablehlo/examples/c++/stablehlo_add.cpp → stablehlo/examples/c++/ExampleAdd.cpp
@@ -14,10 +14,12 @@ limitations under the License.
 ==============================================================================*/
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Quant/QuantOps.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Verifier.h"
 #include "stablehlo/dialect/StablehloOps.h"
+#include "stablehlo/reference/Api.h"
 
 int main() {
   mlir::MLIRContext context;
@@ -31,10 +33,11 @@ int main() {
   module->setName("test_module");
 
   /** create function **/
-  // create function argument and result types
-  mlir::Type arg0 =
+  // create function argument and result types.
+  auto tensorType =
       mlir::RankedTensorType::get({3, 4}, mlir::FloatType::getF32(&context));
-  auto func_type = mlir::FunctionType::get(&context, {arg0, arg0}, {arg0});
+  auto func_type =
+      mlir::FunctionType::get(&context, {tensorType, tensorType}, {tensorType});
 
   // create the function and map arguments.
   llvm::ArrayRef<mlir::NamedAttribute> attrs;
@@ -43,26 +46,34 @@ int main() {
   function.setVisibility(mlir::func::FuncOp::Visibility::Public);
   module->push_back(function);
 
-  // create function block with add operations
+  // create function block with add operations.
   mlir::Block* block = function.addEntryBlock();
   llvm::SmallVector<mlir::Value, 4> arguments(block->args_begin(),
                                               block->args_end());
   mlir::OpBuilder block_builder = mlir::OpBuilder::atBlockEnd(block);
   mlir::Location loc = block_builder.getUnknownLoc();
 
   llvm::SmallVector<mlir::NamedAttribute, 10> attributes;
-  block_builder.create<mlir::stablehlo::AddOp>(loc, arguments, attributes)
-      .getOperation();
-
   mlir::Operation* op =
-      block_builder
-          .create<mlir::stablehlo::AddOp>(loc, arg0, arguments, attributes)
+      block_builder.create<mlir::stablehlo::AddOp>(loc, arguments, attributes)
           .getOperation();
   block_builder.create<mlir::func::ReturnOp>(loc, op->getResult(0));
 
-  // verify the module and dump
+  /** verify and dump the module **/
   assert(mlir::succeeded(mlir::verify(module.get())));
-  module->dump();
 
-  return 0;
+  /* interpret the function "main" with concrete inputs **/
+  auto getConstValue = [&](double val) {
+    return mlir::DenseElementsAttr::get(
+        tensorType,
+        block_builder.getFloatAttr(tensorType.getElementType(), val));
+  };
+
+  auto inputValue1 = getConstValue(10.0);
+  auto inputValue2 = getConstValue(20.0);
+  auto expectedValue = getConstValue(30.0);
+
+  mlir::stablehlo::InterpreterConfiguration config;
+  auto results = evalModule(*module, {inputValue1, inputValue2}, config);
+  return failed(results) || (*results)[0] != expectedValue;
 }
diff --git a/stablehlo/examples/c++/README.md b/stablehlo/examples/c++/README.md
@@ -5,5 +5,41 @@ for how to use StableHLO.
 
 Note: If you have a great example to highlight, we welcome contributions!
 
-* [stablehlo_add](./stablehlo_add.cpp): A simple example that demonstrates
-how to use the StableHLO library to add two numbers.
+* [example-add](./ExampleAdd.cpp): A simple example that demonstrates how to
+  * Use the StableHLO library to add two numbers.
+  * Interpret a StableHLO program using concrete inputs.
+
+```c++
+  // Assume 'module' is an MLIR module with function "main" containing StableHLO
+  // operations.
+  llvm::outs() << "Program:\n " << module << "\n";
+
+  // Create concrete inputs to be used for interpreting "main".
+  auto inputValue1 = mlir::DenseElementsAttr::get(
+      tensorType, block_builder.getFloatAttr(tensorType.getElementType(),
+                                         static_cast<double>(10)));
+  auto inputValue2 = mlir::DenseElementsAttr::get(
+      tensorType, block_builder.getFloatAttr(tensorType.getElementType(),
+                                       static_cast<double>(20)));
+  llvm::outs() << "Inputs: " << inputValue1 << ", " << inputValue2 << "\n";
+
+
+  mlir::stablehlo::InterpreterConfiguration config;
+  auto results = evalModule(module, {inputValue1, inputValue2}, config);
+  llvm::outs() << "Output: " << (*results)[0];
+```
+
+Output:
+
+```mlir
+Program:
+module @test_module {
+  func.func @main(%arg0: tensor<3x4xf32>, %arg1: tensor<3x4xf32>) -> tensor<3x4xf32> {
+    %0 = stablehlo.add %arg0, %arg1 : tensor<3x4xf32>
+    %1 = stablehlo.add %arg0, %arg1 : tensor<3x4xf32>
+    return %1 : tensor<3x4xf32>
+  }
+}
+Inputs: dense<1.000000e+01> : tensor<3x4xf32>, dense<2.000000e+01> : tensor<3x4xf32>
+Output: dense<3.000000e+01> : tensor<3x4xf32>
+```
diff --git a/stablehlo/examples/python/CMakeLists.txt b/stablehlo/examples/python/CMakeLists.txt
@@ -0,0 +1,15 @@
+# Copyright 2024 The StableHLO Authors. All Rights Reserved.
+#
+# 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
+#
+#      https://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.
+
+add_stablehlo_python_test(stablehlo-python-add stablehlo_add.py)
diff --git a/stablehlo/examples/python/README.md b/stablehlo/examples/python/README.md
@@ -0,0 +1,9 @@
+# Python Examples
+
+This directory should contain a series of examples as a starting point
+for how to use StableHLO via our Python bindings.
+
+Note: If you have a great example to highlight, we welcome contributions!
+
+* [stablehlo_add](./stablehlo_add.py): A simple example that demonstrates
+how to use the StableHLO library to add two numbers.
diff --git a/stablehlo/examples/python/stablehlo_add.py b/stablehlo/examples/python/stablehlo_add.py
@@ -0,0 +1,37 @@
+# Copyright 2024 The StableHLO Authors. All Rights Reserved.
+#
+# 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
+#
+#      https://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.
+
+from mlir import ir
+import mlir.dialects.stablehlo as stablehlo
+import mlir.dialects.func as func
+from mlir.ir import Context, Location, InsertionPoint, Module
+import numpy as np
+
+with Context() as ctx, Location.unknown():
+  stablehlo.register_dialect(ctx)
+  module = Module.create()
+
+  with InsertionPoint(module.body):
+
+    @func.func()
+    def main():
+      a_value = ir.DenseElementsAttr.get(np.zeros(shape=[3,4], dtype=np.int64))
+      b_value = ir.DenseElementsAttr.get(np.zeros(shape=[3,4], dtype=np.int64))
+      a = stablehlo.constant(a_value)
+      b = stablehlo.constant(b_value)
+      add = stablehlo.add(a, b)
+      return add
+
+assert main.func_op.verify()
+print(str(module))
diff --git a/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.cpp b/stablehlo/stablehlo/conversions/linalg/transforms/LegalizeToLinalgUtils.cpp
@@ -73,7 +73,7 @@ Value getEmptyTensorFor(OpBuilder &b, Location loc, ShapedType resultType,
   // new tensor initialization operation. This operation only needs the
   // dynamic sizes.
   SmallVector<Value> sizes;
-  if (resultType.hasRank() && !resultType.hasStaticShape()) {
+  if (!resultType.hasStaticShape()) {
     // Ask the op for its output shape.
     auto shapeSource = cast<InferShapedTypeOpInterface>(op);
     SmallVector<Value, 1> reifiedShapes;