移除了 infinicore::py 层

Author: voltjia

python/infinicore/__init__.py

@@ -0,0 +1,58 @@
+from infinicore.device import device
+from infinicore.dtype import (
+    bfloat16,
+    bool,
+    cdouble,
+    cfloat,
+    chalf,
+    complex32,
+    complex64,
+    complex128,
+    double,
+    float,
+    float16,
+    float32,
+    float64,
+    half,
+    int,
+    int8,
+    int16,
+    int32,
+    int64,
+    long,
+    short,
+    uint8,
+)
+from infinicore.ops.matmul import matmul
+from infinicore.tensor import empty, from_blob, ones, zeros
+
+__all__ = [
+    "device",
+    "bfloat16",
+    "bool",
+    "cdouble",
+    "cfloat",
+    "chalf",
+    "complex32",
+    "complex64",
+    "complex128",
+    "double",
+    "float",
+    "float16",
+    "float32",
+    "float64",
+    "half",
+    "int",
+    "int8",
+    "int16",
+    "int32",
+    "int64",
+    "long",
+    "short",
+    "uint8",
+    "matmul",
+    "empty",
+    "from_blob",
+    "ones",
+    "zeros",
+]

python/infinicore/device.py

@@ -0,0 +1,80 @@
+import _infinicore
+
+
+class device:
+    def __init__(self, type=None, index=None):
+        if type is None:
+            type = "cpu"
+
+        if isinstance(type, device):
+            self.type = type.type
+            self.index = type.index
+
+            return
+
+        if ":" in type:
+            if index is not None:
+                raise ValueError(
+                    '`index` should not be provided when `type` contains `":"`.'
+                )
+
+            type, index = type.split(":")
+            index = int(index)
+
+        self.type = type
+
+        self.index = index
+
+        _type, _index = device._to_infinicore_device(type, index if index else 0)
+
+        self._underlying = _infinicore.Device(_type, _index)
+
+    def __repr__(self):
+        return f"device(type='{self.type}'{f', index={self.index}' if self.index is not None else ''})"
+
+    def __str__(self):
+        return f"{self.type}{f':{self.index}' if self.index is not None else ''}"
+
+    @staticmethod
+    def _to_infinicore_device(type, index):
+        all_device_types = tuple(_infinicore.Device.Type.__members__.values())[:-1]
+        all_device_count = tuple(
+            _infinicore.get_device_count(device) for device in all_device_types
+        )
+
+        torch_devices = {
+            torch_type: {
+                infinicore_type: 0
+                for infinicore_type in all_device_types
+                if _TORCH_DEVICE_MAP[infinicore_type] == torch_type
+            }
+            for torch_type in _TORCH_DEVICE_MAP.values()
+        }
+
+        for i, count in enumerate(all_device_count):
+            infinicore_device_type = _infinicore.Device.Type(i)
+            torch_devices[_TORCH_DEVICE_MAP[infinicore_device_type]][
+                infinicore_device_type
+            ] += count
+
+        for infinicore_device_type, infinicore_device_count in torch_devices[
+            type
+        ].items():
+            for i in range(infinicore_device_count):
+                if index == 0:
+                    return infinicore_device_type, i
+
+                index -= 1
+
+
+_TORCH_DEVICE_MAP = {
+    _infinicore.Device.Type.CPU: "cpu",
+    _infinicore.Device.Type.NVIDIA: "cuda",
+    _infinicore.Device.Type.CAMBRICON: "mlu",
+    _infinicore.Device.Type.ASCEND: "npu",
+    _infinicore.Device.Type.METAX: "cuda",
+    _infinicore.Device.Type.MOORE: "musa",
+    _infinicore.Device.Type.ILUVATAR: "cuda",
+    _infinicore.Device.Type.KUNLUN: "cuda",
+    _infinicore.Device.Type.SUGON: "cuda",
+}

python/infinicore/dtype.py

