C++ wrapper

 avatar
unknown
plain_text
4 years ago
9.5 kB
4
Indexable
#include "pybind11/cast.h"

#include <Python.h>

#include "pybind11/numpy.h"
#include "pybind11/pybind11.h"
#include "pybind11/pytypes.h"
#include "pybind11/stl.h"
#include "third_party/absl/status/status.h"
#include "third_party/absl/status/statusor.h"

namespace xla {

namespace py = pybind11;

class JaxCompiledFunction {
 public:
  explicit JaxCompiledFunction(int int_field) : int_field(int_field){};
  int int_field;

  pybind11::handle AsPyHandle();
};

PyObject* JaxCompiledFunction_Type = nullptr;

struct JaxCompiledFunctionObject {
  PyObject_HEAD;
  PyObject* dict;      // Dictionary for __dict__
  PyObject* weakrefs;  // Weak references; for use by the Python interpreter.
  JaxCompiledFunction fun;
};

bool JaxCompiledFunction_Check(py::handle handle) {
  return handle.get_type() == JaxCompiledFunction_Type;
}

JaxCompiledFunction* AsCompiledFunctionUnchecked(py::handle handle) {
  return &(reinterpret_cast<JaxCompiledFunctionObject*>(handle.ptr())->fun);
}

absl::StatusOr<JaxCompiledFunction*> AsCompiledFunction(py::handle handle) {
  if (!JaxCompiledFunction_Check(handle)) {
    return absl::InvalidArgumentError("Expected a CompiledFunction");
  }
  return AsCompiledFunctionUnchecked(handle);
}

py::handle JaxCompiledFunction::AsPyHandle() {
  return reinterpret_cast<PyObject*>(reinterpret_cast<char*>(this) -
                                     offsetof(JaxCompiledFunctionObject, fun));
}

int AccessField(const JaxCompiledFunction& fun) { return fun.int_field; }

int AccessFieldPython(py::handle handle) {
  assert(JaxCompiledFunction_Check(handle));
  JaxCompiledFunction* f = AsCompiledFunctionUnchecked(handle);
  return f->int_field;
}

int AccessFieldsPython(std::vector<py::handle> handles) {
  int v = 0;
  for (const py::handle& handle : handles) {
    assert(JaxCompiledFunction_Check(handle));
    JaxCompiledFunction* f = AsCompiledFunctionUnchecked(handle);
    v += f->int_field;
  }
  return v;
}

int AccessFieldPybind11(py::handle handle) {
  return py::cast<JaxCompiledFunction*>(handle)->int_field;
}

int AccessFieldsPybind11(std::vector<py::handle> handles) {
  int v = 0;
  for (const py::handle& handle : handles) {
    v += py::cast<JaxCompiledFunction*>(handle)->int_field;
  }
  return v;
}

// For the C API: https://docs.python.org/3/c-api/typeobj.html
extern "C" {

PyObject* JaxCompiledFunction_tp_new(PyTypeObject* subtype, PyObject* args,
                                     PyObject* kwds) {
  JaxCompiledFunctionObject* self =
      reinterpret_cast<JaxCompiledFunctionObject*>(
          subtype->tp_alloc(subtype, 0));
  if (!self) return nullptr;
  self->dict = nullptr;
  self->weakrefs = nullptr;
  return reinterpret_cast<PyObject*>(self);
}

void JaxCompiledFunction_tp_dealloc(PyObject* self) {
  PyTypeObject* tp = Py_TYPE(self);
  JaxCompiledFunctionObject* o =
      reinterpret_cast<JaxCompiledFunctionObject*>(self);
  if (o->weakrefs) {
    PyObject_ClearWeakRefs(self);
  }
  Py_CLEAR(o->dict);
  o->fun.~JaxCompiledFunction();
  tp->tp_free(self);
  Py_DECREF(tp);
}

int JaxCompiledFunction_tp_traverse(PyObject* self, visitproc visit,
                                    void* arg) {
  JaxCompiledFunctionObject* o =
      reinterpret_cast<JaxCompiledFunctionObject*>(self);
  Py_VISIT(o->dict);
  // Py_VISIT(o->fun.fun().ptr());
  // Py_VISIT(o->fun.cache_miss().ptr());
  // Py_VISIT(o->fun.get_device().ptr());
  return 0;
}

int JaxCompiledFunction_tp_clear(PyObject* self) {
  JaxCompiledFunctionObject* o =
      reinterpret_cast<JaxCompiledFunctionObject*>(self);
  Py_CLEAR(o->dict);
  // o->fun.ClearPythonReferences();
  return 0;
}

// Implements the Python descriptor protocol so JIT-compiled functions can be
// used as bound methods. See:
// https://docs.python.org/3/howto/descriptor.html#functions-and-methods
PyObject* JaxCompiledFunction_tp_descr_get(PyObject* self, PyObject* obj,
                                           PyObject* type) {
  if (obj == nullptr || obj == Py_None) {
    Py_INCREF(self);
    return self;
  }
  return PyMethod_New(self, obj);
}

// Support d = instance.__dict__.
PyObject* JaxCompiledFunction_get_dict(PyObject* self, void*) {
  JaxCompiledFunctionObject* o =
      reinterpret_cast<JaxCompiledFunctionObject*>(self);
  if (!o->dict) {
    o->dict = PyDict_New();
  }
  Py_XINCREF(o->dict);
  return o->dict;
}

int JaxCompiledFunction_set_dict(PyObject* self, PyObject* new_dict, void*) {
  JaxCompiledFunctionObject* o =
      reinterpret_cast<JaxCompiledFunctionObject*>(self);
  if (!PyDict_Check(new_dict)) {
    PyErr_Format(PyExc_TypeError,
                 "__dict__ must be set to a dictionary, not a '%s'",
                 Py_TYPE(new_dict)->tp_name);
    return -1;
  }
  Py_INCREF(new_dict);
  Py_CLEAR(o->dict);
  o->dict = new_dict;
  return 0;
}

static PyGetSetDef JaxCompiledFunction_tp_getset[] = {
    // Having a __dict__ seems necessary to allow !functool.wraps to override
    // __doc__.
    {const_cast<char*>("__dict__"), JaxCompiledFunction_get_dict,
     JaxCompiledFunction_set_dict, nullptr, nullptr},
    {nullptr, nullptr, nullptr, nullptr, nullptr}};

PyObject* JaxCompiledFunction_tp_call(PyObject* self, PyObject* args,
                                      PyObject* kwargs) {
  // tensorflow::profiler::TraceMe traceme("JaxCompiledFunction::tp_call");
  JaxCompiledFunctionObject* o =
      reinterpret_cast<JaxCompiledFunctionObject*>(self);
  absl::optional<py::kwargs> py_kwargs;
  if (kwargs) {
    py_kwargs = py::reinterpret_borrow<py::kwargs>(kwargs);
  }
  auto obj = py::object(py::str("ddede"));
  return obj.ptr();
  // try {
  //   xla::StatusOr<py::object> out = o->fun.Call(args, std::move(py_kwargs));
  //   if (!out.ok()) {
  //     PyErr_SetString(PyExc_ValueError, out.status().ToString().c_str());
  //     return nullptr;
  //   }
  //   return out.ValueOrDie().release().ptr();
  // } catch (py::error_already_set& e) {
  //   e.restore();
  //   return nullptr;
  // } catch (py::cast_error& e) {
  //   PyErr_SetString(PyExc_ValueError, e.what());
  //   return nullptr;
  // } catch (std::invalid_argument& e) {
  //   PyErr_SetString(PyExc_ValueError, e.what());
  //   return nullptr;
  // }
}

}  // extern "C"

py::object MakeJaxCompiledFunction(int int_field) {
  py::object obj = py::reinterpret_steal<py::object>(JaxCompiledFunction_tp_new(
      reinterpret_cast<PyTypeObject*>(JaxCompiledFunction_Type), nullptr,
      nullptr));
  JaxCompiledFunctionObject* buf =
      reinterpret_cast<JaxCompiledFunctionObject*>(obj.ptr());

  new (&buf->fun) JaxCompiledFunction(int_field);
  return obj;
}

std::vector<py::object> CreateManyObjectsPython(int num_obj) {
  std::vector<py::object> results;
  results.reserve(num_obj);
  for (int i = 0; i < num_obj; ++i) {
    results.push_back(MakeJaxCompiledFunction(i));
  }
  return results;
}

std::vector<py::object> CreateManyObjectsPybind11(int num_obj) {
  std::vector<py::object> results;
  results.reserve(num_obj);
  for (int i = 0; i < num_obj; ++i) {
    results.push_back(py::cast(JaxCompiledFunction(i)));
  }
  return results;
}

PYBIND11_MODULE(cast, m) {
  // Initializes the NumPy API for the use of the types module.
  m.doc() = "Jax XLA function";

  // We need to use heap-allocated type objects because we want to add
  // additional methods dynamically.
  py::object cfun;
  {
    py::str name = py::str("CompiledFunction");
    py::str qualname = py::str("CompiledFunction");
    PyHeapTypeObject* heap_type = reinterpret_cast<PyHeapTypeObject*>(
        PyType_Type.tp_alloc(&PyType_Type, 0));
    // Caution: we must not call any functions that might invoke the GC until
    // PyType_Ready() is called. Otherwise the GC might see a half-constructed
    // type object.
    assert(heap_type);  //  << "Unable to create heap type object";
    heap_type->ht_name = name.release().ptr();
    heap_type->ht_qualname = qualname.release().ptr();
    PyTypeObject* type = &heap_type->ht_type;
    type->tp_name = "CompiledFunction";
    type->tp_basicsize = sizeof(JaxCompiledFunctionObject);
    type->tp_flags =
        Py_TPFLAGS_DEFAULT | Py_TPFLAGS_HEAPTYPE | Py_TPFLAGS_HAVE_GC;
    type->tp_new = JaxCompiledFunction_tp_new;
    type->tp_dealloc = JaxCompiledFunction_tp_dealloc;
    type->tp_dictoffset = offsetof(JaxCompiledFunctionObject, dict);
    type->tp_traverse = JaxCompiledFunction_tp_traverse;
    type->tp_clear = JaxCompiledFunction_tp_clear;
    type->tp_weaklistoffset = offsetof(JaxCompiledFunctionObject, weakrefs);
    type->tp_getset = JaxCompiledFunction_tp_getset;
    type->tp_descr_get = JaxCompiledFunction_tp_descr_get;
    type->tp_call = JaxCompiledFunction_tp_call;
    assert(PyType_Ready(type) == 0);
    JaxCompiledFunction_Type = reinterpret_cast<PyObject*>(type);
    cfun = py::reinterpret_borrow<py::object>(JaxCompiledFunction_Type);
  }

  py::class_<JaxCompiledFunction>(m, "JaxCompiledFunction")
      .def(py::init<int>())
      .def_readwrite("int_field", &JaxCompiledFunction::int_field);

  m.def("MakeJaxCompiledFunction", &MakeJaxCompiledFunction);
  m.def("AccessFieldPython", &AccessFieldPython);
  m.def("AccessFieldPybind11", &AccessFieldPybind11);
  m.def("AccessFieldsPython", &AccessFieldsPython);
  m.def("AccessFieldsPybind11", &AccessFieldsPybind11);

  m.def("CreateManyObjectsPython", &CreateManyObjectsPython);
  m.def("CreateManyObjectsPybind11", &CreateManyObjectsPybind11);

}  // NOLINT(readability/fn_size)

}  // namespace xla
Editor is loading...