Fix build bugs

requirements.txt

@@ -2,4 +2,5 @@ cmake
 ninja
 pybind11
 pytest
+numpy
 torch

setup.py

@@ -9,6 +9,7 @@
 package_name = 'kfunca'
 version = '0.1.1'
 base_dir = os.path.dirname(os.path.abspath(__file__))
+default_parallel_build = 4  # Default number of parallel jobs for building
 
 
 class BuildExt(build_ext):
@@ -19,7 +20,7 @@ def run(self):
         super().run()
 
     def build_cmake(self, ext):
-        cmake = CMake()
+        cmake = CMake(default_parallel_build)
         extdir = pathlib.Path(self.get_ext_fullpath(ext.name)).parent.absolute()
         cmake.generate(rerun=True, output_dir=extdir)
         cmake.build()

src/core/include/half.h

@@ -215,11 +215,11 @@ struct alignas(2) Half {
 #if defined(__CUDACC__)
     DEVICE inline operator float() const;
     DEVICE inline Half(float f);
-    DEVICE inline operator BFloat16() const;
+    DEVICE inline Half(BFloat16 bf);
 #else
     inline operator float() const;
     inline Half(float f);
-    inline operator BFloat16() const;
+    inline Half(BFloat16 bf);
 #endif
 };
 
@@ -229,11 +229,11 @@ struct alignas(2) BFloat16 {
 #if defined(__CUDACC__)
     DEVICE inline operator float() const;
     DEVICE inline BFloat16(float f);
-    DEVICE inline operator Half() const;
+    DEVICE inline BFloat16(Half h);
 #else
     inline operator float() const;
     inline BFloat16(float f);
-    inline operator Half() const;
+    inline BFloat16(Half h);
 #endif
 };
 
@@ -247,8 +247,8 @@ DEVICE inline Half::Half(float f) {
     x = float_to_half(f);
 }
 
-DEVICE inline Half::operator BFloat16() const {
-    return BFloat16(half_to_float(x));
+DEVICE inline Half::Half(BFloat16 bf) {
+    x = float_to_half(bfloat16_to_float(bf.x));
 }
 
 DEVICE inline BFloat16::operator float() const {
@@ -259,8 +259,8 @@ DEVICE inline BFloat16::BFloat16(float f) {
     x = float_to_bfloat16(f);
 }
 
-DEVICE inline BFloat16::operator Half() const {
-    return Half(bfloat16_to_float(x));
+DEVICE inline BFloat16::BFloat16(Half h) {
+    x = float_to_bfloat16(half_to_float(h.x));
 }
 
 #else
@@ -273,8 +273,8 @@ inline Half::Half(float f) {
     x = fp16_ieee_from_fp32_value(f);
 }
 
-inline Half::operator BFloat16() const {
-    return BFloat16(fp16_ieee_to_fp32_value(x));
+inline Half::Half(BFloat16 bf) {
+    x = fp16_ieee_from_fp32_value(f32_from_bits(bf.x));
 }
 
 inline BFloat16::operator float() const {
@@ -285,8 +285,8 @@ inline BFloat16::BFloat16(float f) {
     x = round_to_nearest_even(f);
 }
 
-inline BFloat16::operator Half() const {
-    return Half(f32_from_bits(x));
+inline BFloat16::BFloat16(Half h) {
+    x = round_to_nearest_even(fp16_ieee_to_fp32_value(h.x));
 }
 
 #endif

src/core/include/tensor.h

@@ -30,17 +30,6 @@ Tensor empty_like_reduced(const Tensor &self, int dim, ScalarType dtype);
 Tensor zeros(std::vector<int64_t> shape, ScalarType dtype, int device = 0);
 std::ostream &operator<<(std::ostream &os, const Tensor &t);
 
-template <typename T>
-std::ostream &operator<<(std::ostream &os, const std::vector<T> &vec) {
-    os << "[";
-    for (auto i = 0; i < vec.size(); ++i) {
-        os << vec[i];
-        if (i + 1 != vec.size()) os << ", ";
-    }
-    os << "]";
-    return os;
-}
-
 template <typename T, int vec_size>
 struct d_array {
     T val[vec_size] = {0};

src/core/utils/exception.h

@@ -4,6 +4,7 @@
 #include <string>
 #include <sstream>
 #include <vector>
+#include <cstdint>
 
 #include "xstring.h"
 

src/core/utils/xstring.h

@@ -4,6 +4,17 @@
 #include <string>
 #include <sstream>
 
+template <typename T>
+std::ostream &operator<<(std::ostream &os, const std::vector<T> &vec) {
+    os << "[";
+    for (auto i = 0; i < vec.size(); ++i) {
+        os << vec[i];
+        if (i + 1 != vec.size()) os << ", ";
+    }
+    os << "]";
+    return os;
+}
+
 namespace utils {
 
 inline std::ostream &_str(std::ostream &ss) {

tools/cmake.py

@@ -33,9 +33,10 @@ def which(thefile: str) -> str | None:
 class CMake:
     "Manages cmake."
 
-    def __init__(self, build_dir: str = BUILD_DIR) -> None:
+    def __init__(self, parallel_build=None, build_dir=BUILD_DIR) -> None:
         self._cmake_command = CMake._get_cmake_command()
         self.build_dir = build_dir
+        self.parallel_build = parallel_build
 
     @property
     def _cmake_cache_file(self) -> str:
@@ -100,4 +101,10 @@ def generate(self, rerun: bool, output_dir: str):
         self.run(args, os.environ)
 
     def build(self):
-        self.run(['--build', '.'], os.environ)
+        if not self.parallel_build:
+            self.run(['--build', '.'], os.environ)
+        elif isinstance(self.parallel_build, int):
+            nworkers = str(self.parallel_build)
+            self.run(['--build', '.', '-j', nworkers], os.environ)
+        else:
+            self.run(['--build', '.', '--parallel'], os.environ)