context.h

Go to the documentation of this file.

// Copyright 2018 Delft University of Technology
//
// 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.
 
#pragma once
 
#include <arrow/api.h>
#include <fletcher/common.h>
#include <utility>
#include <vector>
#include <memory>
#include <iostream>
 
#include "fletcher/platform.h"
#include "fletcher/status.h"
 
namespace fletcher {
 
using fletcher::Mode;
 
enum class MemType {
      ANY,
 
      CACHE
};
 
struct DeviceBuffer {
  const uint8_t *host_address = nullptr;
  da_t device_address = D_NULLPTR;
  int64_t size = 0;
 
  MemType memory = MemType::CACHE;
  Mode mode = Mode::READ;
 
  bool available_to_device = false;
  bool was_alloced = false;
 
  DeviceBuffer() = default;
 
  DeviceBuffer(const uint8_t *host_address, int64_t size, MemType type, Mode access_mode)
      : host_address(host_address), size(size), memory(type), mode(access_mode) {}
};
 
class Context {
 public:
  explicit Context(std::shared_ptr<Platform> platform) : platform_(std::move(platform)) {}
 
  ~Context();
 
  static Status Make(std::shared_ptr<Context> *context, const std::shared_ptr<Platform> &platform);
 
  Status QueueRecordBatch(const std::shared_ptr<arrow::RecordBatch> &record_batch,
                          MemType mem_type = MemType::ANY);
 
  size_t GetQueueSize() const;
 
  Status Enable();
 
  std::shared_ptr<Platform> platform() const { return platform_; }
 
  uint64_t num_buffers() const;
 
  DeviceBuffer device_buffer(size_t i) const { return device_buffers_[i]; }
 
  uint64_t num_recordbatches() const { return host_batches_.size(); }
 
  std::shared_ptr<arrow::RecordBatch> recordbatch(size_t i) const { return host_batches_[i]; }
 
 protected:
  std::shared_ptr<Platform> platform_;
  std::vector<std::shared_ptr<arrow::RecordBatch>> host_batches_;
  std::vector<RecordBatchDescription> host_batch_desc_;
  std::vector<MemType> host_batch_memtype_;
  std::vector<DeviceBuffer> device_buffers_;
};
 
}  // namespace fletcher