diff --git a/docs/spec_draft.md b/docs/spec_draft.md index ace88665f58..b00ca303365 100644 --- a/docs/spec_draft.md +++ b/docs/spec_draft.md @@ -399,6 +399,7 @@ syntax. * [convolution](#stablehloconvolution) * [cosine](#stablehlocosine) * [count_leading_zeros](#stablehlocount_leading_zeros) + * [custom_call](#stablehlocustom_call) * [divide](#stablehlodivide) * [dot_general](#stablehlodot_general) * [dynamic_slice](#stablehlodynamic_slice) @@ -2066,6 +2067,47 @@ tensor and produces a `result` tensor. [Back to Ops](#index-of-ops) +## stablehlo.custom_call + +### Semantics + +Encapsulates an implementation-defined operation `call_target_name` that takes +`inputs` and `called_computations` and produces `results`. `has_side_effect`, +`backend_config` and `api_version` may be used to provide additional +implementation-defined metadata. + +### Inputs + +| Name | Type | +|-----------------------|-----------------------------------------------------------------| +| `inputs` | variadic number of values of any supported type | +| `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` | + +### Outputs + +| Name | Type | +|-----------|-------------------------------------------------| +| `results` | variadic number of values of any supported type | + +### 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] +} : (tensor) -> tensor +``` + +[Back to Ops](#index-of-ops) + ## stablehlo.divide ### Semantics diff --git a/docs/status.md b/docs/status.md index 4c0d89e48f0..d503de14ece 100644 --- a/docs/status.md +++ b/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 | yes* | 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 | yes | revisit | revisit | no | no | diff --git a/stablehlo/dialect/StablehloOps.td b/stablehlo/dialect/StablehloOps.td index fa6aab69f88..8ddd52a1308 100644 --- a/stablehlo/dialect/StablehloOps.td +++ b/stablehlo/dialect/StablehloOps.td @@ -2137,44 +2137,21 @@ def StableHLO_CustomCallOp: StableHLO_Op<"custom_call", [DeclareOpInterfaceMethods]> { 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`. + Encapsulates an implementation-defined operation `call_target_name` that + takes `inputs` and `called_computations` 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] + } : (tensor) -> tensor ``` }];