Initial commit of libshortfin. (#119)

This is initial work to create a high level C++ runtime library for
serving (with Python bindings). It incorporates learnings from multiple
priors (IREE Python Runtime, various samples, Turbine runtime, PJRT
runtime) and has several design goals:

* High level C++ API with a Python API that is almost a direct mirror.
* Opinionated concurrency and threading model (i.e. can be used
free-threaded but designed to work with controlled isolate like thread
loops).
* Explicit system setup which can support a variety of bespoke
configurations.
* Safe/low-code APIs for common operations with escape hatches for
esoteric/advanced usage.
* Robust multi-device (both heterogenous and homegenous) from the get-go
from a single process.
* Designed to use free-threaded (gil-lis) Python for high concurrency
situations.
* Minimal deps (i.e. can interop with things like torch, numpy, etc but
has enough bolts to build a complete solution without them).
* Compatibility and preference for async/await style programming (with
direct support for C++ co-routines and Python asyncio).
* Opinionated representation of the local topology, especially focused
on managing execution on directly bus-connected devices (while being
scalable to remote use as an enhancement).

---------

Signed-off-by: Stella Laurenzo <[email protected]>

Author: stellaraccident

.pre-commit-config.yaml

@@ -12,3 +12,7 @@ repos:
     rev: 22.10.0
     hooks:
     -   id: black
+-   repo: https://github.com/pre-commit/mirrors-clang-format
+    rev: 'v18.1.4'
+    hooks:
+    -   id: clang-format

libshortfin/.clang-format

@@ -0,0 +1,8 @@
+# Copyright 2019 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# Follow IREE runtime style, based on Google.
+BasedOnStyle: Google

libshortfin/CMakeLists.txt

@@ -0,0 +1,47 @@
+# Copyright 2024 Advanced Micro Devices, Inc
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions. See
+# https://llvm.org/LICENSE.txt for license information. SPDX-License-Identifier:
+# Apache-2.0 WITH LLVM-exception
+
+cmake_minimum_required(VERSION 3.28)
+
+if(CMAKE_SOURCE_DIR STREQUAL CMAKE_BINARY_DIR)
+  message(
+    FATAL_ERROR
+      "Do not build in-source. Please remove CMakeCache.txt and the CMakeFiles/ directory. Then build out-of-source."
+  )
+endif()
+
+project(
+  "libshortfin"
+  VERSION 0.9
+  LANGUAGES C CXX)
+
+set(CMAKE_C_STANDARD 11)
+set(CMAKE_CXX_STANDARD 20)
+# https://discourse.cmake.org/t/cmake-3-28-cmake-cxx-compiler-clang-scan-deps-notfound-not-found/9244/3
+set(CMAKE_CXX_SCAN_FOR_MODULES 0)
+
+option(SHORTFIN_BUILD_PYTHON_BINDINGS "Builds Python Bindings" OFF)
+
+# Includes.
+list(APPEND CMAKE_MODULE_PATH
+  ${CMAKE_CURRENT_LIST_DIR}/build_tools/cmake/
+)
+include(shortfin_library)
+
+# Dependencies.
+find_package(spdlog)
+find_package(xtensor)
+find_package(IREERuntime)
+
+add_subdirectory(src)
+
+if(SHORTFIN_BUILD_PYTHON_BINDINGS)
+  find_package(Python 3.8 COMPONENTS Interpreter Development.Module REQUIRED)
+  add_subdirectory(bindings/python)
+  set(SHORTFIN_PYTHON_CPP_PREBUILT "TRUE") # See setup.py.
+  configure_file(setup.py setup.py @ONLY)
+  configure_file(pyproject.toml pyproject.toml COPYONLY)
+endif()

libshortfin/README.md

@@ -0,0 +1,73 @@
+# libshortfin - SHARK C++ inference library
+
+## Dev Builds
+
+Library dependencies:
+
+* [spdlog](https://github.com/gabime/spdlog)
+* [xtensor](https://github.com/xtensor-stack/xtensor)
+* [iree runtime](https://github.com/iree-org/iree)
+
+On recent Ubuntu, the primary dependencies can be satisfied via:
+
+```
+apt install libspdlog-dev libxtensor-dev
+```
+
+CMake must be told how to find the IREE runtime, either from a distribution
+tarball, or local build/install dir. For a local build directory, pass:
+
+```
+# Assumes that the iree-build directory is adjacent to this repo.
+-DCMAKE_PREFIX_PATH=$(pwd)/../../iree-build/lib/cmake/IREE
+```
+
+One liner recommended CMake command (note that CMAKE_LINKER_TYPE requires
+cmake>=3.29):
+
+```
+cmake -GNinja -S. -Bbuild \
+    -DCMAKE_C_COMPILER=clang -DCMAKE_CXX_COMPILER=clang++ \
+    -DCMAKE_LINKER_TYPE=LLD \
+    -DCMAKE_PREFIX_PATH=$(pwd)/../../iree-build/lib/cmake/IREE
+```
+
+## Building Python Bindings
+
+If using a Python based development flow, there are two options:
+
+1. `pip install -v .` to build and install the library (TODO: Not yet implemented).
+2. Build with cmake and `-DSHORTFIN_BUILD_PYTHON_BINDINGS=ON` and then
+   from the `build/` directory, run `pip install -v -e .` to create an
+   editable install that will update as you build the C++ project.
+
+If predominantly developing with a C++ based flow, the second option is
+recommended. Your python install should track any source file changes or
+builds without further interaction. Re-installing will be necessary if package
+structure changes significantly.
+
+## Running Tests
+
+The project uses a combination of ctest for native C++ tests and pytest. Much
+of the functionality is only tested via the Python tests, using the
+`_shortfin.lib` internal implementation directly. In order to run these tests,
+you must have installed the Python package as per the above steps.
+
+Which style of test is used is pragmatic and geared at achieving good test
+coverage with a minimum of duplication. Since it is often much more expensive
+to build native tests of complicated flows, many things are only tested via
+Python. This does not preclude having other language bindings later, but it
+does mean that the C++ core of the library must always be built with the
+Python bindings to test the most behavior. Given the target of the project,
+this is not considered to be a significant issue.
+
+# Production Library Building
+
+In order to build a production library, additional build steps are typically
+recommended:
+
+* Compile all deps with the same compiler/linker for LTO compatibility
+* Provide library dependencies manually and compile them with LTO
+* Compile dependencies with `-fvisibility=hidden`
+* Enable LTO builds of libshortfin
+* Set flags to enable symbol versioning

libshortfin/bindings/python/CMakeLists.txt

@@ -0,0 +1,39 @@
+# Copyright 2024 Advanced Micro Devices, Inc
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions. See
+# https://llvm.org/LICENSE.txt for license information. SPDX-License-Identifier:
+# Apache-2.0 WITH LLVM-exception
+
+# libshortfin publishes multiple python packages: - _shortfin: Trampoline
+# __init__.py which looks at environment variables to load an appropriate native
+# library. - _shortfin_default.lib: Native library as a default, uninstrumented
+# build. - _shortfin_tracing.lib: Native library with tracing enabled (TODO). -
+# Others.
+
+# Detect the installed nanobind package and import it into CMake
+execute_process(
+  COMMAND "${Python_EXECUTABLE}" -m nanobind --cmake_dir
+  OUTPUT_STRIP_TRAILING_WHITESPACE
+  OUTPUT_VARIABLE nanobind_ROOT)
+find_package(nanobind CONFIG REQUIRED)
+
+nanobind_add_module(shortfin_python_extension NB_STATIC LTO
+  array_binding.cc
+  lib_ext.cc
+)
+
+set_target_properties(shortfin_python_extension
+                      PROPERTIES OUTPUT_NAME "_shortfin_default/lib")
+
+target_link_libraries(shortfin_python_extension
+  # TODO: This should be configurable as to whether we link to the static
+  # or dynamic version.
+  PRIVATE shortfin
+)
+
+nanobind_add_stub(
+  shortfin_python_extension_stub
+  MODULE _shortfin_default.lib
+  OUTPUT _shortfin_default/lib.pyi
+  DEPENDS shortfin_python_extension
+)

libshortfin/bindings/python/_shortfin/__init__.py

@@ -0,0 +1,12 @@
+# Copyright 2024 Advanced Micro Devices, Inc
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+# The proper way to import this package is via:
+#   from _shortfin import lib as sfl
+
+# TODO: Use environment variables to determine which built variant to
+# import.
+from _shortfin_default import lib

libshortfin/bindings/python/array_binding.cc

@@ -0,0 +1,133 @@
+// Copyright 2024 Advanced Micro Devices, Inc
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "./lib_ext.h"
+#include "./utils.h"
+#include "shortfin/array/api.h"
+
+using namespace shortfin::array;
+
+namespace shortfin::python {
+
+void BindArray(py::module_ &m) {
+  py::class_<DType>(m, "DType")
+      .def_prop_ro("is_boolean", &DType::is_boolean)
+      .def_prop_ro("is_integer", &DType::is_integer)
+      .def_prop_ro("is_float", &DType::is_float)
+      .def_prop_ro("is_complex", &DType::is_complex)
+      .def_prop_ro("bit_count", &DType::bit_count)
+      .def_prop_ro("is_byte_aligned", &DType::is_byte_aligned)
+      .def_prop_ro("dense_byte_count", &DType::dense_byte_count)
+      .def("is_integer_bitwidth", &DType::is_integer_bitwidth)
+      .def(py::self == py::self)
+      .def("__repr__", &DType::name);
+
+  m.attr("opaque8") = DType::opaque8();
+  m.attr("opaque16") = DType::opaque16();
+  m.attr("opaque32") = DType::opaque32();
+  m.attr("opaque64") = DType::opaque64();
+  m.attr("bool8") = DType::bool8();
+  m.attr("int4") = DType::int4();
+  m.attr("sint4") = DType::sint4();
+  m.attr("uint4") = DType::uint4();
+  m.attr("int8") = DType::int8();
+  m.attr("sint8") = DType::sint8();
+  m.attr("uint8") = DType::uint8();
+  m.attr("int16") = DType::int16();
+  m.attr("sint16") = DType::sint16();
+  m.attr("uint16") = DType::uint16();
+  m.attr("int32") = DType::int32();
+  m.attr("sint32") = DType::sint32();
+  m.attr("uint32") = DType::uint32();
+  m.attr("int64") = DType::int64();
+  m.attr("sint64") = DType::sint64();
+  m.attr("uint64") = DType::uint64();
+  m.attr("float16") = DType::float16();
+  m.attr("float32") = DType::float32();
+  m.attr("float64") = DType::float64();
+  m.attr("bfloat16") = DType::bfloat16();
+  m.attr("complex64") = DType::complex64();
+  m.attr("complex128") = DType::complex128();
+
+  py::class_<storage>(m, "storage")
+      .def_static(
+          "allocate_host",
+          [](local::ScopedDevice &device, iree_device_size_t allocation_size) {
+            return storage::AllocateHost(device, allocation_size);
+          },
+          py::arg("device"), py::arg("allocation_size"), py::keep_alive<0, 1>())
+      .def_static(
+          "allocate_device",
+          [](local::ScopedDevice &device, iree_device_size_t allocation_size) {
+            return storage::AllocateDevice(device, allocation_size);
+          },
+          py::arg("device"), py::arg("allocation_size"), py::keep_alive<0, 1>())
+      .def("fill",
+           [](storage &self, py::handle buffer) {
+             Py_buffer py_view;
+             int flags = PyBUF_FORMAT | PyBUF_ND;  // C-Contiguous ND.
+             if (PyObject_GetBuffer(buffer.ptr(), &py_view, flags) != 0) {
+               throw py::python_error();
+             }
+             PyBufferReleaser py_view_releaser(py_view);
+             self.Fill(py_view.buf, py_view.len);
+           })
+      .def("__repr__", &storage::to_s);
+
+  py::class_<base_array>(m, "base_array")
+      .def_prop_ro("dtype", &base_array::dtype)
+      .def_prop_ro("shape", &base_array::shape);
+  py::class_<device_array, base_array>(m, "device_array")
+      .def("__init__", [](py::args, py::kwargs) {})
+      .def_static("__new__",
+                  [](py::handle py_type, class storage storage,
+                     std::span<const size_t> shape, DType dtype) {
+                    return custom_new_keep_alive<device_array>(
+                        py_type, /*keep_alive=*/storage.scope(), storage, shape,
+                        dtype);
+                  })
+      .def_static("__new__",
+                  [](py::handle py_type, local::ScopedDevice &device,
+                     std::span<const size_t> shape, DType dtype) {
+                    return custom_new_keep_alive<device_array>(
+                        py_type, /*keep_alive=*/device.scope(),
+                        device_array::allocate(device, shape, dtype));
+                  })
+      .def_prop_ro("device", &device_array::device,
+                   py::rv_policy::reference_internal)
+      .def_prop_ro("storage", &device_array::storage,
+                   py::rv_policy::reference_internal)
+      .def("__repr__", &device_array::to_s);
+  py::class_<host_array, base_array>(m, "host_array")
+      .def("__init__", [](py::args, py::kwargs) {})
+      .def_static("__new__",
+                  [](py::handle py_type, class storage storage,
+                     std::span<const size_t> shape, DType dtype) {
+                    return custom_new_keep_alive<host_array>(
+                        py_type, /*keep_alive=*/storage.scope(), storage, shape,
+                        dtype);
+                  })
+      .def_static("__new__",
+                  [](py::handle py_type, local::ScopedDevice &device,
+                     std::span<const size_t> shape, DType dtype) {
+                    return custom_new_keep_alive<host_array>(
+                        py_type, /*keep_alive=*/device.scope(),
+                        host_array::allocate(device, shape, dtype));
+                  })
+      .def_static("__new__",
+                  [](py::handle py_type, device_array &device_array) {
+                    return custom_new_keep_alive<host_array>(
+                        py_type, /*keep_alive=*/device_array.device().scope(),
+                        host_array::for_transfer(device_array));
+                  })
+      .def_prop_ro("device", &host_array::device,
+                   py::rv_policy::reference_internal)
+      .def_prop_ro("storage", &host_array::storage,
+                   py::rv_policy::reference_internal)
+      .def("__repr__", &host_array::to_s);
+}
+
+}  // namespace shortfin::python