text: /*
   * Copyright 2019 OmniSci, 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.
   */
  
  #include "QueryEngine/TableFunctions/TableFunctionExecutionContext.h"
  
  #include "Analyzer/Analyzer.h"
  #include "Logger/Logger.h"
  #include "QueryEngine/ColumnFetcher.h"
  #include "QueryEngine/GpuMemUtils.h"
  #include "QueryEngine/TableFunctions/TableFunctionCompilationContext.h"
  
  namespace {
  
  template <typename T>
  const int8_t* create_literal_buffer(T literal,
                                      const ExecutorDeviceType device_type,
                                      std::vector<std::unique_ptr<char[]>>& literals_owner,
                                      CudaAllocator* gpu_allocator) {
    CHECK_LE(sizeof(T), sizeof(int64_t));  // pad to 8 bytes
    switch (device_type) {
      case ExecutorDeviceType::CPU: {
        literals_owner.emplace_back(std::make_unique<char[]>(sizeof(int64_t)));
        std::memcpy(literals_owner.back().get(), &literal, sizeof(T));
        return reinterpret_cast<const int8_t*>(literals_owner.back().get());
      }
      case ExecutorDeviceType::GPU: {
        CHECK(gpu_allocator);
        const auto gpu_literal_buf_ptr = gpu_allocator->alloc(sizeof(int64_t));
        gpu_allocator->copyToDevice(
            gpu_literal_buf_ptr, reinterpret_cast<int8_t*>(&literal), sizeof(T));
        return gpu_literal_buf_ptr;
      }
    }
    UNREACHABLE();
    return nullptr;
  }
  
  size_t get_output_row_count(const TableFunctionExecutionUnit& exe_unit,
                              size_t input_element_count) {
    size_t allocated_output_row_count = 0;
    switch (exe_unit.table_func.getOutputRowSizeType()) {
      case table_functions::OutputBufferSizeType::kConstant:
      case table_functions::OutputBufferSizeType::kUserSpecifiedConstantParameter: {
        allocated_output_row_count = exe_unit.output_buffer_size_param;
        break;
      }
      case table_functions::OutputBufferSizeType::kUserSpecifiedRowMultiplier: {
        allocated_output_row_count =
            exe_unit.output_buffer_size_param * input_element_count;
        break;
      }
      default: {
        UNREACHABLE();
      }
    }
    return allocated_output_row_count;
  }
  
  }  // namespace
  
  ResultSetPtr TableFunctionExecutionContext::execute(
      const TableFunctionExecutionUnit& exe_unit,
      const std::vector<InputTableInfo>& table_infos,
      const TableFunctionCompilationContext* compilation_context,
      const ColumnFetcher& column_fetcher,
      const ExecutorDeviceType device_type,
      Executor* executor) {
    CHECK(compilation_context);
    std::vector<std::shared_ptr<Chunk_NS::Chunk>> chunks_owner;
    std::vector<std::unique_ptr<char[]>> literals_owner;
  
    const int device_id = 0;  // TODO(adb): support multi-gpu table functions
    std::unique_ptr<CudaAllocator> device_allocator;
    if (device_type == ExecutorDeviceType::GPU) {
      auto& data_mgr = executor->catalog_->getDataMgr();
      device_allocator.reset(new CudaAllocator(&data_mgr, device_id));
    }
    std::vector<const int8_t*> col_buf_ptrs;
    std::vector<int64_t> col_sizes;
    std::optional<size_t> output_column_size;
    for (const auto& input_expr : exe_unit.input_exprs) {
      if (auto col_var = dynamic_cast<Analyzer::ColumnVar*>(input_expr)) {
        auto table_id = col_var->get_table_id();
        auto table_info_it = std::find_if(
            table_infos.begin(), table_infos.end(), [&table_id](const auto& table_info) {
              return table_info.table_id == table_id;
            });
        CHECK(table_info_it != table_infos.end());
        auto [col_buf, buf_elem_count] = ColumnFetcher::getOneColumnFragment(
            executor,
            *col_var,
            table_info_it->info.fragments.front(),
            device_type == ExecutorDeviceType::CPU ? Data_Namespace::MemoryLevel::CPU_LEVEL
                                                   : Data_Namespace::MemoryLevel::GPU_LEVEL,
            device_id,
            device_allocator.get(),
            chunks_owner,
            column_fetcher.columnarized_table_cache_);
  
        // We use the size of the first column to be the size of the output column
        if (!output_column_size) {
          output_column_size = buf_elem_count;
        }
  
        col_sizes.push_back(buf_elem_count);
        col_buf_ptrs.push_back(col_buf);
      } else if (const auto& constant_val = dynamic_cast<Analyzer::Constant*>(input_expr)) {
        // TODO(adb): Unify literal handling with rest of system, either in Codegen or as a
        // separate serialization component
        col_sizes.push_back(0);
        const auto const_val_datum = constant_val->get_constval();
        const auto& ti = constant_val->get_type_info();
        if (ti.is_fp()) {
          switch (get_bit_width(ti)) {
            case 32:
              col_buf_ptrs.push_back(create_literal_buffer(const_val_datum.floatval,
                                                           device_type,
                                                           literals_owner,
                                                           device_allocator.get()));
              break;
            case 64:
              col_buf_ptrs.push_back(create_literal_buffer(const_val_datum.doubleval,
                                                           device_type,
                                                           literals_owner,
                                                           device_allocator.get()));
              break;
            default:
              UNREACHABLE();
          }
        } else if (ti.is_integer()) {
          switch (get_bit_width(ti)) {
            case 8:
              col_buf_ptrs.push_back(create_literal_buffer(const_val_datum.tinyintval,
                                                           device_type,
                                                           literals_owner,
                                                           device_allocator.get()));
              break;
            case 16:
              col_buf_ptrs.push_back(create_literal_buffer(const_val_datum.smallintval,
                                                           device_type,
                                                           literals_owner,
                                                           device_allocator.get()));
              break;
            case 32:
              col_buf_ptrs.push_back(create_literal_buffer(const_val_datum.intval,
                                                           device_type,
                                                           literals_owner,
                                                           device_allocator.get()));
              break;
            case 64:
              col_buf_ptrs.push_back(create_literal_buffer(const_val_datum.bigintval,
                                                           device_type,
                                                           literals_owner,
                                                           device_allocator.get()));
              break;
            default:
              UNREACHABLE();
          }
        } else {
          throw std::runtime_error("Literal value " + constant_val->toString() +
                                   " is not yet supported.");
        }
      }
    }
    CHECK_EQ(col_buf_ptrs.size(), exe_unit.input_exprs.size());
    CHECK_EQ(col_sizes.size(), exe_unit.input_exprs.size());
    CHECK(output_column_size);
  
    switch (device_type) {
      case ExecutorDeviceType::CPU:
        return launchCpuCode(exe_unit,
                             compilation_context,
                             col_buf_ptrs,
                             col_sizes,
                             *output_column_size,
                             executor);
      case ExecutorDeviceType::GPU:
        return launchGpuCode(exe_unit,
                             compilation_context,
                             col_buf_ptrs,
                             col_sizes,
                             *output_column_size,
                             /*device_id=*/0,
                             executor);
    }
    UNREACHABLE();
    return nullptr;
  }
  
  ResultSetPtr TableFunctionExecutionContext::launchCpuCode(
      const TableFunctionExecutionUnit& exe_unit,
      const TableFunctionCompilationContext* compilation_context,
      std::vector<const int8_t*>& col_buf_ptrs,
      std::vector<int64_t>& col_sizes,
      const size_t elem_count,
      Executor* executor) {
    // setup the inputs
    const auto byte_stream_ptr = reinterpret_cast<const int8_t**>(col_buf_ptrs.data());
    CHECK(byte_stream_ptr);
  
    // initialize output memory
    auto num_out_columns = exe_unit.target_exprs.size();
    QueryMemoryDescriptor query_mem_desc(
        executor, elem_count, QueryDescriptionType::Projection, /*is_table_function=*/true);
    query_mem_desc.setOutputColumnar(true);
  
    for (size_t i = 0; i < num_out_columns; i++) {
      // All outputs padded to 8 bytes
      query_mem_desc.addColSlotInfo({std::make_tuple(8, 8)});
    }
  
    const auto allocated_output_row_count = get_output_row_count(exe_unit, elem_count);
    auto query_buffers = std::make_unique<QueryMemoryInitializer>(
        exe_unit,
        query_mem_desc,
        /*device_id=*/0,
        ExecutorDeviceType::CPU,
        allocated_output_row_count,
        std::vector<std::vector<const int8_t*>>{col_buf_ptrs},
        std::vector<std::vector<uint64_t>>{{0}},  // frag offsets
        row_set_mem_owner_,
        nullptr,
        executor);
  
    // setup the output
    int64_t output_row_count = allocated_output_row_count;
    auto group_by_buffers_ptr = query_buffers->getGroupByBuffersPtr();
    CHECK(group_by_buffers_ptr);
  
    auto output_buffers_ptr = reinterpret_cast<int64_t*>(group_by_buffers_ptr[0]);
    std::vector<int64_t*> output_col_buf_ptrs;
    for (size_t i = 0; i < num_out_columns; i++) {
      output_col_buf_ptrs.emplace_back(output_buffers_ptr + i * allocated_output_row_count);
    }
  
    // execute
    const auto err = compilation_context->getFuncPtr()(
        byte_stream_ptr, col_sizes.data(), output_col_buf_ptrs.data(), &output_row_count);
    if (err) {
      throw std::runtime_error("Error executing table function: " + std::to_string(err));
    }
    if (exe_unit.table_func.hasNonUserSpecifiedOutputSizeConstant()) {
      if (static_cast<size_t>(output_row_count) != allocated_output_row_count) {
        throw std::runtime_error(
            "Table function with constant sizing parameter must return " +
            std::to_string(allocated_output_row_count) + " (got " +
            std::to_string(output_row_count) + ")");
      }
    } else {
      if (output_row_count < 0 || (size_t)output_row_count > allocated_output_row_count) {
        output_row_count = allocated_output_row_count;
      }
    }
    // Update entry count, it may differ from allocated mem size
    query_buffers->getResultSet(0)->updateStorageEntryCount(output_row_count);
  
    const size_t column_size = output_row_count * sizeof(int64_t);
    const size_t allocated_column_size = allocated_output_row_count * sizeof(int64_t);
  
    int8_t* src = reinterpret_cast<int8_t*>(output_buffers_ptr);
    int8_t* dst = reinterpret_cast<int8_t*>(output_buffers_ptr);
    for (size_t i = 0; i < num_out_columns; i++) {
      if (src != dst) {
        auto t = memmove(dst, src, column_size);
        CHECK_EQ(dst, t);
      }
      src += allocated_column_size;
      dst += column_size;
    }
  
    return query_buffers->getResultSetOwned(0);
  }
  
  namespace {
  enum {
    ERROR_BUFFER,
    COL_BUFFERS,
    COL_SIZES,
    OUTPUT_BUFFERS,
    OUTPUT_ROW_COUNT,
    KERNEL_PARAM_COUNT,
  };
  }
  
  ResultSetPtr TableFunctionExecutionContext::launchGpuCode(
      const TableFunctionExecutionUnit& exe_unit,
      const TableFunctionCompilationContext* compilation_context,
      std::vector<const int8_t*>& col_buf_ptrs,
      std::vector<int64_t>& col_sizes,
      const size_t elem_count,
      const int device_id,
      Executor* executor) {
  #ifdef HAVE_CUDA
    auto num_out_columns = exe_unit.target_exprs.size();
    auto& data_mgr = executor->catalog_->getDataMgr();
    auto gpu_allocator = std::make_unique<CudaAllocator>(&data_mgr, device_id);
    CHECK(gpu_allocator);
    std::vector<CUdeviceptr> kernel_params(KERNEL_PARAM_COUNT, 0);
    // setup the inputs
    auto byte_stream_ptr = gpu_allocator->alloc(col_buf_ptrs.size() * sizeof(int64_t));
    gpu_allocator->copyToDevice(byte_stream_ptr,
                                reinterpret_cast<int8_t*>(col_buf_ptrs.data()),
                                col_buf_ptrs.size() * sizeof(int64_t));
    kernel_params[COL_BUFFERS] = reinterpret_cast<CUdeviceptr>(byte_stream_ptr);
  
    auto col_sizes_ptr = gpu_allocator->alloc(col_sizes.size() * sizeof(int64_t));
    gpu_allocator->copyToDevice(col_sizes_ptr,
                                reinterpret_cast<int8_t*>(col_sizes.data()),
                                col_sizes.size() * sizeof(int64_t));
    kernel_params[COL_SIZES] = reinterpret_cast<CUdeviceptr>(col_sizes_ptr);
  
    kernel_params[ERROR_BUFFER] =
        reinterpret_cast<CUdeviceptr>(gpu_allocator->alloc(sizeof(int32_t)));
    // initialize output memory
    QueryMemoryDescriptor query_mem_desc(
        executor, elem_count, QueryDescriptionType::Projection, /*is_table_function=*/true);
    query_mem_desc.setOutputColumnar(true);
  
    for (size_t i = 0; i < num_out_columns; i++) {
      // All outputs padded to 8 bytes
      query_mem_desc.addColSlotInfo({std::make_tuple(8, 8)});
    }
    const auto allocated_output_row_count = get_output_row_count(exe_unit, elem_count);
    auto query_buffers = std::make_unique<QueryMemoryInitializer>(
        exe_unit,
        query_mem_desc,
        device_id,
        ExecutorDeviceType::GPU,
        allocated_output_row_count,
        std::vector<std::vector<const int8_t*>>{col_buf_ptrs},
        std::vector<std::vector<uint64_t>>{{0}},  // frag offsets
        row_set_mem_owner_,
        gpu_allocator.get(),
        executor);
  
    // setup the output
    int64_t output_row_count = allocated_output_row_count;
  
    kernel_params[OUTPUT_ROW_COUNT] =
        reinterpret_cast<CUdeviceptr>(gpu_allocator->alloc(sizeof(int64_t*)));
    gpu_allocator->copyToDevice(reinterpret_cast<int8_t*>(kernel_params[OUTPUT_ROW_COUNT]),
                                reinterpret_cast<int8_t*>(&output_row_count),
                                sizeof(output_row_count));
  
    // const unsigned block_size_x = executor->blockSize();
    const unsigned block_size_x = 1;
    const unsigned block_size_y = 1;
    const unsigned block_size_z = 1;
    // const unsigned grid_size_x = executor->gridSize();
    const unsigned grid_size_x = 1;
    const unsigned grid_size_y = 1;
    const unsigned grid_size_z = 1;
  
    auto gpu_output_buffers = query_buffers->setupTableFunctionGpuBuffers(
        query_mem_desc, device_id, block_size_x, grid_size_x);
  
    kernel_params[OUTPUT_BUFFERS] = reinterpret_cast<CUdeviceptr>(gpu_output_buffers.first);
  
    // execute
    CHECK_EQ(static_cast<size_t>(KERNEL_PARAM_COUNT), kernel_params.size());
  
    std::vector<void*> param_ptrs;
    for (auto& param : kernel_params) {
      param_ptrs.push_back(&param);
    }
  
    // Get cu func
    const auto gpu_context = compilation_context->getGpuCode();
    CHECK(gpu_context);
    const auto native_code = gpu_context->getNativeCode(device_id);
    auto cu_func = static_cast<CUfunction>(native_code.first);
    checkCudaErrors(cuLaunchKernel(cu_func,
                                   grid_size_x,
                                   grid_size_y,
                                   grid_size_z,
                                   block_size_x,
                                   block_size_y,
                                   block_size_z,
                                   0,  // shared mem bytes
                                   nullptr,
                                   &param_ptrs[0],
                                   nullptr));
    // TODO(adb): read errors
  
    // read output row count from GPU
    gpu_allocator->copyFromDevice(
        reinterpret_cast<int8_t*>(&output_row_count),
        reinterpret_cast<int8_t*>(kernel_params[OUTPUT_ROW_COUNT]),
        sizeof(int64_t));
    if (exe_unit.table_func.hasNonUserSpecifiedOutputSizeConstant()) {
      if (static_cast<size_t>(output_row_count) != allocated_output_row_count) {
        throw std::runtime_error(
            "Table function with constant sizing parameter must return " +
            std::to_string(allocated_output_row_count) + " (got " +
            std::to_string(output_row_count) + ")");
      }
    } else {
      if (output_row_count < 0 || (size_t)output_row_count > allocated_output_row_count) {
        output_row_count = allocated_output_row_count;
      }
    }
  
    // Update entry count, it may differ from allocated mem size
    query_buffers->getResultSet(0)->updateStorageEntryCount(output_row_count);
  
    // Copy back to CPU storage
    query_buffers->copyFromTableFunctionGpuBuffers(&data_mgr,
                                                   query_mem_desc,
                                                   output_row_count,
                                                   gpu_output_buffers,
                                                   device_id,
                                                   block_size_x,
                                                   grid_size_x);
  
    return query_buffers->getResultSetOwned(0);
  #else
    UNREACHABLE();
    return nullptr;
  #endif
  }
  