@@ -0,0 +1,57 @@
+import _infinicore
+
+
+class dtype:
+    def __init__(self, data_type):
+        """An internal method. Please do not use this directly."""
+
+        self._underlying = data_type
+
+    def __repr__(self):
+        repr_map = {
+            _infinicore.DataType.BYTE: "uint8",
+            _infinicore.DataType.BOOL: "bool",
+            _infinicore.DataType.I8: "int8",
+            _infinicore.DataType.I16: "int16",
+            _infinicore.DataType.I32: "int32",
+            _infinicore.DataType.I64: "int64",
+            _infinicore.DataType.U8: "uint8",
+            _infinicore.DataType.U16: "uint16",
+            _infinicore.DataType.U32: "uint32",
+            _infinicore.DataType.U64: "uint64",
+            _infinicore.DataType.F8: "float8",
+            _infinicore.DataType.F16: "float16",
+            _infinicore.DataType.F32: "float32",
+            _infinicore.DataType.F64: "float64",
+            _infinicore.DataType.C16: "complex16",
+            _infinicore.DataType.C32: "complex32",
+            _infinicore.DataType.C64: "complex64",
+            _infinicore.DataType.C128: "complex128",
+            _infinicore.DataType.BF16: "bfloat16",
+        }
+
+        return f"infinicore.{repr_map[self._underlying]}"
+
+
+float32 = dtype(_infinicore.DataType.F32)
+float = float32
+float64 = dtype(_infinicore.DataType.F64)
+double = float64
+complex32 = dtype(_infinicore.DataType.C32)
+chalf = complex32
+complex64 = dtype(_infinicore.DataType.C64)
+cfloat = complex64
+complex128 = dtype(_infinicore.DataType.C128)
+cdouble = complex128
+float16 = dtype(_infinicore.DataType.F16)
+half = float16
+bfloat16 = dtype(_infinicore.DataType.BF16)
+uint8 = dtype(_infinicore.DataType.U8)
+int8 = dtype(_infinicore.DataType.I8)
+int16 = dtype(_infinicore.DataType.I16)
+short = int16
+int32 = dtype(_infinicore.DataType.I32)
+int = int32
+int64 = dtype(_infinicore.DataType.I64)
+long = int64
+bool = dtype(_infinicore.DataType.BOOL)

python/infinicore/ops/__init__.py

@@ -0,0 +1,10 @@
+import _infinicore
+
+from infinicore.tensor import Tensor
+
+
+def matmul(input, other, *, out=None):
+    if out is None:
+        return Tensor(_infinicore.matmul(input._underlying, other._underlying))
+
+    _infinicore.matmul_(out._underlying, input._underlying, other._underlying)

python/infinicore/ops/matmul.py

@@ -0,0 +1,94 @@
+import ctypes
+
+import _infinicore
+
+
+class Tensor:
+    def __init__(self, tensor):
+        """An internal method. Please do not use this directly."""
+
+        self._underlying = tensor
+
+    @property
+    def shape(self):
+        return self._underlying.shape
+
+    @property
+    def dtype(self):
+        return self._underlying.dtype
+
+    @property
+    def device(self):
+        return self._underlying.device
+
+    @property
+    def ndim(self):
+        return self._underlying.device
+
+    def data_ptr(self):
+        return self._underlying.data_ptr
+
+    def size(self, dim=None):
+        if dim is None:
+            return self.shape
+
+        return self.shape[dim]
+
+    def stride(self, dim=None):
+        if dim is None:
+            return self._underlying.strides
+
+        return self._underlying.strides[dim]
+
+    def numel(self):
+        return self._underlying.numel()
+
+    def is_contiguous(self):
+        return self._underlying.is_contiguous()
+
+    def is_is_pinned(self):
+        return self._underlying.is_is_pinned()
+
+    def copy_(self, src):
+        return Tensor(self._underlying.copy_(src._underlying))
+
+    def to(self, *args, **kwargs):
+        return Tensor(
+            self._underlying.to(*tuple(arg._underlying for arg in args), **kwargs)
+        )
+
+    def as_strided(self, size, stride):
+        Tensor(self._underlying.as_strided(size, stride))
+
+    def contiguous(self):
+        return Tensor(self._underlying.contiguous())
+
+    def permute(self, dims):
+        return Tensor(self._underlying.permute(dims))
+
+    def view(self, shape):
+        return Tensor(self._underlying.view(shape))
+
+
+def empty(size, *, dtype=None, device=None, pin_memory=False):
+    return Tensor(
+        _infinicore.empty(size, dtype._underlying, device._underlying, pin_memory)
+    )
+
+
+def zeros(size, *, dtype=None, device=None, pin_memory=False):
+    return Tensor(
+        _infinicore.zeros(size, dtype._underlying, device._underlying, pin_memory)
+    )
+
+
+def ones(size, *, dtype=None, device=None, pin_memory=False):
+    return Tensor(
+        _infinicore.ones(size, dtype._underlying, device._underlying, pin_memory)
+    )
+
+
+def from_blob(data_ptr, size, *, dtype=None, device=None):
+    return Tensor(
+        _infinicore.from_blob(data_ptr, size, dtype._underlying, device._underlying)
+    )

