Repository: Xilinx/inference-server (language: python)

Issue:
Support using custom classes for memory in the pool
The global memory pool and default allocator (#149) allocates memory given just a size. This memory is handled just as a raw pointer. For some workers (e.g. xmodel, tfzendnn and ptzendnn), the inputs and outputs of running an inference must be in custom "tensorbuffer" style classes that require a shape and datatype on construction. Since the batcher is now creating batches in general memory, it requires the worker to copy data into these objects before inference.

This was previously supported with custom buffer classes.


Discussion / hints:

Fix patch:
diff --git a/docs/architecture.rst b/docs/architecture.rst
index 779f4ade3..4a9c09e80 100644
--- a/docs/architecture.rst
+++ b/docs/architecture.rst
@@ -31,7 +31,7 @@ Overview
 --------
 
 :numref:`architecture_overview` shows a high-level view of the Server's architecture.
-At the top, we have a client who can make a request to the Server using HTTP/REST, WebSocket, or the C++ API [#f1]_.
+At the top, we have a client who can make a request to the Server using HTTP/REST, WebSocket, gRPC or the C++ API [#f1]_.
 This request must be made to a particular worker that must be active in the Server.
 Each received request is passed to a batcher associated with the targeted worker independent of its origin.
 The batcher will combine individual requests into a batch and pass it to the worker for execution.
@@ -45,11 +45,8 @@ Ingestion
 
 There are a number of ways to get data into the system for inference.
 In general, each protocol has a custom interface to the client and requires explicit handling in the Server to accept this initial request.
-After receiving the client's request in the Server, each supported protocols' handler packs the request into an implementation of a ``RequestContainer`` object.
-This virtual class is defined in the Server and provides the batcher with a standard way of interacting with data from different ingestion protocols.
-Since all protocols push an implementation of this class to the batcher, they are treated equally by the rest of the Server.
-The only protocol-specific logic applies at the end of the worker when it replies back to the client in a protocol-specific manner.
-This structure also enables easy extension of the Server to add new protocols by extending the ``RequestContainer`` class.
+After receiving the client's request in the Server, each supported protocols' handler converts the incoming data into an ``InferenceRequest`` object.
+This object, plus some metadata, is packed into a ``RequestContainer`` object and passed to the appropriate batcher.
 
 API
 ^^^
@@ -69,6 +66,7 @@ It also doesn't have the notion of load-time parameters.
 In our case, there are parameters we need to pass at load-time, which results in potentially different endpoints if multiple workers with different configurations are loaded at once.
 To maintain compatibility, we do guarantee that the first worker loaded for a particular model, independent of configuration, will be at the same endpoint as the load request.
 Therefore, a KServe client is free to ignore the contents of the response and make requests to the endpoint without resulting in errors.
+There's also a ``modelLoad`` API which behaves more similarly to how KServe expects and is intended for use with using the server with a model repository.
 
 We currently do not support the optional version information associated with a model defined in the KServe specification.
 After a particular worker is loaded, inference requests can be made to it by constructing the appropriate request object and sending it to the prediction endpoint.
@@ -95,6 +93,16 @@ The WebSocket API is custom.
 At this time, the client provides a URL to a video that the worker will retrieve and analyze frame-by-frame and send back to the client but this is subject to change.
 The WebSocket server code is in ``src/amdinfer/servers/websocket_server.*``.
 
+gRPC
+^^^^
+
+The gRPC functionality is provided through a custom implementation of a gRPC server.
+There are some examples in the `gRPC repository on Github <https://github.com/grpc/grpc/blob/master/examples/cpp>`__ that can be used for reference.
+It makes new dynamic objects to keep track of the incoming requests and a state machine is embedded inside to track state.
+A pointer to this object, ``CallData``, is put into the callback for the request so when the worker finishes this request, it will use it to respond to the request.
+After the response, the state machine is marked to finish and the object deallocates itself.
+The gRPC server code is in ``src/amdinfer/servers/grpc_server.*``.
+
 C++ API
 ^^^^^^^
 
@@ -137,17 +145,16 @@ This configuration is determined at compile-time and built into the definition o
 A ``Batch`` is made up of three basic components: ``InferenceRequest`` objects and input/output buffers.
 ``InferenceRequest`` objects are direct C++ implementations of the information present in the KServe API of an inference request structured in a similar format.
 They provide an ingestion-agnostic data format that all workers can process.
-The worker group that the batcher is attached to provides a set of input/output buffers from a pool of buffers that it allocates.
-These buffers are available in a queue for that batcher to pull from when it has incoming requests to batch together.
-Most commonly, each buffer can be used to represent one batch-size worth of contiguous memory but its exact nature depends on the buffer implementation that the worker is using.
+The batcher is responsible for getting memory from the memory pool that is sized appropriately to contain the incoming batch size of data as a buffer.
+The worker provides a list of allocators that the batcher is allowed to use when requesting memory.
+Most commonly, each buffer can be used to represent one batch-size worth of contiguous memory but its exact nature depends on the buffer implementation that the worker has requested.
 In this case, the batcher's job is to take individual requests and move its data into one slot of this buffer and construct the corresponding ``InferenceRequest`` object.
 Batchers have some flexibility with how these batches are constructed, which is why multiple batcher implementations are possible and supported in the AMD Inference Server.
 For example, one batcher may allow partial batches to be pushed on after enough time whereas this may not be allowed by another batcher.
 
-Batchers use the ``RequestContainer`` object's ``getRequest()`` method to help create batches.
-This method must be implemented by each interface and governs how, given some buffers and counters, the particular ingestion method's data should be converted to an ``InferenceRequest`` and its data is copied over to the buffers.
-THis method allows batchers to process all ingestion methods without knowing about the details of how the data may be stored internally in the ``RequestContainer``.
-
+Assuming that contiguous batches are expected, the batcher should request memory from the pool on the first request of a new batch.
+As new requests come in, their data is copied over to the newly allocated memory so it's contiguous for downstream processing.
+The memory that was used initially by the ingestion layer can now be freed.
 
 .. _architectureWorkers:
 
@@ -171,14 +178,12 @@ This class defines the lifecycle methods of the worker that are called by the Se
 This lifecycle is defined as follows:
 
 #.	``init()``: perform low-cost initialization of the worker
-#.	``allocate()``: allocate memory buffers that are used to hold input and output data for the worker. `Buffering` is further discussed below.
 #.	``acquire()``: acquire any hardware accelerators/resources and/or perform any high-cost initialization for the worker
 #.	``run()``: the main body of the worker performs the chosen computations on incoming batches
 #.	``release()``: release any hardware accelerators/resources
-#.	``deallocate()``: free the memory buffers allocated by this worker
 #.	``destroy()``: perform any final operations prior to shutdown
 
-The first three steps set up the worker while the latter three tear it down and are performed in this order by the Server.
+The first two steps set up the worker while the latter two tear it down and are performed in this order by the Server.
 The body of these methods must be provided by each worker implementation in the corresponding ``doX()`` methods (e.g. ``doInit()``).
 At load-time, the server will create an instance of the worker using its ``getWorker()`` method:
 
@@ -188,14 +193,22 @@ At load-time, the server will create an instance of the worker using its ``getWo
         amdinfer::workers::Worker* getWorker() { return new amdinfer::workers::MyWorkerClass(); }
     }
 
-This instance is saved internally and the first three methods above are called to initialize the worker.
+This instance is saved internally and the first two methods above are called to initialize the worker.
 The worker's batcher is also started by the server at this time.
 Finally, the worker's ``run()`` method is started as a separate thread with the batcher's output queue passed as the input queue to the worker.
 This method performs the body of the work.
 In an infinite loop, this method should wait for incoming batches from the worker's input queue, process the requests, and respond to the clients.
 
-To unload a worker, the Manager sends a ``nullptr`` to the worker, which should terminate the ``run()`` thread.
-This thread is joined and the last three lifecycle methods are called to safely clean up the worker.
+To unload a worker, the State sends a ``nullptr`` to the worker, which should terminate the ``run()`` thread.
+This thread is joined and the last two lifecycle methods are called to safely clean up the worker.
+
+Workers must also define a ``getAllocators()`` method to choose which allocators can be used by the batcher when it's preparing the incoming batch.
+
+.. code-block:: c++
+
+    std::vector<MemoryAllocators> MyWorkerClass::getAllocators() const {
+        return {MemoryAllocators::Cpu};
+    }
 
 Improving Performance
 ^^^^^^^^^^^^^^^^^^^^^
@@ -237,8 +250,15 @@ XModel
     ..
 
 As perhaps the most complex worker thus far, the architecture of the XModel worker is examined here in greater detail.
-The XModel worker is intended to run an arbitrary XModel specified by the user on a Xilinx FPGA [#f3]_.
+The XModel worker is intended to run an arbitrary XModel specified by the user on a Xilinx FPGA [#f2]_.
 We take a look at the lifecycle of this worker in the following sections.
+The XModel worker is using the ``VartTensor`` allocator.
+
+.. code-block:: c++
+
+    std::vector<MemoryAllocators> Xmodel::getAllocators() const {
+        return {MemoryAllocators::VartTensor};
+    }
 
 Initialization
 """"""""""""""
@@ -249,13 +269,6 @@ In the future, we may support running an arbitrary number of subgraphs but this
 Using this subgraph, we create a *Runner*, which is a thread-safe object defined in the Vitis-AI runtime and is responsible for submitting requests to the FPGA.
 These objects are all saved as part of the internal state of the worker.
 
-Allocation
-""""""""""
-
-We use a special buffer backend for the this worker: the VartTensorBuffer.
-This custom type provides better compatibility with using the Runner as that expects ``vart::TensorBuffer`` objects to pass data to the FPGA.
-Therefore, this worker creates buffers using this backend and passes them to the Manager.
-
 Acquisition
 """""""""""
 
@@ -282,57 +295,25 @@ Cleanup
 There is almost no special cleanup required as the Vitis-AI objects that are part of the worker's state are smart pointers and are cleaned by the worker's destructor.
 THe only non-default implementation of the clean-up functions is to stop the internal thread pool and join the threads.
 
-Buffering
----------
-
-.. _fig_buffering:
-.. figure:: assets/buffer_lifecycle.png
-    :alt: Diagram showing the buffer lifecycle
-    :height: 400px
-    :align: left
-
-    The buffer lifecycle
-
-Buffers are used to hold data internally within the server after receiving a request.
-The implementations of buffers are in ``src/amdinfer/buffers``.
-
-The lifecycle of buffers is shown in :numref:`fig_buffering`.
-In ``allocate()``, the worker creates a buffer pool made of some number of buffers.
-Using a buffer pool saves the cost of constantly allocating dynamic memory for each new request.
-Instead, we can reuse the same set of buffers that are allocated by the worker at one time.
-They are initially provided by the worker to the Manager which maintains a queue of buffers for storing the coalesced requests for one batch.
-The buffers of all the workers in one group are maintained in this common queue.
-They are consumed from the pool as the batcher creates batches and then the worker returns them to the pool after finishing work on a batch.
-If the batcher needs a buffer but there are none available, the batcher can block execution until a buffer becomes available.
-Thus, the number of buffers in the pool controls the number of active batches for a particular worker group.
-Currently, there's no mechanism to change the number of buffers in the pool at run-time short of allocating a new worker or sending a large request that forces the automatic allocation of more buffers.
-In the future, the number of buffers may be controllable from the Manager and dynamically managed depending on the number of requests.
-
-Multiple kinds of buffer backends are supported by providing the appropriate wrappers.
-For example, a simple implementation may use buffers allocated in CPU memory.
-For more advanced sharing of data and to minimize data movement, buffers may be allocated in shared memory or on hardware accelerators.
-Buffer backends extend the ``Buffer`` class.
-They provide methods to write different data types into the buffer and access the underlying data at some offset.
-
-Manager
--------
+Shared State
+------------
 
-The shared state of the AMD Inference Server is maintained by the Manager: the active workers, their buffer pools, the endpoints and load-time parameters associated with them and is visualized+ in :numref:`architecture_detail`.
-This information enables the ingestion protocols to query the Manager to retrieve a pointer to the correct batcher to use to push the ``RequestContainer`` object to the right one corresponding to the targeted worker.
-To manage multiple versions of workers that may be running with different configurations, the Manager stores the load-time parameters, if any, and compares new parameters with ones its seen before to determine whether the newly loaded worker should be part of an existing worker group or a new one.
+The shared state of the AMD Inference Server is maintained by the ``shared_state``—the active workers, their buffer pools, the endpoints, the load-time parameters associated with them, global memory pool and the repository—is visualized in :numref:`architecture_detail`.
+This information enables the ingestion protocols to query the State to retrieve a pointer to the correct batcher to use to push the ``RequestContainer`` object to the right one corresponding to the targeted worker.
+To manage multiple versions of workers that may be running with different configurations, the State stores the load-time parameters, if any, and compares new parameters with ones its seen before to determine whether the newly loaded worker should be part of an existing worker group or a new one.
 In the case that it's assigned to an existing worker group, the previously allocated endpoint is returned to the client.
 If a new worker group is created, a new endpoint is reserved for this worker group and returned to the client.
-The implementation is in ``src/amdinfer/core/manager.*``.
+The implementation is in ``src/amdinfer/core/shared_state.*``, ``src/amdinfer/core/endpoints.*`` and ``src/amdinfer/core/model_repository.*``.
 
-Loading a new worker results in the creation of a new ``WorkerInfo`` (see ``src/amdinfer/core/worker_info.*``) object which the Manager uses internally to hold all the information associated with the worker.
+Loading a new worker results in the creation of a new ``WorkerInfo`` (see ``src/amdinfer/core/worker_info.*``) object which the State uses internally to hold all the information associated with the worker.
 The worker class instance, its batcher, and its buffer pool are all stored in this object.
 The ``WorkerInfo`` object provides two methods to create new workers: its constructor and an ``addAndStartWorker()`` method.
 The former is used for a brand-new worker and creates queues for the buffer pools and initializes the private members of the class.
 The latter loads the shared library associated with the worker, creates and saves the instance of the worker class, and starts its ``run()`` method in a new thread.
 
-The Manager also provides methods to safely modify the shared state such as loading or unloading new workers.
+The State also provides methods to safely modify the shared state such as loading or unloading new workers.
 Such actions must be taken with care because there are many threads that may need to modify state or make decisions based on the current state.
-The Manager uses a queue and a separate thread for this purpose.
+The State uses a queue and a separate thread for this purpose.
 All methods that modify state enqueue requests to this queue.
 These methods may be called from a multi-threaded context and so multiple duplicate or contradictory requests are possible.
 The queue enforces serialization and defines an ordering for all incoming requests so they can be processed in this order by the new thread.
@@ -367,13 +348,12 @@ Look at :ref:`metrics:metrics` for more information.
 Tracing
 ^^^^^^^
 
-The Server uses :github:`jaeger-client-cpp <jaegertracing/jaeger-client-cpp>` [#f2]_ to provide tracing.
+The Server uses :github:`opentelemetry-cpp <open-telemetry/opentelemetry-cpp>` to provide tracing.
 Tracing tracks the time taken for different sections of the architecture to process a single request.
-This data can be visualized in the Jaeger UI.
+This data can be visualized in the Jaeger UI or Grafana.
 Tracing data can be disabled at compile-time with a CMake option.
 
 Look at :ref:`tracing:tracing` for more information.
 
-.. [#f1] Some methods are only available through HTTP at this time. Using the C++ API requires compiling an application linked against ``libamdinfer.so`` rather than making requests to a server.
-.. [#f2] This library is deprecating and will be replaced with OpenTelemetry as recommended by Jaeger.
-.. [#f3] There are currently some restrictions on what may be run such as the number of input/output tensors.
+.. [#f1] Using the C++ API requires compiling an application linked against ``libamdinfer.so`` rather than making requests to a server.
+.. [#f2] There are currently some restrictions on what may be run such as the number of input tensors.
diff --git a/docs/assets/buffer_lifecycle.png b/docs/assets/buffer_lifecycle.png
deleted file mode 100644
index 59a340deb..000000000
--- a/docs/assets/buffer_lifecycle.png
+++ /dev/null
@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:aba28f0b81199a13e3bbe68b8720beebdde4dad7f704479ed108420eeb1a3095
-size 36245
diff --git a/src/amdinfer/batching/CMakeLists.txt b/src/amdinfer/batching/CMakeLists.txt
index 7abe885f4..01519ddf1 100644
--- a/src/amdinfer/batching/CMakeLists.txt
+++ b/src/amdinfer/batching/CMakeLists.txt
@@ -39,7 +39,7 @@ endforeach()
 target_link_libraries(softBatcher INTERFACE timer)
 
 add_library(batching STATIC ${BATCHING_OBJS})
-target_link_libraries(batching INTERFACE batch batcher ${targets})
+target_link_libraries(batching INTERFACE batch batcher ${targets} buffers)
 set_target_options(batching)
 
 set(BATCHING_OBJS ${BATCHING_OBJS} PARENT_SCOPE)
diff --git a/src/amdinfer/batching/hard.cpp b/src/amdinfer/batching/hard.cpp
index 7f6fdc823..8c332e4b7 100644
--- a/src/amdinfer/batching/hard.cpp
+++ b/src/amdinfer/batching/hard.cpp
@@ -95,9 +95,7 @@ void HardBatcher::doRun(const std::vector<MemoryAllocators>& allocators) {
         // output_buffers.reserve(output_sizes.size());
         // std::vector<size_t> output_offset(output_buffers.size(), 0);
         for (const auto& input : inputs) {
-          input_buffers.push_back(pool_->get(
-            allocators,
-            input.getSize() * input.getDatatype().size() * batch_size_));
+          input_buffers.push_back(pool_->get(allocators, input, batch_size_));
         }
         // for(const auto& tensor_size : output_sizes) {
         //   output_buffers.push_back(pool_->get(allocators, tensor_size));
diff --git a/src/amdinfer/batching/soft.cpp b/src/amdinfer/batching/soft.cpp
index 19347034e..d256a299c 100644
--- a/src/amdinfer/batching/soft.cpp
+++ b/src/amdinfer/batching/soft.cpp
@@ -127,9 +127,7 @@ void SoftBatcher::doRun(const std::vector<MemoryAllocators>& allocators) {
         // output_buffers.reserve(output_sizes.size());
         // std::vector<size_t> output_offset(output_buffers.size(), 0);
         for (const auto& input : inputs) {
-          input_buffers.push_back(pool_->get(
-            allocators,
-            input.getSize() * input.getDatatype().size() * batch_size_));
+          input_buffers.push_back(pool_->get(allocators, input, batch_size_));
         }
         // for(const auto& tensor_size : output_sizes) {
         //   output_buffers.push_back(pool_->get(allocators, tensor_size));
diff --git a/src/amdinfer/buffers/CMakeLists.txt b/src/amdinfer/buffers/CMakeLists.txt
index dd4b8c9e5..6cfeb3deb 100644
--- a/src/amdinfer/buffers/CMakeLists.txt
+++ b/src/amdinfer/buffers/CMakeLists.txt
@@ -15,6 +15,10 @@
 
 list(APPEND targets cpu)
 
+if(${AMDINFER_ENABLE_VITIS})
+  list(APPEND targets vart_tensor)
+endif()
+
 add_library(buffer OBJECT buffer.cpp)
 target_include_directories(buffer PRIVATE ${AMDINFER_INCLUDE_DIRS})
 set_target_options(buffer)
diff --git a/src/amdinfer/buffers/buffer.cpp b/src/amdinfer/buffers/buffer.cpp
index 63140af46..1cd094e1f 100644
--- a/src/amdinfer/buffers/buffer.cpp
+++ b/src/amdinfer/buffers/buffer.cpp
@@ -24,9 +24,13 @@
 
 namespace amdinfer {
 
+Buffer::Buffer(MemoryAllocators allocator) : allocator_(allocator) {}
+
 size_t Buffer::write(void* data, size_t offset, size_t size) {
   std::memcpy(this->data(offset), data, size);
   return offset + size;
 }
 
+MemoryAllocators Buffer::getAllocator() const { return allocator_; }
+
 }  // namespace amdinfer
diff --git a/src/amdinfer/buffers/buffer.hpp b/src/amdinfer/buffers/buffer.hpp
index f3476c5d7..0b7c98af7 100644
--- a/src/amdinfer/buffers/buffer.hpp
+++ b/src/amdinfer/buffers/buffer.hpp
@@ -26,6 +26,8 @@
 #include <cstring>  // for memcpy
 #include <string>   // for string
 
+#include "amdinfer/core/memory_pool/memory_allocator.hpp"
+
 // IWYU is creating a cycle with adding/removing this header
 // IWYU pragma: no_include <algorithm>
 
@@ -37,6 +39,8 @@ namespace amdinfer {
  */
 class Buffer {
  public:
+  explicit Buffer(MemoryAllocators allocator);
+
   /// Destroy the Buffer object
   virtual ~Buffer() = default;
 
@@ -86,6 +90,11 @@ class Buffer {
       return this->write(&value, offset, sizeof(T));
     }
   }
+
+  MemoryAllocators getAllocator() const;
+
+ private:
+  MemoryAllocators allocator_;
 };
 
 }  // namespace amdinfer
diff --git a/src/amdinfer/buffers/cpu.cpp b/src/amdinfer/buffers/cpu.cpp
index a99d6d0bb..44cb2517a 100644
--- a/src/amdinfer/buffers/cpu.cpp
+++ b/src/amdinfer/buffers/cpu.cpp
@@ -26,7 +26,7 @@
 namespace amdinfer {
 
 CpuBuffer::CpuBuffer(void* data, MemoryAllocators allocator)
-  : allocator_(allocator), data_(static_cast<std::byte*>(data)) {}
+  : Buffer(allocator), data_(static_cast<std::byte*>(data)) {}
 
 void* CpuBuffer::data(size_t offset) { return data_ + offset; }
 
diff --git a/src/amdinfer/buffers/cpu.hpp b/src/amdinfer/buffers/cpu.hpp
index 438b0513e..ad5a40769 100644
--- a/src/amdinfer/buffers/cpu.hpp
+++ b/src/amdinfer/buffers/cpu.hpp
@@ -25,13 +25,14 @@
 
 #include "amdinfer/buffers/buffer.hpp"   // IWYU pragma: export
 #include "amdinfer/core/data_types.hpp"  // for DataType
-#include "amdinfer/core/memory_pool/pool.hpp"
-#include "amdinfer/util/queue.hpp"  // for BufferPtrsQueue
+#include "amdinfer/util/queue.hpp"       // for BufferPtrsQueue
 
 namespace amdinfer {
 
+enum class MemoryAllocators;
+
 /**
- * @brief VectorBuffer uses vectors for storing data
+ * @brief CpuBuffer uses vectors for storing data
  *
  */
 class CpuBuffer : public Buffer {
@@ -51,10 +52,7 @@ class CpuBuffer : public Buffer {
    */
   void* data(size_t offset) override;
 
-  MemoryAllocators getAllocator() const { return allocator_; }
-
  private:
-  MemoryAllocators allocator_;
   std::byte* data_;
 };
 
diff --git a/src/amdinfer/buffers/vart_tensor.cpp b/src/amdinfer/buffers/vart_tensor.cpp
new file mode 100644
index 000000000..86ee92de0
--- /dev/null
+++ b/src/amdinfer/buffers/vart_tensor.cpp
@@ -0,0 +1,58 @@
+// Copyright 2023 Advanced Micro Devices, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/**
+ * @file
+ * @brief Implements the VartTensorBuffer class
+ */
+
+#include "amdinfer/buffers/vart_tensor.hpp"
+
+#include <vart/tensor_buffer.hpp>
+#include <vector>                 // for vector
+#include <xir/tensor/tensor.hpp>  // for Tensor
+
+namespace amdinfer {
+
+VartTensorBuffer::VartTensorBuffer(void* data, MemoryAllocators allocator)
+  : Buffer(allocator), data_(static_cast<vart::TensorBuffer*>(data)) {}
+
+void* VartTensorBuffer::data(size_t offset) {
+  const auto* tensor = data_->get_tensor();
+
+  // Some DPUs need a shape argument to data() to get the data properly.
+  // This argument should be the same size as the tensor (by default,
+  // [batch, h, w, c]). This first argument is the batch index. The other
+  // indices should be zero to get the start of the batch
+  std::vector<int> shape(tensor->get_shape().size(), 0);
+
+  auto dims = tensor->get_shape();
+  auto size = tensor->get_shape().size();
+
+  // convert the offset to a shape based on the tensor shape
+  for (auto k = 0U; k < size; k++) {
+    auto stride = 1;
+    for (auto m = k + 1; m < size; m++) {
+      stride *= dims[m];
+    }
+    shape[k] = static_cast<int>(offset / stride);
+    offset -= shape[k] * stride;
+  }
+
+  return reinterpret_cast<void*>(data_->data(shape).first);
+}
+
+vart::TensorBuffer* VartTensorBuffer::getTensorBuffer() { return data_; }
+
+}  // namespace amdinfer
diff --git a/src/amdinfer/buffers/vart_tensor.hpp b/src/amdinfer/buffers/vart_tensor.hpp
new file mode 100644
index 000000000..a6f91ac3a
--- /dev/null
+++ b/src/amdinfer/buffers/vart_tensor.hpp
@@ -0,0 +1,58 @@
+// Copyright 2023 Advanced Micro Devices, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/**
+ * @file
+ * @brief Defines the VartTensorBuffer class
+ */
+
+#ifndef GUARD_AMDINFER_BUFFERS_VART_TENSOR
+#define GUARD_AMDINFER_BUFFERS_VART_TENSOR
+
+#include <cstddef>  // for size_t
+
+#include "amdinfer/buffers/buffer.hpp"
+
+namespace vart {
+class TensorBuffer;
+}  // namespace vart
+
+namespace amdinfer {
+
+class VartTensorBuffer : public Buffer {
+ public:
+  /**
+   * @brief Construct a new Vector Buffer object
+   *
+   * @param elements number of elements to store
+   * @param data_type type of element to store
+   */
+  VartTensorBuffer(void* data, MemoryAllocators allocator);
+
+  /**
+   * @brief Returns a pointer to the underlying data
+   *
+   * @return void*
+   */
+  void* data(size_t offset) override;
+
+  vart::TensorBuffer* getTensorBuffer();
+
+ private:
+  vart::TensorBuffer* data_;
+};
+
+}  // namespace amdinfer
+
+#endif  // GUARD_AMDINFER_BUFFERS_VART_TENSOR
diff --git a/src/amdinfer/clients/native.cpp b/src/amdinfer/clients/native.cpp
index ffbbbdc0b..059867b59 100644
--- a/src/amdinfer/clients/native.cpp
+++ b/src/amdinfer/clients/native.cpp
@@ -26,6 +26,7 @@
 #include <string>   // for string
 #include <utility>  // for move
 
+#include "amdinfer/buffers/buffer.hpp"   // for BufferPtr
 #include "amdinfer/build_options.hpp"    // for AMDINFER_ENABLE_TRACING
 #include "amdinfer/core/exceptions.hpp"  // for invalid_argument
 #include "amdinfer/core/parameters.hpp"  // for ParameterMap
@@ -89,7 +90,7 @@ InferenceRequestPtr getRequest(const InferenceRequest& req,
   int i = 0;
   for (const auto& input : inputs) {
     auto size = input.getSize() * input.getDatatype().size();
-    auto buffer = pool->get({MemoryAllocators::Cpu}, size);
+    auto buffer = pool->get({MemoryAllocators::Cpu}, input, 1);
     buffer->write(input.getData(), 0, size);
     request->setInputTensorData(i, buffer->data(0));
     i++;
diff --git a/src/amdinfer/core/data_types_internal.cpp b/src/amdinfer/core/data_types_internal.cpp
index 93546d9d8..a56b8ff84 100644
--- a/src/amdinfer/core/data_types_internal.cpp
+++ b/src/amdinfer/core/data_types_internal.cpp
@@ -92,13 +92,13 @@ xir::DataType mapTypeToXir(DataType type) {
     case DataType::Uint16:
     case DataType::Uint32:
     case DataType::Uint64:
-      retval.type = xir::DataType::UINT;
+      retval.type = xir::DataType::XUINT;
       break;
     case DataType::Int8:
     case DataType::Int16:
     case DataType::Int32:
     case DataType::Int64:
-      retval.type = xir::DataType::INT;
+      retval.type = xir::DataType::XINT;
       break;
     // case DataType::Fp16 fall through to default handler
     case DataType::Fp32:
diff --git a/src/amdinfer/core/memory_pool/CMakeLists.txt b/src/amdinfer/core/memory_pool/CMakeLists.txt
index e3f0b9607..39bc8f1c8 100644
--- a/src/amdinfer/core/memory_pool/CMakeLists.txt
+++ b/src/amdinfer/core/memory_pool/CMakeLists.txt
@@ -15,6 +15,10 @@
 
 set(targets cpu_allocator pool)
 
+if(${AMDINFER_ENABLE_VITIS})
+  list(APPEND targets vart_tensor_allocator)
+endif()
+
 foreach(target ${targets})
   add_library(${target} OBJECT ${target}.cpp)
   target_include_directories(${target} PRIVATE ${AMDINFER_INCLUDE_DIRS})
@@ -23,8 +27,15 @@ foreach(target ${targets})
   list(APPEND MEMORY_POOL_OBJS $<TARGET_OBJECTS:${target}>)
 endforeach()
 
+if(${AMDINFER_ENABLE_VITIS})
+  target_link_libraries(
+    vart_tensor_allocator INTERFACE vart_tensorBuffer vart::runner
+                                    data_types_internal
+  )
+endif()
+
 add_library(memory_pool STATIC ${MEMORY_POOL_OBJS})
-target_link_libraries(memory_pool INTERFACE ${targets})
+target_link_libraries(memory_pool INTERFACE ${targets} cpuBuffer)
 set_target_options(memory_pool)
 
 set(MEMORY_POOL ${MEMORY_POOL_OBJS} PARENT_SCOPE)
diff --git a/src/amdinfer/core/memory_pool/cpu_allocator.cpp b/src/amdinfer/core/memory_pool/cpu_allocator.cpp
index 4b55bf15f..077ab6345 100644
--- a/src/amdinfer/core/memory_pool/cpu_allocator.cpp
+++ b/src/amdinfer/core/memory_pool/cpu_allocator.cpp
@@ -23,14 +23,19 @@
 #include <iostream>
 #include <vector>
 
+#include "amdinfer/buffers/cpu.hpp"
 #include "amdinfer/core/exceptions.hpp"
+#include "amdinfer/core/predict_api.hpp"
 
 namespace amdinfer {
 
 CpuAllocator::CpuAllocator(size_t block_size, size_t max_allocate)
   : max_allocate_(max_allocate), block_size_(block_size) {}
 
-void* CpuAllocator::get(size_t size) {
+BufferPtr CpuAllocator::get(const InferenceRequestInput& tensor,
+                            size_t batch_size) {
+  auto size = tensor.getSize() * tensor.getDatatype().size() * batch_size;
+
   const std::lock_guard lock{mutex_};
   auto best = headers_.end();
   const auto end = headers_.end();
@@ -46,7 +51,7 @@ void* CpuAllocator::get(size_t size) {
     if (best->size == size) {
       best->free = false;
       // std::cout << "Matched " << size << " bytes\n";
-      return best->address;
+      return std::make_unique<CpuBuffer>(best->address, MemoryAllocators::Cpu);
     }
     const auto& new_block =
       headers_.emplace(best, best->address, size, false, best->block_id);
@@ -54,7 +59,8 @@ void* CpuAllocator::get(size_t size) {
     best->size -= size;
     best->address += size;
     // std::cout << "Partitioned " << size << " bytes\n";
-    return new_block->address;
+    return std::make_unique<CpuBuffer>(new_block->address,
+                                       MemoryAllocators::Cpu);
   }
 
   auto size_to_allocate = std::max(size, block_size_);
@@ -77,7 +83,7 @@ void* CpuAllocator::get(size_t size) {
   }
 
   // std::cout << "Allocated " << size << " bytes\n";
-  return retval;
+  return std::make_unique<CpuBuffer>(retval, MemoryAllocators::Cpu);
 }
 
 void CpuAllocator::put(const void* address) {
@@ -96,22 +102,29 @@ void CpuAllocator::put(const void* address) {
   }
 
   // if the previous is free and from the same block, merge the two
-  if (auto prev = std::prev(found);
-      prev->block_id == found->block_id && prev->free) {
-    prev->size += found->size;
-    headers_.erase(found);
-    found = prev;
+  if (found != std::begin(headers_)) {
+    auto prev = std::prev(found);
+    if (prev->block_id == found->block_id && prev->free) {
+      prev->size += found->size;
+      headers_.erase(found);
+      found = prev;
+    }
   }
 
   // if the next is free and from the same block, merge the two
-  if (auto next = std::next(found);
-      next->block_id == found->block_id && next->free) {
-    next->size += found->size;
-    next->address = found->address;
-    headers_.erase(found);
+  if (found != std::prev(std::end(headers_))) {
+    auto next = std::next(found);
+    if (next->block_id == found->block_id && next->free) {
+      next->size += found->size;
+      next->address = found->address;
+      headers_.erase(found);
+    } else {
+      found->free = true;
+    }
   } else {
     found->free = true;
   }
+
   // std::cout << "Freed memory\n";
 }
 
diff --git a/src/amdinfer/core/memory_pool/cpu_allocator.hpp b/src/amdinfer/core/memory_pool/cpu_allocator.hpp
index 15b98297f..7db3a4f40 100644
--- a/src/amdinfer/core/memory_pool/cpu_allocator.hpp
+++ b/src/amdinfer/core/memory_pool/cpu_allocator.hpp
@@ -35,7 +35,8 @@ class CpuAllocator : public MemoryAllocator {
  public:
   explicit CpuAllocator(size_t block_size, size_t max_allocated = -1);
 
-  [[nodiscard]] void* get(size_t size) override;
+  [[nodiscard]] BufferPtr get(const InferenceRequestInput& tensor,
+                              size_t batch_size) override;
   void put(const void* address) override;
 
   // void free(const void* address);
@@ -47,7 +48,7 @@ class CpuAllocator : public MemoryAllocator {
   size_t block_id_ = 0;
   std::mutex mutex_;
   std::list<MemoryHeader> headers_;
-  std::vector<std::vector<std::byte>> data_;
+  std::list<std::vector<std::byte>> data_;
 };
 
 }  // namespace amdinfer
diff --git a/src/amdinfer/core/memory_pool/memory_allocator.hpp b/src/amdinfer/core/memory_pool/memory_allocator.hpp
index df3857f2f..089f71863 100644
--- a/src/amdinfer/core/memory_pool/memory_allocator.hpp
+++ b/src/amdinfer/core/memory_pool/memory_allocator.hpp
@@ -22,8 +22,12 @@
 
 #include <cstddef>  // for size_t
 
+#include "amdinfer/declarations.hpp"
+
 namespace amdinfer {
 
+enum class MemoryAllocators { Cpu, VartTensor };
+
 struct MemoryHeader {
   std::byte* address;
   bool free;
@@ -39,7 +43,8 @@ class MemoryAllocator {
   virtual ~MemoryAllocator() = default;
 
   // these methods are thread-safe
-  [[nodiscard]] virtual void* get(size_t size) = 0;
+  [[nodiscard]] virtual BufferPtr get(const InferenceRequestInput& tensor,
+                                      size_t batch_size) = 0;
   virtual void put(const void* address) = 0;
 };
 
diff --git a/src/amdinfer/core/memory_pool/pool.cpp b/src/amdinfer/core/memory_pool/pool.cpp
index f07ef14b0..c7f0f2810 100644
--- a/src/amdinfer/core/memory_pool/pool.cpp
+++ b/src/amdinfer/core/memory_pool/pool.cpp
@@ -22,6 +22,7 @@
 #include "amdinfer/buffers/cpu.hpp"
 #include "amdinfer/core/exceptions.hpp"
 #include "amdinfer/core/memory_pool/cpu_allocator.hpp"
+#include "amdinfer/core/memory_pool/vart_tensor_allocator.hpp"
 
 namespace amdinfer {
 
@@ -30,28 +31,28 @@ const size_t kDefaultCpuBlockSize = 1'048'576;  // arbitrarily 1MiB
 MemoryPool::MemoryPool() {
   allocators_.try_emplace(MemoryAllocators::Cpu,
                           std::make_unique<CpuAllocator>(kDefaultCpuBlockSize));
+#ifdef AMDINFER_ENABLE_VITIS
+  allocators_.try_emplace(MemoryAllocators::VartTensor,
+                          std::make_unique<VartTensorAllocator>());
+#endif
 }
 
 std::unique_ptr<Buffer> MemoryPool::get(
-  const std::vector<MemoryAllocators>& allocators, size_t size) const {
+  const std::vector<MemoryAllocators>& allocators,
+  const InferenceRequestInput& tensor, size_t batch_size) const {
   for (const auto& allocator : allocators) {
-    void* address = nullptr;
     try {
-      address = allocators_.at(allocator)->get(size);
+      return allocators_.at(allocator)->get(tensor, batch_size);
     } catch (const runtime_error&) {
       continue;
     }
-    if (address != nullptr) {
-      return std::make_unique<CpuBuffer>(address, allocator);
-    }
   }
   throw runtime_error("Memory could not be allocated");
 }
 
 void MemoryPool::put(std::unique_ptr<Buffer> memory) const {
-  auto* buffer = dynamic_cast<CpuBuffer*>(memory.get());
-  const auto allocator = buffer->getAllocator();
-  const auto* address = buffer->data(0);
+  const auto allocator = memory->getAllocator();
+  const auto* address = memory->data(0);
   allocators_.at(allocator)->put(address);
 }
 
diff --git a/src/amdinfer/core/memory_pool/pool.hpp b/src/amdinfer/core/memory_pool/pool.hpp
index 922ed588b..d7395608e 100644
--- a/src/amdinfer/core/memory_pool/pool.hpp
+++ b/src/amdinfer/core/memory_pool/pool.hpp
@@ -25,22 +25,19 @@
 #include <unordered_map>
 #include <vector>
 
-#include "amdinfer/buffers/buffer.hpp"
 #include "amdinfer/build_options.hpp"
 #include "amdinfer/core/memory_pool/memory_allocator.hpp"
+#include "amdinfer/declarations.hpp"
 
 namespace amdinfer {
 
-enum class MemoryAllocators { Cpu };
-
-using Memory = std::pair<MemoryAllocators, void*>;
-
 class MemoryPool {
  public:
   MemoryPool();
 
   std::unique_ptr<Buffer> get(const std::vector<MemoryAllocators>& allocators,
-                              size_t size) const;
+                              const InferenceRequestInput& tensor,
+                              size_t batch_size) const;
   void put(std::unique_ptr<Buffer> memory) const;
 
  private:
diff --git a/src/amdinfer/core/memory_pool/vart_tensor_allocator.cpp b/src/amdinfer/core/memory_pool/vart_tensor_allocator.cpp
new file mode 100644
index 000000000..de8db7265
--- /dev/null
+++ b/src/amdinfer/core/memory_pool/vart_tensor_allocator.cpp
@@ -0,0 +1,97 @@
+// Copyright 2023 Advanced Micro Devices, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/**
+ * @file
+ * @brief
+ */
+
+#include "amdinfer/core/memory_pool/vart_tensor_allocator.hpp"
+
+#include <cassert>
+#include <iostream>
+#include <vector>
+#include <xir/util/data_type.hpp>  // for DataType
+
+#include "amdinfer/buffers/vart_tensor.hpp"
+#include "amdinfer/core/data_types_internal.hpp"
+#include "amdinfer/core/exceptions.hpp"
+#include "amdinfer/core/predict_api.hpp"
+
+namespace amdinfer {
+
+VartTensorAllocator::VartTensorAllocator(size_t max_allocate)
+  : max_allocate_(max_allocate) {}
+
+BufferPtr VartTensorAllocator::get(const InferenceRequestInput& tensor,
+                                   size_t batch_size) {
+  const auto& name = tensor.getName();
+  const auto& shape = tensor.getShape();
+  const auto datatype = tensor.getDatatype();
+
+  const std::lock_guard lock{mutex_};
+  auto found = headers_.end();
+  const auto end = headers_.end();
+  for (auto it = headers_.begin(); it != end; it++) {
+    if (it->free && it->name == name && it->shape == shape &&
+        it->datatype == datatype && it->batch_size == batch_size) {
+      found = it;
+      break;
+    }
+  }
+
+  if (found != end) {
+    found->free = false;
+    return std::make_unique<VartTensorBuffer>(found->address,
+                                              MemoryAllocators::VartTensor);
+  }
+
+  auto size_to_allocate = tensor.getSize() * datatype.size() * batch_size;
+  if (allocated_ + size_to_allocate > max_allocate_) {
+    throw runtime_error("Too much requested");
+  }
+
+  auto xir_type = mapTypeToXir(datatype);
+  std::vector<int> xir_shape{shape.begin(), shape.end()};
+  tensors_.emplace_back(xir::Tensor::create(name, xir_shape, xir_type));
+  buffers_.emplace_back(tensors_.back().get());
+
+  allocated_ += size_to_allocate;
+
+  auto* retval = reinterpret_cast<std::byte*>(&(buffers_.back()));
+
+  headers_.emplace_back(retval, false, name, shape, datatype, batch_size);
+
+  return std::make_unique<VartTensorBuffer>(retval,
+                                            MemoryAllocators::VartTensor);
+}
+
+void VartTensorAllocator::put(const void* address) {
+  const std::lock_guard lock{mutex_};
+  const auto end = headers_.end();
+  auto found = headers_.end();
+  for (auto it = headers_.begin(); it != end; it++) {
+    if (it->address == address) {
+      found->free = true;
+      // std::cout << "Freed memory\n";
+      return;
+    }
+  }
+
+  if (found == end) {
+    throw runtime_error("Address not found");
+  }
+}
+
+}  // namespace amdinfer
diff --git a/src/amdinfer/core/memory_pool/vart_tensor_allocator.hpp b/src/amdinfer/core/memory_pool/vart_tensor_allocator.hpp
new file mode 100644
index 000000000..598ce06c8
--- /dev/null
+++ b/src/amdinfer/core/memory_pool/vart_tensor_allocator.hpp
@@ -0,0 +1,83 @@
+// Copyright 2023 Advanced Micro Devices, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/**
+ * @file
+ * @brief
+ */
+
+#ifndef GUARD_AMDINFER_CORE_MEMORY_POOL_VART_TENSOR_ALLOCATOR
+#define GUARD_AMDINFER_CORE_MEMORY_POOL_VART_TENSOR_ALLOCATOR
+
+#include "amdinfer/build_options.hpp"
+
+#ifdef AMDINFER_ENABLE_VITIS
+
+#include <cstddef>
+#include <list>
+#include <mutex>
+#include <vart/experimental/runner_helper.hpp>  // for CpuFlatTensorBufferOwned
+#include <vector>
+#include <xir/tensor/tensor.hpp>  // for Tensor
+
+#include "amdinfer/core/memory_pool/memory_allocator.hpp"
+#include "amdinfer/core/predict_api.hpp"
+
+namespace amdinfer {
+
+struct VartTensorHeader {
+  std::byte* address;
+  bool free;
+
+  std::string name;
+  std::vector<uint64_t> shape;
+  DataType datatype;
+  size_t batch_size;
+
+  VartTensorHeader(std::byte* address, bool free, const std::string& name,
+                   const std::vector<uint64_t>& shape, DataType datatype,
+                   size_t batch_size)
+    : address(address),
+      free(free),
+      name(name),
+      shape(shape),
+      datatype(datatype),
+      batch_size(batch_size) {}
+};
+
+class VartTensorAllocator : public MemoryAllocator {
+ public:
+  explicit VartTensorAllocator(size_t max_allocated = -1);
+
+  [[nodiscard]] BufferPtr get(const InferenceRequestInput& tensor,
+                              size_t batch_size) override;
+  void put(const void* address) override;
+
+  // void free(const void* address);
+
+ private:
+  size_t allocated_ = 0;
+  size_t max_allocate_;
+  std::mutex mutex_;
+
+  std::list<VartTensorHeader> headers_;
+  std::list<std::unique_ptr<xir::Tensor>> tensors_;
+  std::list<vart::CpuFlatTensorBufferOwned> buffers_;
+};
+
+}  // namespace amdinfer
+
+#endif
+
+#endif  // GUARD_AMDINFER_CORE_MEMORY_POOL_VART_TENSOR_ALLOCATOR
diff --git a/src/amdinfer/servers/grpc_server.cpp b/src/amdinfer/servers/grpc_server.cpp
index cfdbeab41..a4889989c 100644
--- a/src/amdinfer/servers/grpc_server.cpp
+++ b/src/amdinfer/servers/grpc_server.cpp
@@ -317,8 +317,7 @@ InferenceRequestInput getInput(
   input.setParameters(mapProtoToParameters(req.parameters()));
 
   auto size = input.getSize();
-  auto buffer =
-    pool->get({MemoryAllocators::Cpu}, size * input.getDatatype().size());
+  auto buffer = pool->get({MemoryAllocators::Cpu}, input, 1);
   input.setData(buffer->data(0));
   // auto* dest = static_cast<std::byte*>(input_buffer->data(offset));
   AMDINFER_LOG_TRACE(observer.logger, "Writing " + std::to_string(size) +
diff --git a/src/amdinfer/servers/http_server.cpp b/src/amdinfer/servers/http_server.cpp
index eef0ac05f..f7ad48b81 100644
--- a/src/amdinfer/servers/http_server.cpp
+++ b/src/amdinfer/servers/http_server.cpp
@@ -32,6 +32,7 @@
 #include <utility>        // for move
 #include <vector>         // for vector
 
+#include "amdinfer/buffers/buffer.hpp"         // for BufferPtr
 #include "amdinfer/build_options.hpp"          // for AMDINFER_ENABLE_TRACING
 #include "amdinfer/clients/http_internal.hpp"  // for propagate, errorHtt...
 #include "amdinfer/core/exceptions.hpp"        // for runtime_error, inva...
@@ -367,8 +368,7 @@ InferenceRequestInput getInput(const Json::Value &json,
     input.setParameters(std::make_unique<ParameterMap>());
   }
 
-  auto size = input.getDatatype().size() * util::containerProduct(shape_vector);
-  auto buffer = pool->get({MemoryAllocators::Cpu}, size);
+  auto buffer = pool->get({MemoryAllocators::Cpu}, input, 1);
   if (!json.isMember("data")) {
     throw invalid_argument("No 'data' key present in request input");
   }
diff --git a/src/amdinfer/workers/xmodel.cpp b/src/amdinfer/workers/xmodel.cpp
index c5c6efc95..d858d1e25 100644
--- a/src/amdinfer/workers/xmodel.cpp
+++ b/src/amdinfer/workers/xmodel.cpp
@@ -44,6 +44,7 @@
 
 #include "amdinfer/batching/batcher.hpp"          // for BatchPtr, Batch
 #include "amdinfer/buffers/buffer.hpp"            // for Buffer
+#include "amdinfer/buffers/vart_tensor.hpp"       // for VartTensorBuffer
 #include "amdinfer/build_options.hpp"             // for AMDINFER_ENABLE_ME...
 #include "amdinfer/core/data_types.hpp"           // for DataType
 #include "amdinfer/core/data_types_internal.hpp"  // for mapXirToType
@@ -97,7 +98,7 @@ std::thread XModel::spawn(BatchPtrQueue* input_queue) {
 }
 
 std::vector<MemoryAllocators> XModel::getAllocators() const {
-  return {MemoryAllocators::Cpu};
+  return {MemoryAllocators::VartTensor};
 }
 
 vart::RunnerExt* XModel::getRunner() {
@@ -212,8 +213,6 @@ void XModel::doRun(BatchPtrQueue* input_queue) {
         }
 #endif
 
-        std::queue<std::pair<uint32_t, int>> futures;
-
         auto input_buffers = batch->getInputBuffers();
         // const auto& output_buffers = batch->getOutputBuffers();
 
@@ -221,27 +220,41 @@ void XModel::doRun(BatchPtrQueue* input_queue) {
           logTraceBuffer(getLogger(), buffer->data(0));
         }
 
-        auto inputs_ptr = getRunner()->get_inputs();
-        if (inputs_ptr.size() != input_buffers.size()) {
-          throw invalid_argument("Input size does not match model");
+        std::vector<vart::TensorBuffer*> inputs_ptr;
+        for (const auto& buffer : input_buffers) {
+          auto* vart = dynamic_cast<VartTensorBuffer*>(buffer.get());
+          inputs_ptr.emplace_back(vart->getTensorBuffer());
         }
-        for (auto j = 0U; j < input_buffers.size(); ++j) {
-          auto* tensor = inputs_ptr.at(j);
-          auto& buffer = input_buffers.at(j);
 
-          const auto [addr, size] = tensor->data();
-          auto* tensor_addr = reinterpret_cast<void*>(addr);
+        BufferPtrs output_buffers;
+        auto output_tensors = this->getRunner()->get_output_tensors();
+        for (const auto* tensor : output_tensors) {
+          auto xir_shape = tensor->get_shape();
+          std::vector<size_t> shape{xir_shape.begin(), xir_shape.end()};
+          auto xir_type = tensor->get_data_type();
+          auto type = mapXirToType(xir_type);
+          InferenceRequestInput input(nullptr, shape, type, tensor->get_name());
+          // the shape includes the batch size so use external batch size 1
+          output_buffers.push_back(
+            pool_->get({MemoryAllocators::VartTensor}, input, 1));
+        }
 
-          std::memcpy(tensor_addr, buffer->data(0), size);
+        std::vector<vart::TensorBuffer*> outputs_ptr;
+        for (const auto& buffer : output_buffers) {
+          auto* vart = dynamic_cast<VartTensorBuffer*>(buffer.get());
+          outputs_ptr.emplace_back(vart->getTensorBuffer());
         }
-        auto outputs_ptr = getRunner()->get_outputs();
+
+        // auto outputs_ptr = this->getRunner()->get_outputs();
 
         for (auto* input : inputs_ptr) {
           const auto* tensor = input->get_tensor();
-          input->sync_for_write(
-            0, tensor->get_element_num() / (tensor->get_shape())[0]);
+          auto num = tensor->get_element_num();
+          auto batches = (tensor->get_shape())[0];
+          input->sync_for_write(0, num / batches);
         }
 
+        std::queue<std::pair<uint32_t, int>> futures;
         futures.push(getRunner()->execute_async(inputs_ptr, outputs_ptr));
 
         std::vector<InferenceResponse> responses;


Test patch:
diff --git a/tests/performance/batching/CMakeLists.txt b/tests/performance/batching/CMakeLists.txt
index 071685fa8..53fe7cb9b 100644
--- a/tests/performance/batching/CMakeLists.txt
+++ b/tests/performance/batching/CMakeLists.txt
@@ -20,8 +20,8 @@ list(
     tests_libs
     "$<TARGET_OBJECTS:batcher>~fake_observation\
       ~$<TARGET_OBJECTS:fake_worker_info_buffers_finite>~predict_api~parameters\
-      ~data_types\~$<TARGET_OBJECTS:softBatcher>~cpuBuffer~buffer~memory_pool\
-      ~batch~timer"
+      ~data_types\~batching~buffer~memory_pool\
+      ~data_types_internal~batch~timer"
 )
 
 amdinfer_add_benchmarks("${tests}" "${tests_libs}")
diff --git a/tests/performance/batching/benchmark_soft.cpp b/tests/performance/batching/benchmark_soft.cpp
index a2b86d600..51991474f 100644
--- a/tests/performance/batching/benchmark_soft.cpp
+++ b/tests/performance/batching/benchmark_soft.cpp
@@ -31,6 +31,7 @@
 #include <vector>    // for vector
 
 #include "amdinfer/batching/soft.hpp"          // for BatchPtr, SoftBatcher
+#include "amdinfer/buffers/buffer.hpp"         // for BufferPtr
 #include "amdinfer/core/data_types.hpp"        // for DataType, DataType::U...
 #include "amdinfer/core/memory_pool/pool.hpp"  // for MemoryPool
 #include "amdinfer/core/parameters.hpp"        // for ParameterMap
diff --git a/tests/unit/batching/CMakeLists.txt b/tests/unit/batching/CMakeLists.txt
index 5e3c2a12d..ae4c4a157 100644
--- a/tests/unit/batching/CMakeLists.txt
+++ b/tests/unit/batching/CMakeLists.txt
@@ -20,11 +20,12 @@ list(
     tests_libs
     "$<TARGET_OBJECTS:batcher>~fake_observation~\
       $<TARGET_OBJECTS:fake_worker_info_buffers_infinite>~predict_api~\
-      parameters~data_types~$<TARGET_OBJECTS:softBatcher>~timer~memory_pool~batch~cpuBuffer~buffer"
+      parameters~data_types~batching~timer~memory_pool~\
+      batch~buffer~data_types_internal"
     "$<TARGET_OBJECTS:batcher>~fake_observation~\
       $<TARGET_OBJECTS:fake_worker_info_buffers_finite>~predict_api~data_types~\
-      parameters~$<TARGET_OBJECTS:softBatcher>~cpuBuffer~buffer~\
-      timer~memory_pool~batch"
+      parameters~batching~buffer~\
+      timer~memory_pool~batch~data_types_internal"
 )
 
 amdinfer_add_unit_tests("${tests}" "${tests_libs}")
diff --git a/tests/unit/batching/test_soft_batching.cpp b/tests/unit/batching/test_soft_batching.cpp
index 58c7ef008..efd2e7974 100644
--- a/tests/unit/batching/test_soft_batching.cpp
+++ b/tests/unit/batching/test_soft_batching.cpp
@@ -111,7 +111,9 @@ class UnitSoftBatcherFixture : public testing::TestWithParam<BatchConfig> {
 
     this->batcher_->start({MemoryAllocators::Cpu});
 
-    buffer_ = pool_.get({MemoryAllocators::Cpu}, data_size);
+    InferenceRequestInput input{nullptr, data_shape_, DataType::Uint8};
+
+    buffer_ = pool_.get({MemoryAllocators::Cpu}, input, data_size);
 
     auto* data = static_cast<uint8_t*>(buffer_->data(0));
 
diff --git a/tests/unit/core/memory_pool/CMakeLists.txt b/tests/unit/core/memory_pool/CMakeLists.txt
index affca5303..f88b7f7ed 100644
--- a/tests/unit/core/memory_pool/CMakeLists.txt
+++ b/tests/unit/core/memory_pool/CMakeLists.txt
@@ -14,6 +14,12 @@
 
 list(APPEND tests cpu_allocator pool)
 
-list(APPEND tests_libs "cpu_allocator" "memory_pool~buffer~cpuBuffer")
+list(
+  APPEND tests_libs
+         "cpu_allocator~buffers~cpuBuffer~predict_api~data_types~parameters~\
+           data_types_internal"
+         "memory_pool~buffers~cpuBuffer~predict_api~data_types~parameters~\
+           data_types_internal"
+)
 
 amdinfer_add_unit_tests("${tests}" "${tests_libs}")
diff --git a/tests/unit/core/memory_pool/test_cpu_allocator.cpp b/tests/unit/core/memory_pool/test_cpu_allocator.cpp
index 2f3d033b1..7ae755dee 100644
--- a/tests/unit/core/memory_pool/test_cpu_allocator.cpp
+++ b/tests/unit/core/memory_pool/test_cpu_allocator.cpp
@@ -22,9 +22,11 @@
 // #include <string>   // for string, basic_string, alloc...
 // #include <vector>   // for vector
 
+#include "amdinfer/buffers/buffer.hpp"  // for BufferPtr
 #include "amdinfer/core/exceptions.hpp"
 #include "amdinfer/core/memory_pool/cpu_allocator.hpp"  // for CpuSimpl...
-#include "amdinfer/testing/gtest.hpp"  // for AssertionResult,...
+#include "amdinfer/core/predict_api.hpp"  // for InferenceRequestInput
+#include "amdinfer/testing/gtest.hpp"     // for AssertionResult,...
 
 namespace amdinfer {
 
@@ -32,53 +34,66 @@ namespace amdinfer {
 TEST(UnitCpuAllocator, Basic) {
   CpuAllocator allocator{sizeof(int)};
 
-  const auto* foo = static_cast<int*>(allocator.get(sizeof(int)));
+  InferenceRequestInput input{nullptr, {1}, DataType::Int32};
 
-  const auto* bar = static_cast<int*>(allocator.get(sizeof(int)));
-  ASSERT_NE(foo, bar);
+  const auto buffer_0 = allocator.get(input, 1);
+  const auto buffer_1 = allocator.get(input, 1);
 
-  allocator.put(foo);
-  allocator.put(bar);
+  const auto* address_0 = static_cast<int*>(buffer_0->data(0));
+  const auto* address_1 = static_cast<int*>(buffer_1->data(0));
+  ASSERT_NE(address_0, address_1);
+
+  allocator.put(address_0);
+  allocator.put(address_1);
 }
 
 // NOLINTNEXTLINE(cert-err58-cpp, cppcoreguidelines-owning-memory)
 TEST(UnitCpuAllocator, Correctness) {
   CpuAllocator allocator{sizeof(int) * 4, sizeof(int) * 4};
+  InferenceRequestInput input{nullptr, {1}, DataType::Int32};
+
+  const auto buffer_0 = allocator.get(input, 1);
+  const auto buffer_1 = allocator.get(input, 1);
 
-  const auto* foo = static_cast<int*>(allocator.get(sizeof(int)));
-  const auto* bar = static_cast<int*>(allocator.get(sizeof(int)));
+  const auto* address_0 = static_cast<int*>(buffer_0->data(0));
+  const auto* address_1 = static_cast<int*>(buffer_1->data(0));
 
-  ASSERT_EQ(foo + 1, bar);
+  ASSERT_EQ(address_0 + 1, address_1);
 
-  allocator.put(foo);
-  const auto* baz = static_cast<int*>(allocator.get(sizeof(int)));
+  allocator.put(address_0);
+  const auto buffer_2 = allocator.get(input, 1);
+  const auto* address_2 = static_cast<int*>(buffer_2->data(0));
 
-  ASSERT_EQ(foo, baz);
+  ASSERT_EQ(address_0, address_2);
 
-  allocator.put(baz);
-  allocator.put(bar);
+  allocator.put(address_2);
+  allocator.put(address_1);
 
-  const auto* bard = static_cast<int*>(allocator.get(sizeof(int) * 4));
-  ASSERT_EQ(foo, bard);
+  const auto buffer_3 = allocator.get(input, 4);
+  const auto* address_3 = static_cast<int*>(buffer_3->data(0));
+  ASSERT_EQ(address_0, address_3);
 }
 
 // NOLINTNEXTLINE(cert-err58-cpp, cppcoreguidelines-owning-memory)
 TEST(UnitCpuAllocator, ExceedingMax) {
   CpuAllocator allocator{sizeof(int), sizeof(int)};
+  InferenceRequestInput input{nullptr, {1}, DataType::Int32};
 
-  std::ignore = static_cast<int*>(allocator.get(sizeof(int)));
+  std::ignore = allocator.get(input, 1);
   // NOLINTNEXTLINE(cppcoreguidelines-avoid-goto, hicpp-avoid-goto)
-  EXPECT_THROW_CHECK(
-    std::ignore = static_cast<int*>(allocator.get(sizeof(int)));
-    , EXPECT_STREQ(e.what(), "Too much requested");, runtime_error);
+  EXPECT_THROW_CHECK(std::ignore = allocator.get(input, 1);
+                     , EXPECT_STREQ(e.what(), "Too much requested");
+                     , runtime_error);
 }
 
 // NOLINTNEXTLINE(cert-err58-cpp, cppcoreguidelines-owning-memory)
 TEST(UnitCpuAllocator, BadFree) {
   CpuAllocator allocator{sizeof(int)};
+  InferenceRequestInput input{nullptr, {1}, DataType::Int32};
 
-  const auto* foo = static_cast<int*>(allocator.get(sizeof(int)));
-  const auto* bad_address = foo + 1;
+  const auto buffer_0 = allocator.get(input, 1);
+  const auto* address_0 = static_cast<int*>(buffer_0->data(0));
+  const auto* bad_address = address_0 + 1;
   // NOLINTNEXTLINE(cppcoreguidelines-avoid-goto, hicpp-avoid-goto)
   EXPECT_THROW_CHECK(allocator.put(bad_address);
                      , EXPECT_STREQ(e.what(), "Address not found");
diff --git a/tests/unit/core/memory_pool/test_pool.cpp b/tests/unit/core/memory_pool/test_pool.cpp
index 3fd33a374..b045aa09b 100644
--- a/tests/unit/core/memory_pool/test_pool.cpp
+++ b/tests/unit/core/memory_pool/test_pool.cpp
@@ -19,17 +19,20 @@
 
 #include <tuple>
 
+#include "amdinfer/buffers/buffer.hpp"  // for BufferPtr
 #include "amdinfer/core/exceptions.hpp"
 #include "amdinfer/core/memory_pool/pool.hpp"
-#include "amdinfer/testing/gtest.hpp"  // for AssertionResult,...
+#include "amdinfer/core/predict_api.hpp"  // for InferenceRequestInput
+#include "amdinfer/testing/gtest.hpp"     // for AssertionResult,...
 
 namespace amdinfer {
 
 // NOLINTNEXTLINE(cert-err58-cpp, cppcoreguidelines-owning-memory)
 TEST(UnitPool, Basic) {
   MemoryPool pool;
+  InferenceRequestInput input{nullptr, {1}, DataType::Int32};
 
-  auto buffer = pool.get({MemoryAllocators::Cpu}, sizeof(int));
+  auto buffer = pool.get({MemoryAllocators::Cpu}, input, 1);
 
   pool.put(std::move(buffer));
 }