Skip to content

Commit

Permalink
Skeleton of the StableHLO spec (#23)
Browse files Browse the repository at this point in the history
Add a skeleton of the StableHLO spec

Here is the skeleton of the StableHLO specification: Presented as a draft to signify that this is a preliminary version of the specification and will follow additions in future w.r.t different ops.
  • Loading branch information
sdasgup3 authored Aug 24, 2022
1 parent dcdcfe4 commit 2725978
Showing 1 changed file with 129 additions and 0 deletions.
129 changes: 129 additions & 0 deletions docs/spec_draft.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,129 @@
# StableHLO Specification Draft

## Types

Following are the supported element types in StableHLO:

* **Integer types**

* Signed integer with two’s complement representation. Referred to in the
document as `si<N>`, where the bit-width N ∊ {4, 8, 16, 32, 64}
* Unsigned integer referred to in the document as `ui<N>`, where the
bit-width N ∊ {4, 8, 16, 32, 64}
* **Boolean types** referred to in the document as `pred`
* **Floating-point types**
* Single precision `f32`, double precision `f64` and half precision `f16`
floating-points complying with [IEEE 754
format](https://ieeexplore.ieee.org/document/8766229).
* Bfloat16 `bf16` floating-point complying with [Brain Floating-Point Format](https://cloud.google.com/blog/products/ai-machine-learning/bfloat16-the-secret-to-high-performance-on-cloud-tpus).
Provides the same number of exponent bits as `f32`, so that it matches its
dynamic range, but with greatly reduced precision.
* **Complex types** represents a pair of floating-point types. Supported ones
are `c64` (represents paired `f32`) and `c128` (represents paired `f64`).

StableHLO supports a Tensor type `tensor`, to model the type of a n-dimensional
array, represented in the opset as `tensor<SxE>` such that

* Shape `S` is a list of number of elements in each of the
dimensions and represented, in increasing order of the corresponding dimension
number, as an array of values of type `ui64`. A 0-value represents missing or

This comment has been minimized.

Copy link
@jpienaar

jpienaar Aug 25, 2022

Member

So StableHLO doesn't support unknown sizes? (Missing is confusing for 0 extent)

This comment has been minimized.

Copy link
@sdasgup3

sdasgup3 Aug 25, 2022

Author Member

StableHLO will support dynamic shapes and the spec with also evolve with that.

(Missing is confusing for 0 extent)

Point noted. I will clarify this is the next version.

This comment has been minimized.

Copy link
@sdasgup3

sdasgup3 Aug 29, 2022

Author Member

Please find the clarification
#52

empty data in the corresponding dimension.
* Element type `E` is any one of the supported element types mentioned above.

## Programs

StableHLO programs consist of functions. Each function has arguments and results
of supported types and a list of ops in static single-assignment (SSA) form
which is terminated by a return op which produces the results of the function.
StableHLO ops take operands and produce results.

```
ml_program.func @example_func(%arg: tensor<4x16xf32>) -> tensor<4x16xf32> {
%1 = stablehlo.floor %arg : tensor<4x16xf32>
%2 = stablehlo.ceil %arg : tensor<4x16xf32>
%3 = stablehlo.add %1, %2 : tensor<4x16xf32>
ml_program.return %3 : tensor<4x16xf32>
}
```

A program is executed by passing argument values to a given function and
computing result values. Result values of a function are computed by evaluating
the graph of ops rooted in the corresponding return op. The evaluation order is
implementation-defined, as long as ops are evaluated before their uses. Possible
execution orders of the above example program are `%1``%2``%3``return`
or `%2``%1``%3``return`.

## Constants

The section describes the constants supported in StableHLO along with their
syntax.

* **Integer Constants** Standard integers, e.g. `123`, are constants of the
integer type (signed or unsigned). Negative numbers can be used with signed
integer types.
* **Boolean Constants** `true` and `false` are both valid constants of the
`pred` type.
* **Floating-point Constants** Floating-point constants use standard decimal
notation, e.g. `123.421`, exponential notation, e.g. `1.23421e+2`, or a more
precise hexadecimal notation, e.g. `0x42f6d78d`.
* **Complex Constants** Complex constants are represented as a pair of real
and imaginary values of `f32` or `f64` types, e.g. `(12.34, 56,78)`.

## Structure of an Op’s Specification

This comment has been minimized.

Copy link
@jpienaar

jpienaar Aug 25, 2022

Member

These seem to be very similar to what is autogenerated for a dialect, having authoritative source would be good and automatic generation of docs (e.g., https://www.tensorflow.org/mlir/hlo_ops that is generated today) from the dialect is appealing to ensure up to date. At the top level here one could just have an include and for web page generation inline it.

This comment has been minimized.

Copy link
@sdasgup3

sdasgup3 Aug 25, 2022

Author Member

Point taken, In fact, we also have plans, along the same lines, to have a single source of truth, which very well could be the ODS .td file.

I note that this is draft specification which helped us to "Bootstrap StablehloOps.td from MHLO" and "Document Specification" independently.


The specification of an op comprises of the following components (in the order
described below)

* **Syntax** Operation mnemonic and its signature.
* **Semantics** Semantics of the operation.
* **Operands** Meaning of operand(s) and their type(s).
* **Results** Meaning of the result(s) and the type(s).
* **Constraints** Type constraints on the operand(s), result(s).
* **Examples** Examples demonstrating the working of the op.

## stablehlo.add

`stablehlo.add(lhs, rhs) -> result`

### Semantics

Performs element-wise addition of two tensors `lhs` and `rhs` and produces a
`result` tensor. For integer element types, if the element-wise sum has an
unsigned/signed overflow/underflow, the result is implementation defined and one
of the followings:

* mathematical result modulo $2^n$, where n is the bit width of the result.
* saturation to $2^{n-1} - 1$ (or $-2^{n-1}$) for signed overflow (or signed
underflow) and saturation to $2^n - 1$ (or $0$) for unsigned overflow (or
unsigned underflow).

For floating-point element types, corner cases are defined by the IEEE-754
specification.

### Arguments

| Operand Name(s) | Type |
|-|-|
| `lhs` | `tensor` of Integer, Floating-point and Complex element types |
| `rhs` | `tensor` of Integer, Floating-point and Complex element types |

### Results

| Result Name | Type |
|-|-|
| `result` | `tensor` of Integer, Floating-point, and Complex element types |

### Constraints

* `lhs`, `rhs` have the same type.
* Supported shapes: all static shapes.
* `result` must have the same type as that of `lhs` (or `rhs`).

### Examples

```c
// %x: [[1, 2], [3, 4]]
// %y: [[5, 6], [7, 8]]
%z = stablehlo.add %x, %y : tensor<2x2xf32>
// %z: [[6, 8], [10, 12]]
```

0 comments on commit 2725978

Please sign in to comment.