mlx-sys-burn 0.2.2

Low-level interface and binding generation for the mlx library (fork with additional operations for burn-mlx)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

#ifndef MLX_IO_PRIVATE_H
#define MLX_IO_PRIVATE_H

#include <iostream>
#include "mlx/mlx.h"

namespace {

class CReader : public mlx::core::io::Reader {
 public:
  void* desc;
  mlx_io_vtable vtable;

  CReader(void* desc, mlx_io_vtable vtable) : desc(desc), vtable(vtable) {};
  virtual bool is_open() const override {
    return vtable.is_open(desc);
  };
  virtual bool good() const override {
    return vtable.good(desc);
  };
  virtual size_t tell() override {
    return vtable.tell(desc);
  }
  virtual void seek(
      int64_t off,
      std::ios_base::seekdir way = std::ios_base::beg) override {
    switch (way) {
      case std::ios_base::beg:
        return vtable.seek(desc, off, SEEK_SET);
        break;
      case std::ios_base::cur:
        return vtable.seek(desc, off, SEEK_CUR);
        break;
      case std::ios_base::end:
        return vtable.seek(desc, off, SEEK_END);
        break;
      default:
        throw std::runtime_error("mlx_io_reader: invalid seek way");
    }
  }
  virtual void read(char* data, size_t n) override {
    return vtable.read(desc, data, n);
  };
  virtual void read(char* data, size_t n, size_t offset) override {
    return vtable.read_at_offset(desc, data, n, offset);
  };
  virtual std::string label() const override {
    return vtable.label(desc);
  };
  virtual ~CReader() {
    vtable.free(desc);
  }
};

class CWriter : public mlx::core::io::Writer {
 public:
  void* desc;
  mlx_io_vtable vtable;

  CWriter(void* desc, mlx_io_vtable vtable) : desc(desc), vtable(vtable) {};
  virtual bool is_open() const override {
    return vtable.is_open(desc);
  };
  virtual bool good() const override {
    return vtable.good(desc);
  };
  virtual size_t tell() override {
    return vtable.tell(desc);
  }
  virtual void seek(
      int64_t off,
      std::ios_base::seekdir way = std::ios_base::beg) override {
    switch (way) {
      case std::ios_base::beg:
        return vtable.seek(desc, off, SEEK_SET);
        break;
      case std::ios_base::cur:
        return vtable.seek(desc, off, SEEK_CUR);
        break;
      case std::ios_base::end:
        return vtable.seek(desc, off, SEEK_END);
        break;
      default:
        throw std::runtime_error("mlx_io_writer: invalid seek way");
    }
  }
  virtual void write(const char* data, size_t n) override {
    return vtable.write(desc, data, n);
  };
  virtual std::string label() const override {
    return vtable.label(desc);
  };
  virtual ~CWriter() {
    vtable.free(desc);
  }
};

struct creader_holder {
  std::shared_ptr<CReader> ptr;
};

inline mlx_io_reader mlx_io_reader_new_(void* uctx, mlx_io_vtable vtable) {
  return mlx_io_reader(
      {new creader_holder({std::make_shared<CReader>(uctx, vtable)})});
}

inline std::shared_ptr<CReader> mlx_io_reader_get_(mlx_io_reader d) {
  if (!d.ctx) {
    throw std::runtime_error("expected a non-empty mlx_io_reader");
  }
  return static_cast<creader_holder*>(d.ctx)->ptr;
}

inline void mlx_io_reader_free_(mlx_io_reader io) {
  if (io.ctx) {
    delete static_cast<creader_holder*>(io.ctx);
  }
}

struct cwriter_holder {
  std::shared_ptr<CWriter> ptr;
};

inline mlx_io_writer mlx_io_writer_new_(void* uctx, mlx_io_vtable vtable) {
  return mlx_io_writer(
      {new cwriter_holder({std::make_shared<CWriter>(uctx, vtable)})});
}

inline std::shared_ptr<CWriter> mlx_io_writer_get_(mlx_io_writer d) {
  if (!d.ctx) {
    throw std::runtime_error("expected a non-empty mlx_io_writer");
  }
  return static_cast<cwriter_holder*>(d.ctx)->ptr;
}

inline void mlx_io_writer_free_(mlx_io_writer io) {
  if (io.ctx) {
    delete static_cast<cwriter_holder*>(io.ctx);
  }
}

} // namespace

#endif