kernel.h

/* Copyright 2016 Carnegie Mellon University
 *
 * 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 "scanner/api/frame.h"
#include "scanner/util/common.h"
#include "scanner/util/memory.h"
#include "scanner/util/profiler.h"

#include <vector>

namespace scanner {

struct Element {
  Element() = default;
  Element(const Element&) = default;
  Element(Element&&) = default;
  Element& operator=(const Element&) = default;

  Element(u8* buffer, size_t size);
  Element(Frame* frame);

  inline Frame* as_frame() { return reinterpret_cast<Frame*>(buffer); }
  inline const Frame* as_const_frame() const {
    return reinterpret_cast<Frame*>(buffer);
  }
  inline FrameInfo* as_frame_info() {
    return reinterpret_cast<FrameInfo*>(buffer);
  }
  inline const FrameInfo* as_const_frame_info() const {
    return reinterpret_cast<FrameInfo*>(buffer);
  }

  inline bool is_null() const {
    return buffer == nullptr;
  }

  u8* buffer;
  size_t size;
  bool is_frame;
  // @brief the index of the element in the input domain
  i64 index;
};

using Elements = std::vector<Element>;

using BatchedElements = std::vector<Elements>;

using StenciledElements = std::vector<Elements>;

using StenciledBatchedElements = std::vector<std::vector<Elements>>;

inline size_t num_rows(const Elements& column) { return column.size(); }

inline void insert_element(Elements& column, u8* buffer, size_t size) {
  column.push_back(::scanner::Element{buffer, size});
}

inline void insert_frame(Elements& column, Frame* frame) {
  column.push_back(::scanner::Element{frame});
}

inline void insert_element(Element& element, u8* buffer, size_t size) {
  element = ::scanner::Element{buffer, size};
}

inline void insert_frame(Element& element, Frame* frame) {
  element = ::scanner::Element{frame};
}

inline Element add_element_ref(DeviceHandle device, Element& element) {
  if (element.is_null()) {
    return Element();
  }
  Element ele;
  if (element.is_frame) {
    Frame* frame = element.as_frame();
    add_buffer_ref(device, frame->data);
    // Copy frame because Frame is not referenced counted
    ele = ::scanner::Element{new Frame(frame->as_frame_info(), frame->data)};
  } else {
    add_buffer_ref(device, element.buffer);
    ele = element;
  }
  ele.index = element.index;
  return ele;
}

inline void delete_element(DeviceHandle device, Element& element) {
  if (element.is_null()) {
    return;
  }
  if (element.is_frame) {
    Frame* frame = element.as_frame();
    delete_buffer(device, frame->data);
    delete frame;
  } else {
    delete_buffer(device, element.buffer);
  }
}

struct KernelConfig {
  std::vector<DeviceHandle> devices;  
  std::vector<std::string> input_columns;
  std::vector<proto::ColumnType> input_column_types;
  std::vector<std::string> output_columns;
  std::vector<proto::ColumnType> output_column_types;
  std::vector<u8> args;  
  i32 node_id;

  static KernelConfig dummy() {
    KernelConfig config;
    config.devices.push_back(CPU_DEVICE);
    return config;
  }
};

class BaseKernel {
 public:
  static const i32 UnlimitedDevices = 0;
  BaseKernel(const KernelConfig& config);

  virtual ~BaseKernel(){};

  virtual void validate(proto::Result* result) { result->set_success(true); }

  virtual void fetch_resources(proto::Result* result) { result->set_success(true); }

  virtual void setup_with_resources(proto::Result* result) { result->set_success(true); }

  virtual void new_stream(const std::vector<u8>& args) {};

  virtual void reset(){};

  virtual void execute_kernel(const StenciledBatchedElements& input_columns,
                              BatchedElements& output_columns) = 0;

  virtual void set_profiler(Profiler* profiler) { profiler_ = profiler; }

  Profiler* profiler_ = nullptr;
};


class StenciledBatchedKernel : public BaseKernel {
 public:
  static const i32 UnlimitedDevices = 0;
  StenciledBatchedKernel(const KernelConfig& config);

  virtual ~StenciledBatchedKernel(){};

  virtual void execute_kernel(const StenciledBatchedElements& input_columns,
                              BatchedElements& output_columns) override;

 protected:
  virtual void execute(const StenciledBatchedElements& input_columns,
                       BatchedElements& output_columns) = 0;

  Profiler* profiler_ = nullptr;
};

class BatchedKernel : public BaseKernel {
 public:
  BatchedKernel(const KernelConfig& config);

  virtual ~BatchedKernel(){};

  virtual void execute_kernel(const StenciledBatchedElements& input_columns,
                              BatchedElements& output_columns);
 protected:
  virtual void execute(const BatchedElements& input_columns,
                       BatchedElements& output_columns) = 0;
};

class StenciledKernel : public BaseKernel {
 public:
  StenciledKernel(const KernelConfig& config);

  virtual ~StenciledKernel(){};

  virtual void execute_kernel(const StenciledBatchedElements& input_columns,
                              BatchedElements& output_columns);
 protected:
  virtual void execute(const StenciledElements& input_columns,
                       Elements& output_columns) = 0;
};

class Kernel : public BaseKernel {
 public:
  Kernel(const KernelConfig& config);

  virtual ~Kernel(){};

  virtual void execute_kernel(const StenciledBatchedElements& input_columns,
                              BatchedElements& output_columns);
 protected:
  virtual void execute(const Elements& input_columns,
                       Elements& output_columns) = 0;
};

class VideoKernel {
 protected:
  void check_frame(const DeviceHandle& device, const Element& element);

  void check_frame_info(const DeviceHandle& device, const Element& element);

  virtual void new_frame_info(){};

  FrameInfo frame_info_{};
};

namespace internal {

class KernelBuilder;

using KernelConstructor =
    std::function<BaseKernel*(const KernelConfig& config)>;

class KernelRegistration {
 public:
  KernelRegistration(const KernelBuilder& builder);
};

class KernelBuilder {
 public:
  friend class KernelRegistration;

  KernelBuilder(const std::string& name, KernelConstructor constructor)
    : name_(name),
      constructor_(constructor),
      device_type_(DeviceType::CPU),
      num_devices_(1),
      can_batch_(false),
      preferred_batch_size_(1) {}

  KernelBuilder& device(DeviceType device_type) {
    device_type_ = device_type;
    return *this;
  }

  KernelBuilder& num_devices(i32 devices) {
    num_devices_ = devices;
    return *this;
  }

  KernelBuilder& input_device(const std::string& input_name,
                              DeviceType device_type) {
    input_devices_[input_name] = device_type;
    return *this;
  }

  KernelBuilder& output_device(const std::string& output_name,
                               DeviceType device_type) {
    output_devices_[output_name] = device_type;
    return *this;
  }

  KernelBuilder& batch(i32 preferred_batch_size = 1) {
    can_batch_ = true;
    preferred_batch_size_ = preferred_batch_size;
    return *this;
  }

 private:
  std::string name_;
  KernelConstructor constructor_;
  DeviceType device_type_;
  i32 num_devices_;
  std::map<std::string, DeviceType> input_devices_;
  std::map<std::string, DeviceType> output_devices_;
  bool can_batch_;
  i32 preferred_batch_size_;
};
}

#define REGISTER_KERNEL(name__, kernel__) \
  REGISTER_KERNEL_HELPER(__COUNTER__, name__, kernel__)

#define REGISTER_KERNEL_HELPER(uid__, name__, kernel__) \
  REGISTER_KERNEL_UID(uid__, name__, kernel__)

#define REGISTER_KERNEL_UID(uid__, name__, kernel__)                         \
  static ::scanner::internal::KernelRegistration kernel_registration_##uid__ \
      __attribute__((unused)) = ::scanner::internal::KernelBuilder(          \
          #name__, [](const ::scanner::KernelConfig& config) {             \
            return new kernel__(config);                                     \
          })
}