issue/410 Feature: Add infinicore python package

.github/workflows/build.yml

@@ -31,7 +31,7 @@ jobs:
         xmake-version: latest
 
     - name: Build & Install
-      run: python scripts/install.py --omp=y
+      run: python scripts/install.py --omp=y -y
 
     - name: install python packages
       run: |

MANIFEST.in

+global-include *

include/infinicore.hpp

+#ifndef __INFINICORE_API_HPP__
+#define __INFINICORE_API_HPP__
+
+#include "infinicore/tensor.hpp"
+
+#endif

include/infinicore/device.hpp

+#ifndef __INFINICORE_DEVICE_API_HPP__
+#define __INFINICORE_DEVICE_API_HPP__
+
+#include <cstdint>
+#include <string>
+
+namespace infinicore {
+
+class Device {
+public:
+    using Index = std::size_t;
+
+    enum class Type {
+        cpu,
+        cuda,
+        meta,
+    };
+
+    Device(const Type &type, const Index &index = 0);
+
+    const Type &get_type() const;
+
+    const Index &get_index() const;
+
+    std::string to_string() const;
+
+    static std::string to_string(const Type &type);
+
+private:
+    Type type_;
+
+    Index index_;
+};
+
+} // namespace infinicore
+
+#endif

include/infinicore/dtype.hpp

+#ifndef __INFINICORE_DTYPE_API_HPP__
+#define __INFINICORE_DTYPE_API_HPP__
+
+#include <infinicore.h>
+
+namespace infinicore {
+
+enum class DataType {
+    bfloat16 = INFINI_DTYPE_BF16,
+    float16 = INFINI_DTYPE_F16,
+    float32 = INFINI_DTYPE_F32,
+    float64 = INFINI_DTYPE_F64,
+    int32 = INFINI_DTYPE_I32,
+    int64 = INFINI_DTYPE_I64,
+    uint8 = INFINI_DTYPE_U8,
+};
+
+std::string to_string(const DataType &dtype);
+
+} // namespace infinicore
+
+#endif

include/infinicore/tensor.hpp

+#ifndef __INFINICORE_TENSOR_API_HPP__
+#define __INFINICORE_TENSOR_API_HPP__
+
+#include <vector>
+
+#include "device.hpp"
+#include "dtype.hpp"
+
+namespace infinicore {
+
+class Tensor {
+public:
+    using Size = std::size_t;
+
+    using Stride = std::ptrdiff_t;
+
+    using Shape = std::vector<Size>;
+
+    using Strides = std::vector<Stride>;
+
+    Tensor(const Shape &shape, const DataType &dtype, const Device &device);
+
+    const Shape &get_shape() const;
+
+    const DataType &get_dtype() const;
+
+    const Device &get_device() const;
+
+private:
+    Shape shape_;
+
+    DataType dtype_;
+
+    Device device_;
+};
+
+} // namespace infinicore
+
+#endif

pyproject.toml

+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "InfiniCore"
+version = "0.1.0"
+description = "InfiniCore 是一个跨平台统一编程工具集，为不同芯片平台的功能（包括计算、运行时、通信等）提供统一 C 语言接口。"
+readme = "README.md"
+dependencies = []
+requires-python = ">=3.10"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+]
+
+[project.urls]
+Homepage = "https://github.com/InfiniTensor/InfiniCore"
+Issues = "https://github.com/InfiniTensor/InfiniCore/issues"
+
+[tool.ruff]
+src = [".", "src"]
+
+[tool.ruff.lint]
+select = ["E4", "E7", "E9", "F", "I"]

setup.py

+import glob
+import os
+import subprocess
+from pathlib import Path
+
+from setuptools import setup
+from setuptools.command.build_py import build_py
+
+INSTALLATION_DIR = os.getenv("INFINI_ROOT", str(Path.home() / ".infini"))
+
+LIB_DIR = os.path.join(INSTALLATION_DIR, "lib")
+
+PACKAGE_NAME = "infinicore"
+
+PACKAGE_DIR = os.path.join(INSTALLATION_DIR, PACKAGE_NAME)
+
+
+class BuildPy(build_py):
+    def run(self):
+        subprocess.run(["xmake", "build", "-y"])
+        subprocess.run(["xmake", "install"])
+        built_lib = glob.glob(os.path.join(LIB_DIR, f"{PACKAGE_NAME}.*"))[0]
+        os.makedirs(PACKAGE_DIR, exist_ok=True)
+        self.copy_file(built_lib, PACKAGE_DIR)
+
+
+setup(
+    cmdclass={"build_py": BuildPy},
+    package_dir={"": "."},
+)

