Add spec for CustomCallOp

docs/spec_draft.md

@@ -393,6 +393,7 @@ syntax.
    * [constant](#stablehloconstant)
    * [cosine](#stablehlocosine)
    * [count_leading_zeros](#stablehlocount_leading_zeros)
+   * [custom_call](#stablehlocustom_call)
    * [divide](#stablehlodivide)
    * [dynamic_slice](#stablehlodynamic_slice)
    * [dynamic_update_slice](#stablehlodynamic_update_slice)
@@ -1798,6 +1799,92 @@ tensor and produces a `result` tensor.
 
 [Back to Ops](#index-of-ops)
 
+## stablehlo.custom_call
+
+### Semantics
+
+Invokes exteral code with `inputs` as arguments and produces `results`. The
+exact mechanism is backend-specific. The inputs are described as follows:
+  * `inputs` are passed as agruments to the external API.
+  * `call_target_name` may be used as an identifier for the external code.
+  * `has_side_effect` is true if the custom call has side-effects.
+  * `backend_config` may be used to encode additional backend-specific
+    information.
+  * `api_version` specifies the version of the custom call API.
+  * `called_computations` used to apply functions within the scope of the parent
+    module.
+  * `operand_layouts` specifies the layout of the operand.
+  * `result_layouts` specifies the layout of the result.
+  * `output_operand_aliases` captures the alias relationship of `results` to one of the operands.
+    * `operand_index` denotes the index of the operand `inputs` that relates to `results`.
+    * `operand_tuple_indices` are used to index into the operand at `operand_index`, if it is a tuple/nested tuple.
+    * `output_tuple_indices` are used to index into `results`, if it is a tuple/nested tuple.
+
+### Inputs
+
+| Name                    | Type                                                                |
+|-------------------------|---------------------------------------------------------------------|
+| `inputs`                | variadic number of tensors of any supported type, tokens, or tuples |
+| `call_target_name`      | constant of type `string`                                           |
+| `has_side_effect`       | constant of type `i1`                                               |
+| `backend_config`        | constant of type `string`                                           |
+| `api_version`           | enum of `API_VERSION_ORIGINAL`, `API_VERSION_STATUS_RETURNING`,     |
+|                         | and `API_VERSION_STATUS_RETURNING_UNIFIED`                          |
+| `called_computations`   | variadic number of `function` references                            |
+| `operand_layouts`       | variadic number of 1-dimensional tensor constants of type `si64`    |
+| `result_layouts`        | variadic number of 1-dimensional tensor constants of type `si64`    |
+| `output_tuple_indices`  | 1-dimensional tensor constant of type `si64`                        |
+| `operand_index`         | constant of type `si64`                                             |
+| `operand_tuple_indices` | 1-dimensional tensor constant of type `si64`                        |
+
+### Outputs
+
+| Name      | Type                                                                |
+|-----------|---------------------------------------------------------------------|
+| `results` | variadic number of tensors of any supported type, tokens, or tuples |
+
+### Constraints
+
+  * (C1) `operand_layouts` & `result_layouts` attributes can be specified under
+          the following constraints:
+    * Either both `operand_layouts` and `result_layouts` are specified or none.
+    * None of the `inputs` are of tuple type.
+    * None of the `results` are of tuple type except the common case of single
+      tuple result packing non-tuple values is allowed.
+    * `operand_layouts[i]` = `layout(inputs[i])`, for all `i`.
+    * `result_layouts[i]` = `layout(results[i])`, for all `i`.
+    * Only tensor types can have non-empty layout.
+    * Layout of an `N`-ranked tensor must be a permutation of `[0, N)`.
+  * (C2) `output_operand_aliases` can be specified under the following
+         constraints:
+    * 0 $\le$ `operand_index` $\lt$ `size(inputs)`.
+    * `type(inputs[operand_index])` = `tuple`.
+    * `shape(inputs[operand_index][i0][i1]..[iN-1])` =
+      `shape(results[j0][j1]..[jM-1]])`, where,
+      `i0, i1, ..` = `operand_tuple_indices[0], operand_tuple_indices[1], ..`,
+      `j0, j1, ..` = `operand_tuple_indices[0], operand_tuple_indices[1], ..`,
+      `size(operand_tuple_indices) = N`, and `size(output_tuple_indices) = M`.
+
+### Examples
+
+```mlir
+%results = "stablehlo.custom_call"(%inputs0) {
+  call_target_name = "foo",
+  has_side_effect = false,
+  backend_config = "bar",
+  api_version = 1 : i32,
+  called_computations = [@foo],
+  operand_layouts = [dense<> : tensor<0xindex>],
+  result_layouts = [dense<> : tensor<0xindex>],
+  output_operand_aliases = [
+    #stablehlo.output_operand_alias<output_tuple_indices = [],
+                                    operand_index = 0,
+                                    operand_tuple_indices = []>]
+} : (tensor<f32>) -> tensor<f32>
+```
+
+[Back to Ops](#index-of-ops)
+
 ## stablehlo.divide
 
 ### Semantics

docs/status.md

@@ -74,7 +74,7 @@ one of the following tracking labels.
 | create_token             | no            | yes*         | yes*           | yes             | no          |
 | cross-replica-sum        | no            | revisit      | revisit        | no              | no          |
 | cstr_reshapable          | no            | revisit      | no             | yes             | no          |
-| custom_call              | no            | revisit      | infeasible     | yes             | no          |
+| custom_call              | yes           | yes          | infeasible     | yes             | no          |
 | divide                   | yes           | yes          | yes            | yes             | no          |
 | dot                      | no            | revisit      | revisit        | yes             | no          |
 | dot_general              | no            | yes*         | yes*           | no              | no          |

stablehlo/dialect/StablehloOps.td

@@ -2086,44 +2086,26 @@ def StableHLO_CustomCallOp: StableHLO_Op<"custom_call",
     [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>]> {
   let summary = "CustomCall operator";
   let description = [{
-    A custom call invokes code external to XLA. The `args` are passed to the
-    external code, and the external code is expected to produce a result of the
-    given type. The exact mechanism is backend-specific. For example, in the CPU
-    backend, a call instruction is emitted which targets a symbol with the name
-    `call_target_name`.
+    Invokes exteral code with `inputs` as arguments and produces `results`.
 
-    `call_target_name` and `backend_config` can be arbitrary strings, but
-    `call_target_name` should be short as it may be used in labels.
-    `backend_config` can encode arbitrarily large amounts of information.
-
-    `has_side_effect` must be true if the custom call has side-effects.
-    `api_version` specifies the version of the API used by the custom call
-    function.
-
-    A custom call may apply functions within the scope of the parent module.
-    They can be referenced using `called_computations` attribute.
-
-    A custom call can also have layout constraints on operands and results which
-    can be specified as optional `operand_layouts` and `result_layouts`
-    attributes. The layout attribute is an array of rank-1 index tensors and the
-    i-th layout attribute specifies the layout for i-th operand/result.
-
-    The `operand_layouts` & `result_layouts` attributes can be specified under
-    the following constraints:
-    1) Either both `operand_layouts` and `result_layouts` are specified or none.
-    2) None of the operands are of tuple type.
-    3) None of the results are of tuple type except the common case of single
-       tuple result packing non-tuple values is allowed. In this case the i-th
-       `result_layouts` attribute specifies the layout of i-th element in the
-       result tuple.
-
-    See https://www.tensorflow.org/xla/operation_semantics#customcall.
+    See:
+    https://github.com/openxla/stablehlo/blob/main/docs/spec_draft.md#stablehlocustom_call
 
     Example:
-
     ```mlir
-    %1 = stablehlo.custom_call @foo(%arg0, %arg1) {backend_config = "bar", has_side_effect = true}
-          : (tensor<2x3xf32>, tensor<5x5xf32>) -> tensor<1x2x3xf32>
+    %results = "stablehlo.custom_call"(%inputs0) {
+      call_target_name = "foo",
+      has_side_effect = false,
+      backend_config = "bar",
+      api_version = 1 : i32,
+      called_computations = [@foo],
+      operand_layouts = [dense<> : tensor<0xindex>],
+      result_layouts = [dense<> : tensor<0xindex>],
+      output_operand_aliases = [
+        #stablehlo.output_operand_alias<output_tuple_indices = [],
+                                        operand_index = 0,
+                                        operand_tuple_indices = []>]
+    } : (tensor<f32>) -> tensor<f32>
     ```
   }];