## XLA Implementation Status

List of XLA(mhlo-mlir) Ops that SPU supports:

    The list of mhlo ops is obtained from this file:
        https://github.com/tensorflow/mlir-hlo/blob/master/include/mlir-hlo/Dialect/mhlo/IR/hlo_ops.td

    General limitation with SPU:
        * Dynamic shape is not supported
        * Complex number is not supported
        * SPU only supports fixed-point numbers, so no-finite is not supported

### XLA nullary ops

|  Op Name       | supported(fully/partial/no) | notes       |
| :------------: | :-------------------------: | ----------- |
| `constant`     | fully                       | Always yields a public value
| `iota`         | fully                       | Always yields a public value
| `dynamic_iota` | no                          |
| `create_token` | no                          |

Count: Total = 4, fully supported = 2

### XLA unary element-wise ops

|  Op Name       | supported(fully/partial/no) | notes       |
| :------------: | :-------------------------: | ----------- |
| `abs`          | fully                       |
| `cbrt`         | no                          |
| `ceil`         | fully                       |
| `convert`      | fully                       |
| `count_leading_zeros`| no                    |
| `cosine`       | no                          |
| `exponential`  | fully                       |
| `exponential_minus_one`| fully               |
| `floor`        | fully                       |
| `imag`         | no                          |
| `is_finite`    | no                          |
| `log`          | fully                       |
| `log_plus_one` | fully                       |
| `logistic`     | fully                       |
| `not`          | fully                       |
| `negate`       | fully                       |
| `popcnt`       | not                         |
| `real`         | not                         |
| `round_nearest_afz`| not                     |
| `rsqrt`        | fully                       |
| `sign`         | partial                     |
| `sine`         | not                         |
| `sqrt`         | fully                       |
| `tanh`         | fully                       |

Count: Total = 24, fully supported = 12, partial = 0

### XLA binary element-wise ops

|  Op Name       | supported(fully/partial/no) | notes       |
| :------------: | :-------------------------: | ----------- |
| `add`          | fully                       |
| `atan2`        | no                          |
| `complex`      | no                          |
| `compare`      | fully                       |
| `divide`       | fully                       |
| `maximum`      | fully                       |
| `minimum`      | fully                       |
| `multiply`     | fully                       |
| `power`        | fully                       |
| `remainder`    | fully                       |
| `shift_left`   | partial                     | rhs must be a public
| `shift_right_arithmetic` | partial           | rhs must be a public
| `shift_right_logical` | partial              | rhs must be a public
| `subtract`     | fully                       |

Count: Total = 14, fully supported = 9, partial = 3

### XLA binary logical element-wise ops

|  Op Name       | supported(fully/partial/no) | notes       |
| :------------: | :-------------------------: | ----------- |
| `and`          | fully                       |
| `or`           | fully                       |
| `xor`          | fully                       |

Count: Total = 3, fully supported = 3

### XLA communication ops

|  Op Name       | supported(fully/partial/no) | notes       |
| :------------: | :-------------------------: | ----------- |
| `infeed`       | no                          |
| `outfeed`      | no                          |
| `send`         | no                          |
| `recv`         | no                          |

Count: Total = 4, fully supported = 0

### XLA parallelism related ops

|  Op Name       | supported(fully/partial/no) | notes       |
| :------------: | :-------------------------: | ----------- |
| `replica_id`   | no                          |

Count: Total = 1, fully supported = 0

### XLA control flow ops

|  Op Name       | supported(fully/partial/no) | notes       |
| :------------: | :-------------------------: | ----------- |
| `after_all`    | no                          |
| `if`           | partial                     | condition variable must be a public scalar
| `case`         | no                          |
| `while`        | partial                     | condition region must return a public scalar
| `all_gather`   | no                          |
| `all_reduce`   | no                          |
| `reduce_scatter` | no                        |
| `all_to_all`   | no                          |
| `reduce`       | fully                       | inherits limitations from reduce function

Count: Total = 9, fully supported = 1, partial = 2

### XLA tuple ops

|  Op Name       | supported(fully/partial/no) | notes       |
| :------------: | :-------------------------: | ----------- |
| `get_tuple_element` | fully                  |
| `tuple`        | fully                       |

Count: Total = 2, fully supported = 2

### XLA other ops

|  Op Name       | supported(fully/partial/no) | notes       |
| :------------: | :-------------------------: | ----------- |
| `slice`        | fully                       |
| `dynamic-slice`| partial                     | start_indices must be public values
| `dynamic-update-slice`| partial              | start_indices must be public values
| `batch_norm_grad`| fully                     | Rely on XLA's batchnorm_expander pass
| `batch_norm_inference`| fully                | Rely on XLA's batchnorm_expander pass
| `batch_norm_training` | fully                | Rely on XLA's batchnorm_expander pass
| `bitcast_convert` | partial                  | Only supports convert to type of same size
| `broadcast`    | fully                       |
| `broadcast_in_dim` | fully                   |
| `dynamic_broadcast_in_dim` | no              |
| `cholesky`     | fully                       | Rely on CholeskyExpander pass
| `clamp`        | fully                       |
| `concatenate`  | fully                       |
| `collective_permute` | no                    |
| `convolution`  | fully                       | 
| `copy`         | no                          |
| `cross-replica-sum` | no                     |
| `custom_call`  | no                          |
| `dot`          | fully                       | 
| `dot_general`  | fully                       |
| `einsum`       | fully                       |
| `unary_einsum` | fully                       |
| `fft`          | no                          |
| `gather`       | fully                       |
| `get_dimension_size` | no                    |
| `map`          | fully                       | Rely on XLA's MapInliner pass
| `reshape`      | fully                       |
| `dynamic_reshape` | no                       |
| `scatter`      | no                          |
| `select`       | fully                       |
| `select_and_scatter` | fully                 |
| `set_dimension_size` | no                    |
| `sort`         | fully                       |
| `reverse`      | fully                       |
| `pad`          | fully                       |
| `trace`        | no                          |
| `transpose`    | fully                       |
| `triangular_solve` | fully                   | Rely on XLA's TriangularSolverExpander pass
| `reduce_window`| fully                       |
| `return`       | fully                       |
| `torch_index_select` | no                    |
| `optimization_barrier` | no                  |

Count: Total = 42, fully supported = 26, partial = 3

### XLA RNG ops

|  Op Name       | supported(fully/partial/no) | notes       |
| :------------: | :-------------------------: | ----------- |
| `rng_uniform`  | partial                     | Bound [a, b) must all be public scalar, result is also a public tensor
| `rng_normal`   | no                          |
| `rng_bit_generator` | no                     |

Count: Total = 3, fully supported = 0, partial = 1

### XLA quantize op

|  Op Name       | supported(fully/partial/no) | notes       |
| :------------: | :-------------------------: | ----------- |
| `dequantize`   | no                          |

Count: Total = 1, fully supported = 0, partial = 0

### XLA miscellaneous ops

|  Op Name       | supported(fully/partial/no) | notes       |
| :------------: | :-------------------------: | ----------- |
| `fusion`       | no                          |
| `bitcast`      | no                          | Internal op to XLA/GPU
| `reduce_precision` | no                      |
| `real_dynamic_slice` | no                    |
| `dynamic_pad`  | no                          |
| `dynamic_gather` | no                        |
| `dynamic_conv` | no                          |
| `print`        | no                          |
| `compute_reshape_shape` | no                 |
| `cstr_reshapable` | no                       |

Count: Total = 10, fully supported = 0, partial = 0