src/infinicore/device.cc

+#include <infinicore.hpp>
+
+namespace infinicore {
+
+Device::Device(const Type &type, const Index &index) : type_{type}, index_{index} {}
+
+const Device::Type &Device::get_type() const {
+    return type_;
+}
+
+const Device::Index &Device::get_index() const {
+    return index_;
+}
+
+std::string Device::to_string() const {
+    return to_string(type_) + ":" + std::to_string(index_);
+}
+
+std::string Device::to_string(const Type &type) {
+    switch (type) {
+    case Type::cpu:
+        return "cpu";
+    case Type::cuda:
+        return "cuda";
+    case Type::meta:
+        return "meta";
+    }
+
+    // TODO: Add error handling.
+    return "";
+}
+
+} // namespace infinicore

src/infinicore/dtype.cc

+#include <infinicore.hpp>
+
+namespace infinicore {
+
+std::string to_string(const DataType &dtype) {
+    std::string str{"infinicore."};
+
+    switch (dtype) {
+    case DataType::bfloat16:
+        str += "bfloat16";
+        break;
+    case DataType::float16:
+        str += "float16";
+        break;
+    case DataType::float32:
+        str += "float32";
+        break;
+    case DataType::float64:
+        str += "float64";
+        break;
+    case DataType::int32:
+        str += "int32";
+        break;
+    case DataType::int64:
+        str += "int64";
+        break;
+    case DataType::uint8:
+        str += "uint8";
+        break;
+    }
+
+    return str;
+}
+
+} // namespace infinicore

src/infinicore/infinicore.cc

+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+#include <infinicore.hpp>
+
+namespace py = pybind11;
+
+namespace infinicore {
+
+PYBIND11_MODULE(infinicore, m) {
+    py::enum_<DataType>(m, "dtype")
+        .value("bfloat16", DataType::bfloat16)
+        .value("float16", DataType::float16)
+        .value("float32", DataType::float32)
+        .value("float64", DataType::float64)
+        .value("int32", DataType::int32)
+        .value("int64", DataType::int64)
+        .value("uint8", DataType::uint8)
+        .export_values();
+
+    py::class_<Device>(m, "Device")
+        .def(py::init<const Device::Type &, const Device::Index &>(),
+             py::arg("type"), py::arg("index") = 0)
+        .def_property_readonly("type", &Device::get_type)
+        .def_property_readonly("index", &Device::get_index)
+        .def("__repr__", static_cast<std::string (Device::*)() const>(&Device::to_string));
+
+    py::class_<Tensor>(m, "Tensor")
+        .def(py::init<const Tensor::Shape &, const DataType &, const Device &>(),
+             py::arg("shape"), py::arg("dtype") = DataType::float32, py::arg("device") = Device{Device::Type::cpu})
+        .def_property_readonly("shape", &Tensor::get_shape)
+        .def_property_readonly("dtype", &Tensor::get_dtype)
+        .def_property_readonly("device", &Tensor::get_device);
+}
+
+} // namespace infinicore

src/infinicore/tensor.cc

+#include <infinicore.hpp>
+
+namespace infinicore {
+
+Tensor::Tensor(const Shape &shape, const DataType &dtype, const Device &device) : shape_{shape}, dtype_{dtype}, device_{device} {}
+
+const Tensor::Shape &Tensor::get_shape() const {
+    return shape_;
+}
+
+const DataType &Tensor::get_dtype() const {
+    return dtype_;
+}
+
+const Device &Tensor::get_device() const {
+    return device_;
+}
+
+} // namespace infinicore

xmake.lua

 add_rules("mode.debug", "mode.release")
+add_requires("pybind11")
+
 -- Define color codes
 local GREEN = '\27[0;32m'
 local YELLOW = '\27[1;33m'
@@ -310,11 +312,23 @@ target("infiniccl")
     set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini"))
 target_end()
 
-target("all")
+target("infinicore_c_api")
     set_kind("phony")
     add_deps("infiniop", "infinirt", "infiniccl")
     after_build(function (target) print(YELLOW .. "[Congratulations!] Now you can install the libraries with \"xmake install\"" .. NC) end)
 target_end()
 
+target("infinicore")
+    add_rules("python.library", {soabi = true})
+    add_packages("pybind11")
+
+    set_kind("shared")
+    add_deps("infinicore_c_api")
+
+    add_files("src/infinicore/*.cc")
+
+    set_installdir(os.getenv("INFINI_ROOT") or (os.getenv(is_host("windows") and "HOMEPATH" or "HOME") .. "/.infini"))
+target_end()
+
 -- Tests
 includes("xmake/test.lua")