Add interpreter for ScatterOp (#1488)
We have the following constraints in the spec:
```
(I1) `inputs` variadic number of tensors.
(I2) `scatter_indices` tensor of integer type.
(I3) `updates` variadic number of tensors.
(I4) `update_window_dims` 1-dimensional tensor constant of type `si64`.
(I5) `inserted_window_dims` 1-dimensional tensor constant of type `si64`.
(I6) `scatter_dims_to_operand_dims` 1-dimensional tensor constant of type `si64`.
(I7) `index_vector_dim` constant of type `si64`.
(I8) `indices_are_sorted` constant of type `i1`.
(I9) `unique_indices` constant of type `i1`.
(I10) `update_computation` function.
(C1) All `inputs` have the same shape.
(C2) rank(`inputs`[0]) = size(`update_window_dims`) +
size(`inserted_window_dims`).
(C3) All `updates` have the same shape.
(C4) `shape(updates[0])` =
`combine(update_scatter_dim_sizes, update_window_dim_sizes)` where:
* `update_scatter_dim_sizes` = `shape(scatter_indices)` except that
the dimension size of `scatter_indices` corresponding to
`index_vector_dim` is not included.
* `update_window_dim_sizes` <= `shape(inputs[0])` except that
the dimension sizes in `inputs[0]` corresponding to `inserted_window_dims`
are not included.
* `combine` puts `update_scatter_dim_sizes` at axes corresponding to
`update_scatter_dims` and `update_window_dim_sizes` at axes corresponding
to `update_window_dims`.
(C5) N = size(`inputs`) = size(`updates`) and N >= 1.
(C6) `element_type(updates[k]) = element_type(inputs[k])` for all k $\in$
[0, N).
(C7) All dimensions in `update_window_dims` are unique and sorted.
(C8) For all i in [0, size(`update_window_dims`)), 0 <=
`update_window_dims`[i] < rank(`updates`[0]).
(C9) All dimensions in `inserted_window_dims` are unique and sorted.
(C10) For all i in [0, size(`inserted_window_dims`)), 0 <=
`inserted_window_dims`[i] < rank(`inputs`[0]).
(C11) size(`scatter_dims_to_operand_dims`) =
`index_vector_dim` < rank(`scatter_indices`) ?
dim(`scatter_indices`, `index_vector_dim`) : 1.
(C12) All dimensions in `scatter_dims_to_operand_dims` are unique.
(C13) For all i in [0, size(`scatter_dims_to_operand_dims`)), 0 <=
`scatter_dims_to_operand_dims`[i] < rank(`inputs`[0]).
(C14) 0 <= `index_vector_dim` <= rank(`scatter_indices`).
(C15) `update_computation` has type
`(tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>)`
where `Ek = element_type(inputs[k])` for all k in [0, N).
(C16) `inputs[k]` and `result[k]` have the same type for all k in [0, N).
```
These constraints will be comprehensively covered by the following
tests:
```
I1: a) `inputs` variadic number of tensors. (Covered by ODS).
I2: a) `scatter_indices` tensor of integer type. (Covered by ODS).
I3: a) `updates` variadic number of tensors. (Covered by ODS).
I4: a) `update_window_dims` 1-dimensional tensor constant of type `si64`. (Covered by ODS).
I5: a) `inserted_window_dims` 1-dimensional tensor constant of type `si64`. (Covered by ODS).
I6: a) `scatter_dims_to_operand_dims` 1-dimensional tensor constant of type `si64`. (Covered by ODS).
I7: a) `index_vector_dim` constant of type `si64`. (Covered by ODS).
I8: a) `indices_are_sorted` constant of type `i1`. (Covered by ODS).
I9: a) `unique_indices` constant of type `i1`. (Covered by ODS).
I10: a) `update_computation` function. (Covered by ODS).
C1: a) Not all `inputs` have the same shape.
C2: a) rank(`inputs`[0]) != size(`update_window_dims`) +
size(`inserted_window_dims`).
C3: a) Not all `updates` have the same shape.
C4: a) `shape(updates[0])` !=
`combine(update_scatter_dim_sizes, update_window_dim_sizes)`.
b) `update_scatter_dim_sizes` != `shape(scatter_indices)` except that the
dimension size of `scatter_indices` corresponding to `index_vector_dim`
is not included.
c) `update_window_dim_sizes` <= `shape(inputs[0])` except that the dimension
sizes in `inputs[0]` corresponding to `inserted_window_dims` are not
included.
where `combine` puts `update_scatter_dim_sizes` at axes corresponding to
`update_scatter_dims` and `update_window_dim_sizes` at axes corresponding to
`update_window_dims`.
C5: a) N != size(`inputs`). (Covered by ODS).
b) N != size(`updates`). (Covered by ODS).
c) N < 1. (Covered by ODS).
C6: a) `element_type(updates[k]) != element_type(inputs[k])` for any k in [0, N).
C7: a) Dimensions in `update_window_dims` are not unique.
b) Dimensions in `update_window_dims` are not sorted.
C8: a) For any i in [0, size(`update_window_dims`)), `update_window_dims`[i] < 0.
b) For any i in [0, size(`update_window_dims`)), `update_window_dims`[i] >= rank(`updates`[0]).
C9: a) Dimensions in `inserted_window_dims` are not unique.
b) Dimensions in `inserted_window_dims` are not sorted.
C10: a) For any i in [0, size(`inserted_window_dims`)), `inserted_window_dims`[i] < 0.
b) For any i in [0, size(`inserted_window_dims`)), >= rank(`inputs`[0]).
C11: a) size(`scatter_dims_to_operand_dims`) !=
`index_vector_dim` < rank(`scatter_indices`) ?
dim(`scatter_indices`, `index_vector_dim`) : 1.
C12: a) Dimensions in `scatter_dims_to_operand_dims` are not unique.
C13: a) For any i in [0, size(`scatter_dims_to_operand_dims`)), `scatter_dims_to_operand_dims`[i] < 0.
b) For any i in [0, size(`scatter_dims_to_operand_dims`)), `scatter_dims_to_operand_dims`[i] >= rank(`inputs`[0]).
C14: a) `index_vector_dim` < 0.
b) `index_vector_dim` > rank(`scatter_indices`).
C15: a) `update_computation` does not have type
`(tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>)`
where `Ek = element_type(inputs[k])` for any k $\in$ [0, N).
C16: a) type(`inputs[k]`) != type(`result[k]`) for any k $\in$ [0, N).
```
If we drop the "Covered by ODS" pieces, this will leave us with the
following test cases:
```
C1a: Not all `inputs` have the same shape.
C2a: rank(`inputs`[0]) != size(`update_window_dims`) + size(`inserted_window_dims`).
C3a: Not all `updates` have the same shape.
C4a: `shape(updates[0])` !=
`combine(update_scatter_dim_sizes, update_window_dim_sizes)`.
C4b: `update_scatter_dim_sizes` != `shape(scatter_indices)` except that the
dimension size of `scatter_indices` corresponding to `index_vector_dim`
is not included.
C4c: `update_window_dim_sizes` <= `shape(inputs[0])` except that the dimension
sizes in `inputs[0]` corresponding to `inserted_window_dims` are not
included.
where `combine` puts `update_scatter_dim_sizes` at axes corresponding to
`update_scatter_dims` and `update_window_dim_sizes` at axes corresponding to
`update_window_dims`.
C6a: `element_type(updates[k]) != element_type(inputs[k])` for any k in [0, N).
C7a: Dimensions in `update_window_dims` are not unique.
C7b: Dimensions in `update_window_dims` are not sorted.
C8a: For any i in [0, size(`update_window_dims`)), `update_window_dims`[i] < 0.
C8b: For any i in [0, size(`update_window_dims`)), `update_window_dims`[i] >= rank(`updates`[0]).
C9a: Dimensions in `inserted_window_dims` are not unique.
C9b: Dimensions in `inserted_window_dims` are not sorted.
C10a: For any i in [0, size(`inserted_window_dims`)), `inserted_window_dims`[i] < 0.
C10b: For any i in [0, size(`inserted_window_dims`)), >= rank(`inputs`[0]).
C11a: size(`scatter_dims_to_operand_dims`) !=
`index_vector_dim` < rank(`scatter_indices`) ?
dim(`scatter_indices`, `index_vector_dim`) : 1.
C12a: Dimensions in `scatter_dims_to_operand_dims` are not unique.
C13a: For any i in [0, size(`scatter_dims_to_operand_dims`)), `scatter_dims_to_operand_dims`[i] < 0.
C13b: For any i in [0, size(`scatter_dims_to_operand_dims`)), `scatter_dims_to_operand_dims`[i] >= rank(`inputs`[0]).
C14a: `index_vector_dim` < 0.
C14b: `index_vector_dim` > rank(`scatter_indices`).
C15a: `update_computation` does not have type
`(tensor<E0>, ..., tensor<EN-1>, tensor<E0>, ..., tensor<EN-1>) -> (tensor<E0>, ..., tensor<EN-1>)`
where `Ek = element_type(inputs[k])` for any k $\in$ [0, N).
C16a: type(`inputs[k]`) != type(`result[k]`) for any k $\in$ [0, N).
```
Notes:
* Some missing verifications were added.
* Updates typo in spec wording (i.e. For any k -> For all k).
* Updates notation `do` -> `di`, `ds` -> `dj`
closes #987StableHLO is an operation set for high-level operations (HLO) in machine learning (ML) models. Essentially, it's a portability layer between different ML frameworks and ML compilers: ML frameworks that produce StableHLO programs are compatible with ML compilers that consume StableHLO programs.
Our goal is to simplify and accelerate ML development by creating more interoperability between various ML frameworks (such as TensorFlow, JAX and PyTorch) and ML compilers (such as XLA and IREE).
StableHLO is based on the MHLO dialect and enhances it with additional functionality, including serialization and versioning. We use MLIR bytecode as serialization format and provide backward and forward compatibility guarantees. This ensures compatibility between frameworks and compilers, even as StableHLO continues to evolve.
This repository includes the StableHLO specification along with an MLIR-based implementation in C++ and Python, which you can use to define StableHLO programs for consumption by compilers such as XLA and IREE.
Here's how to build the StableHLO repo on Linux or macOS:
CMake is our primary build tool, so before you begin make sure that you have CMake and Ninja installed.
If you're using Linux, we recommend installing lld as well - we have observed it to be noticeably faster than alternatives on our typical software and hardware configurations.
# On Linux sudo apt install cmake ninja-build lld # On macOS brew install cmake ninja
Set the LLVM_ENABLE_LLD shell variable depending on your preferences. We recommend setting it to ON on Linux and to OFF on macOS.
[[ "$(uname)" != "Darwin" ]] && LLVM_ENABLE_LLD="ON" || LLVM_ENABLE_LLD="OFF"
Clone the StableHLO repo and the LLVM repository:
git clone https://github.com/openxla/stablehlo
cd stablehlo && git clone https://github.com/llvm/llvm-project.git
Cloning the LLVM repository may take a few minutes.
Make sure you check out the correct commit in the LLVM repository:
(cd llvm-project && git fetch && git checkout $(cat ../build_tools/llvm_version.txt))
You need to do this every time llvm_version.txt changes.
Configure and build MLIR:
build_tools/build_mlir.sh ${PWD}/llvm-project/ ${PWD}/llvm-build
This will take a considerable amount of time. For example, on a MacBook Pro with an M1 Pro chip, building MLIR took around 10 minutes at the moment of writing.
Again, you need to do this every time llvm_version.txt changes.
Build StableHLO as a standalone library:
mkdir -p build && cd build cmake .. -GNinja \ -DLLVM_ENABLE_LLD="$LLVM_ENABLE_LLD" \ -DCMAKE_BUILD_TYPE=Release \ -DLLVM_ENABLE_ASSERTIONS=On \ -DMLIR_DIR=${PWD}/../llvm-build/lib/cmake/mlir
Now you can make sure it works by running some tests:
ninja check-stablehlo-tests
You should see results like this:
Testing Time: 5.99s
Passed: 47
This runs all the tests in stablehlo/tests/.
Building an amazing portability layer between ML frameworks and ML compilers requires collaboration across the whole ML industry, so we're happy to have your help on the StableHLO project.
We're using GitHub issues / pull requests to organize development and openxla-discuss to have longer discussions. We also have a #stablehlo channel on the OpenXLA Discord server.