python/infinicore/tensor.py

@@ -1,5 +1,5 @@
-import glob
 import os
+import shutil
 import subprocess
 from pathlib import Path
 
@@ -10,9 +10,9 @@
 
 LIB_DIR = os.path.join(INSTALLATION_DIR, "lib")
 
-PACKAGE_NAME = "infinicore"
+PYTHON_PACKAGE_DIR = os.path.join("python", "infinicore")
 
-PACKAGE_DIR = os.path.join(INSTALLATION_DIR, PACKAGE_NAME)
+LIB_PACKAGE_DIR = os.path.join(LIB_DIR, "infinicore")
 
 
 class BuildPy(build_py):
@@ -21,12 +21,11 @@ def run(self):
         subprocess.run(["xmake", "install"])
         subprocess.run(["xmake", "build", "-y", "infinicore"])
         subprocess.run(["xmake", "install", "infinicore"])
-        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)
 
+        if os.path.exists(LIB_PACKAGE_DIR):
+            shutil.rmtree(LIB_PACKAGE_DIR)
 
-setup(
-    cmdclass={"build_py": BuildPy},
-    package_dir={"": "."},
-)
+        shutil.copytree(PYTHON_PACKAGE_DIR, LIB_PACKAGE_DIR)
+
+
+setup(cmdclass={"build_py": BuildPy})

setup.py

@@ -0,0 +1,16 @@
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+#include "infinicore.hpp"
+
+namespace py = pybind11;
+
+namespace infinicore::context {
+
+inline void bind(py::module &m) {
+    m.def("get_device", &getDevice);
+    m.def("get_device_count", &getDeviceCount);
+}
+
+} // namespace infinicore::context

src/infinicore/pybind11/context.hpp

@@ -0,0 +1,34 @@
+#pragma once
+
+#include <pybind11/pybind11.h>
+
+#include "infinicore.hpp"
+
+namespace py = pybind11;
+
+namespace infinicore::device {
+
+inline void bind(py::module &m) {
+    py::class_<Device> device(m, "Device");
+
+    py::enum_<Device::Type>(device, "Type")
+        .value("CPU", Device::Type::CPU)
+        .value("NVIDIA", Device::Type::NVIDIA)
+        .value("CAMBRICON", Device::Type::CAMBRICON)
+        .value("ASCEND", Device::Type::ASCEND)
+        .value("METAX", Device::Type::METAX)
+        .value("MOORE", Device::Type::MOORE)
+        .value("ILUVATAR", Device::Type::ILUVATAR)
+        .value("KUNLUN", Device::Type::KUNLUN)
+        .value("SUGON", Device::Type::SUGON)
+        .value("COUNT", Device::Type::COUNT);
+
+    device
+        .def(py::init<const Device::Type &, const Device::Index &>(),
+             py::arg("type") = Device::Type::CPU, py::arg("index") = 0)
+        .def_property_readonly("type", &Device::getType)
+        .def_property_readonly("index", &Device::getIndex)
+        .def("__str__", static_cast<std::string (Device::*)() const>(&Device::toString));
+}
+
+} // namespace infinicore::device