[Torch] Fixed calls into torch.nexus

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 <typename T>
   nxs_status tilizeAndCopyToDevice(T *data_ptr, bool blocking);

plugins/tenstorrent/tt_command.cpp

@@ -3,6 +3,7 @@
 #include <rt_buffer.h>
 
 #include <tt-metalium/bfloat16.hpp>
+#include <tt-metalium/tensor_accessor_args.hpp>
 
 #include <algorithm>
 
@@ -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));
 

plugins/tenstorrent/tt_command.h

@@ -6,11 +6,13 @@
 #include <rt_command.h>
 #include <tt_library.h>
 #include <tt_device.h>
+#include <tt_buffer.h>
 
 class TTRuntime;
 
 class TTCommand : public nxs::rt::Command<TTKernel *, nxs_int, nxs_int> {
   TTRuntime *rt;
+  std::array<TTBuffer*, NXS_KERNEL_MAX_ARGS> buffers;
 
  public:
   TTCommand(TTRuntime *rt = nullptr, TTKernel *kernel = nullptr,
@@ -23,6 +25,12 @@ class TTCommand : public nxs::rt::Command<TTKernel *, nxs_int, nxs_int> {
 
   ~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,

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<TTBuffer>(buffer_id);
   if (!buffer) return NXS_InvalidBuffer;
 
+  command->setBufferArgument(argument_index, buffer);
   return command->setScalar(argument_index, buffer->getAddress(), "Buffer", argument_settings);
 }
 

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) {