1 /* Copyright 2022 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef XLA_RUNTIME_RUNTIME_H_ 17 #define XLA_RUNTIME_RUNTIME_H_ 18 19 #include <stdint.h> 20 21 namespace xla { 22 namespace runtime { 23 24 //===----------------------------------------------------------------------===// 25 // XLA Runtime <-> XLA Executable integration API. 26 //===----------------------------------------------------------------------===// 27 28 // This API enables compiled XLA executables to call back into the XLA runtime 29 // for: 30 // 31 // - Returning results back to the caller by writing to result storage 32 // allocated in the call frame. 33 // - User friendly error reporting integrated with the high level execution 34 // model (errors do not crash the compiled binary). 35 // - Invoking custom calls registered with the XLA runtime. 36 // 37 // XLA compilation pipeline sets up passes to convert the regular functions 38 // to the so called "XLA entrypoint functions" integrated with the runtime using 39 // the API defined below, e.g. instead of conventional returns all results are 40 // returned via the `GetResultStorage` API. At MLIR level these operations 41 // correspond to the `rt` dialect, and converted to LLVM using the `rt-to-llvm` 42 // conversion pass. 43 // 44 // Runtime errors are reported back to the runtime via the `SetError` API. 45 // The compilation pipeline will automatically convert assertions in the 46 // entrypoint function into run-time errors. 47 48 // TODO(ezhulenev): Rename KernelContext to ExecutionContext. 49 50 // Opaque runtime execution context passed as the first argument to compiled 51 // executables and passed back to all runtime API methods. 52 using KernelContext = struct KernelContext; 53 54 // Returns a pointer to the memory location of the result at the given index. 55 void *GetResultStorage(KernelContext *, int64_t); 56 57 // Sets execution context to an error state. 58 void SetError(KernelContext *, const char *); 59 60 // Calls the custom call function registered with the runtime. Returns true 61 // if the custom call was successful. 62 bool CustomCall(KernelContext *, const char *target, void **args, void **attrs); 63 64 } // namespace runtime 65 } // namespace xla 66 67 #endif // XLA_RUNTIME_RUNTIME_H_ 68