diff --git a/plugins/tenstorrent/tt_buffer.h b/plugins/tenstorrent/tt_buffer.h index 6063ad1..7ff2d0e 100644 --- a/plugins/tenstorrent/tt_buffer.h +++ b/plugins/tenstorrent/tt_buffer.h @@ -27,6 +27,7 @@ class TTBuffer : public nxs::rt::Buffer { nxs_ulong size() const { return getSizeBytes(); } nxs_uint *getAddress() { return &address; } + ttmd::MeshBuffer *getMeshBuffer() { return buffer.get(); } template nxs_status tilizeAndCopyToDevice(T *data_ptr, bool blocking); diff --git a/plugins/tenstorrent/tt_command.cpp b/plugins/tenstorrent/tt_command.cpp index 7baba88..7fd5de7 100644 --- a/plugins/tenstorrent/tt_command.cpp +++ b/plugins/tenstorrent/tt_command.cpp @@ -3,6 +3,7 @@ #include #include +#include #include @@ -50,6 +51,14 @@ nxs_status TTCommand::runCommand(TTDevice *device, nxs_int stream, ttmd::MeshWor } } + for (TTBuffer *buffer: buffers) { + if (buffer) { + NXSLOG_INFO("TensorAccessor Buffer: {}", (intptr_t)buffer); + ttm::TensorAccessorArgs(buffer->getMeshBuffer()).append_to(ctas); + } + } + + NXSLOG_INFO("Add semaphores: {}", ctas.size()); ctas.push_back(ttm::CreateSemaphore(program, cores, 0)); ctas.push_back(ttm::CreateSemaphore(program, cores, 0)); diff --git a/plugins/tenstorrent/tt_command.h b/plugins/tenstorrent/tt_command.h index 777ebba..87320b9 100644 --- a/plugins/tenstorrent/tt_command.h +++ b/plugins/tenstorrent/tt_command.h @@ -6,11 +6,13 @@ #include #include #include +#include class TTRuntime; class TTCommand : public nxs::rt::Command { TTRuntime *rt; + std::array buffers; public: TTCommand(TTRuntime *rt = nullptr, TTKernel *kernel = nullptr, @@ -23,6 +25,12 @@ class TTCommand : public nxs::rt::Command { ~TTCommand() = default; + void setBufferArgument(nxs_int argument_index, TTBuffer *buffer) { + if (argument_index < 0 || argument_index >= NXS_KERNEL_MAX_ARGS) return; + NXSLOG_INFO("Set buffer argument: {} -> {}", argument_index, (intptr_t)buffer); + buffers[argument_index] = buffer; + } + nxs_status runCommand(nxs_int stream) override { assert(0); return NXS_Success; } nxs_status runCommand(TTDevice *device, nxs_int stream, ttmd::MeshWorkload &workload, ttmd::MeshCoordinateRange &dev_range, diff --git a/plugins/tenstorrent/tt_runtime.cpp b/plugins/tenstorrent/tt_runtime.cpp index cba40c0..879fd30 100644 --- a/plugins/tenstorrent/tt_runtime.cpp +++ b/plugins/tenstorrent/tt_runtime.cpp @@ -414,6 +414,7 @@ extern "C" nxs_status NXS_API_CALL nxsSetCommandArgument(nxs_int command_id, auto buffer = rt->get(buffer_id); if (!buffer) return NXS_InvalidBuffer; + command->setBufferArgument(argument_index, buffer); return command->setScalar(argument_index, buffer->getAddress(), "Buffer", argument_settings); } diff --git a/python/src/system.cc b/python/src/system.cc index 5c373d8..25e16db 100644 --- a/python/src/system.cc +++ b/python/src/system.cc @@ -50,7 +50,7 @@ static DevPtr getPointer(PyObject *obj) { result.device_id = 0; } Py_DECREF(device_m); - if (result.runtime_name == "knexus") { + if (result.runtime_name == "nexus") { return result; } } @@ -159,9 +159,9 @@ static Buffer make_buffer(py::object tensor, Device device = Device(), } auto data_ptr = getPointer(tensor.ptr()); - if (data_ptr.runtime_name == "knexus") { - // call torch.knexus.get_knexus_buffer(obj) - auto get_knexus_buffer = import_from("torch.knexus", "get_knexus_buffer"); + if (data_ptr.runtime_name == "nexus") { + // call torch.nexus.get_nexus_buffer(obj) + auto get_knexus_buffer = import_from("torch.nexus", "get_nexus_buffer"); PyObject *get_knexus_buffer_ret = PyObject_CallFunctionObjArgs(get_knexus_buffer, tensor.ptr(), NULL); Py_DECREF(get_knexus_buffer); if (!get_knexus_buffer_ret) {