madspace 0.3.1__cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

madspace/_madspace_py_loader.py

@@ -0,0 +1,111 @@
+import ctypes
+import logging
+import os
+import platform
+
+# pre-load libmadspace
+ctypes.CDLL(
+    os.path.join(
+        os.path.dirname(os.path.abspath(__file__)),
+        "lib",
+        "libmadspace.dylib" if platform.system() == "Darwin" else "libmadspace.so",
+    ),
+    mode=ctypes.RTLD_GLOBAL,
+)
+
+from ._madspace_py import *
+
+
+def _init():
+    """
+    Monkey-patch classes for a more pythonic experience.
+    """
+    set_lib_path(os.path.join(os.path.dirname(os.path.abspath(__file__)), "lib"))
+
+    def call_and_convert(runtime, args):
+        if len(args) == 0:
+            tensorlib = "numpy"
+        else:
+            tensorlib = type(args[0]).__module__
+        outputs = runtime.call(args)
+        # Convert outputs, lazy-loading torch or numpy
+        if tensorlib == "torch":
+            import torch
+
+            return tuple(torch.from_dlpack(out) for out in outputs)
+        else:
+            import numpy
+
+            return tuple(numpy.from_dlpack(out) for out in outputs)
+
+    def runtime_call(self, *args):
+        outputs = call_and_convert(self, args)
+        if len(outputs) == 1:
+            return outputs[0]
+        else:
+            return outputs
+
+    def function_call(self, *args):
+        if not hasattr(self, "runtime"):
+            self.runtime = FunctionRuntime(self)
+        outputs = call_and_convert(self.runtime, args)
+        if len(outputs) == 1:
+            return outputs[0]
+        else:
+            return outputs
+
+    def function_generator_call(self, *args):
+        if not hasattr(self, "runtime"):
+            self.runtime = FunctionRuntime(self.function())
+        outputs = call_and_convert(self.runtime, args)
+        if len(outputs) == 1:
+            return outputs[0]
+        else:
+            return outputs
+
+    def map_forward(self, inputs, conditions=[]):
+        if not hasattr(self, "forward_runtime"):
+            self.forward_runtime = FunctionRuntime(self.forward_function())
+        outputs = call_and_convert(self.forward_runtime, [*inputs, *conditions])
+        return outputs[:-1], outputs[-1]
+
+    def map_inverse(self, inputs, conditions=[]):
+        if not hasattr(self, "inverse_runtime"):
+            self.inverse_runtime = FunctionRuntime(self.inverse_function())
+        outputs = call_and_convert(self.inverse_runtime, [*inputs, *conditions])
+        return outputs[:-1], outputs[-1]
+
+    def tensor_numpy(tensor):
+        import numpy  # Lazy-load numpy, to make it optional dependency
+
+        return numpy.from_dlpack(tensor)
+
+    def tensor_torch(tensor):
+        import torch  # Lazy-load torch, to make it optional dependency
+
+        return torch.from_dlpack(tensor)
+
+    py_logger = logging.getLogger("madspace")
+
+    def log_handler(level, message):
+        match level:
+            case Logger.level_debug:
+                py_logger.debug(message)
+            case Logger.level_info:
+                py_logger.info(message)
+            case Logger.level_warning:
+                py_logger.warning(message)
+            case Logger.level_error:
+                py_logger.error(message)
+
+    FunctionRuntime.__call__ = runtime_call
+    Function.__call__ = function_call
+    FunctionGenerator.__call__ = function_generator_call
+    Mapping.map_forward = map_forward
+    Mapping.map_inverse = map_inverse
+    Tensor.numpy = tensor_numpy
+    Tensor.torch = tensor_torch
+    # Logger.set_log_handler(log_handler)
+
+
+_init()

madspace/include/madspace/constants.h

@@ -0,0 +1,17 @@
+#pragma once
+
+namespace madspace {
+
+// mathematical constants
+
+inline constexpr double PI = 3.14159265358979323846;
+inline constexpr double LOG_TWO = 0.69314718055994530942;
+inline constexpr double SQRT_HALF = 0.70710678118654752440;
+inline constexpr double TWO_DIV_SQRT_PI = 1.12837916709551257390;
+
+// spline flow settings
+
+inline constexpr double MIN_BIN_SIZE = 1e-3;
+inline constexpr double MIN_DERIVATIVE = 1e-3;
+
+} // namespace madspace

madspace/include/madspace/madcode/function.h

@@ -0,0 +1,102 @@
+#pragma once
+
+#include <iostream>
+#include <map>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include <nlohmann/json.hpp>
+
+#include "instruction.h"
+
+namespace madspace {
+
+struct InstructionCall {
+    InstructionPtr instruction;
+    ValueVec inputs;
+    ValueVec outputs;
+};
+
+class Function {
+public:
+    friend class FunctionBuilder;
+
+    Function() = default;
+
+    const ValueVec& inputs() const { return _inputs; }
+    const ValueVec& outputs() const { return _outputs; }
+    const ValueVec& locals() const { return _locals; }
+    const std::unordered_map<std::string, Value>& globals() const { return _globals; }
+    const std::vector<InstructionCall>& instructions() const { return _instructions; }
+
+    void save(const std::string& file) const;
+    static Function load(const std::string& file);
+
+private:
+    Function(
+        const ValueVec& inputs,
+        const ValueVec& outputs,
+        const ValueVec& locals,
+        const std::unordered_map<std::string, Value>& globals,
+        const std::vector<InstructionCall>& instructions
+    ) :
+        _inputs(inputs),
+        _outputs(outputs),
+        _locals(locals),
+        _globals(globals),
+        _instructions(instructions) {}
+
+    ValueVec _inputs;
+    ValueVec _outputs;
+    ValueVec _locals;
+    std::unordered_map<std::string, Value> _globals;
+    std::vector<InstructionCall> _instructions;
+};
+
+std::ostream& operator<<(std::ostream& out, const Value& value);
+std::ostream& operator<<(std::ostream& out, const ValueVec& list);
+std::ostream& operator<<(std::ostream& out, const InstructionCall& call);
+std::ostream& operator<<(std::ostream& out, const Function& func);
+
+void to_json(nlohmann::json& j, const InstructionCall& call);
+void to_json(nlohmann::json& j, const Function& call);
+void from_json(const nlohmann::json& j, Function& call);
+
+class FunctionBuilder {
+public:
+    FunctionBuilder(
+        const std::vector<Type> _input_types, const std::vector<Type> _output_types
+    );
+    FunctionBuilder(const Function& function);
+    Value input(int index) const;
+    ValueVec input_range(int start_index, int end_index) const;
+    void output(int index, Value value);
+    void output_range(int start_index, const ValueVec& values);
+    Value
+    global(const std::string& name, DataType dtype, const std::vector<int>& shape);
+    ValueVec instruction(const std::string& name, const ValueVec& args);
+    ValueVec instruction(InstructionPtr instruction, const ValueVec& args);
+    Function function();
+
+    Value sum(const ValueVec& values);
+    Value product(const ValueVec& values);
+
+#include "function_builder_mixin.h"
+
+private:
+    std::vector<Type> output_types;
+    ValueVec inputs;
+    std::vector<std::optional<Value>> outputs;
+    std::map<LiteralValue, Value> literals;
+    ValueVec locals;
+    std::unordered_map<std::string, Value> globals;
+    std::vector<InstructionCall> instructions;
+    std::map<std::vector<std::size_t>, std::vector<std::size_t>> instruction_cache;
+    std::vector<int> local_sources;
+    std::vector<bool> instruction_used;
+
+    void register_local(Value& val);
+};
+
+} // namespace madspace