// Copyright 2019 Google LLC
//
// 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
//
//      https://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.

// Utilities used by the bytecode_dispatch routines to aid in working with the
// bytecode stream and kernel dispatch.

#ifndef IREE_HAL_INTERPRETER_BYTECODE_DISPATCH_UTIL_H_
#define IREE_HAL_INTERPRETER_BYTECODE_DISPATCH_UTIL_H_

#include "absl/base/attributes.h"
#include "absl/container/inlined_vector.h"
#include "base/status.h"
#include "hal/buffer_view.h"
#include "hal/heap_buffer.h"
#include "hal/interpreter/bytecode_kernels.h"
#include "rt/function.h"
#include "rt/stack.h"
#include "schemas/bytecode/interpreter_bytecode_v0.h"
#include "vm/bytecode_reader.h"
#include "vm/type.h"

// TODO(benvanik): move to dedicated config file/build flags.
#define IREE_SUPPORT_F32 1
#define IREE_SUPPORT_F64 1

namespace iree {
namespace hal {

// Returns true if the contents of the BufferView are bitwise non-zero.
// Returns false if there is no buffer, the buffer is empty, or the contents are
// bitwise zero.
bool BufferViewIsTrue(const BufferView& buffer_view);

Status ValidateElementwiseUnaryOp(BufferView* src_local, BufferView* dst_local);
Status ValidateElementwiseBinaryOp(BufferView* lhs_local, BufferView* rhs_local,
                                   BufferView* dst_local);
Status ValidateElementwiseTernaryOp(BufferView* a_local, BufferView* b_local,
                                    BufferView* c_local, BufferView* dst_local);
Status ValidateMatMulOpI(BufferView* lhs_local, BufferView* rhs_local,
                         BufferView* bias_local,
                         BufferView* multiplier_mantissa_local,
                         BufferView* multiplier_exponent_local,
                         BufferView* dst_local);
Status ValidateMatMulOpF(BufferView* lhs_local, BufferView* rhs_local,
                         BufferView* bias_local, BufferView* dst_local);

template <typename KERNEL, typename T, typename... ARGS>
Status ApplyUnaryOp(BufferView* src_local, BufferView* dst_local,
                    ARGS... args) {
  // TODO(benvanik): avoid mapping by changing buffer type?
  ASSIGN_OR_RETURN(auto src_buffer,
                   src_local->buffer->MapMemory<T>(MemoryAccess::kRead));
  ASSIGN_OR_RETURN(auto dst_buffer, dst_local->buffer->MapMemory<T>(
                                        MemoryAccess::kDiscardWrite));
  return KERNEL::Execute(src_buffer.contents(), dst_buffer.mutable_contents(),
                         args...);
}

template <typename KERNEL, typename T, typename... ARGS>
Status ApplyBinaryOp(BufferView* lhs_local, BufferView* rhs_local,
                     BufferView* dst_local, ARGS... args) {
  ASSIGN_OR_RETURN(auto lhs_buffer,
                   lhs_local->buffer->MapMemory<T>(MemoryAccess::kRead));
  ASSIGN_OR_RETURN(auto rhs_buffer,
                   rhs_local->buffer->MapMemory<T>(MemoryAccess::kRead));
  ASSIGN_OR_RETURN(auto dst_buffer, dst_local->buffer->MapMemory<T>(
                                        MemoryAccess::kDiscardWrite));
  return KERNEL::Execute(lhs_buffer.contents(), rhs_buffer.contents(),
                         dst_buffer.mutable_contents(), args...);
}

template <typename KERNEL, typename T, typename... ARGS>
Status ApplyTernaryOp(BufferView* a_local, BufferView* b_local,
                      BufferView* c_local, BufferView* dst_local,
                      ARGS... args) {
  ASSIGN_OR_RETURN(auto a_buffer,
                   a_local->buffer->MapMemory<T>(MemoryAccess::kRead));
  ASSIGN_OR_RETURN(auto b_buffer,
                   b_local->buffer->MapMemory<T>(MemoryAccess::kRead));
  ASSIGN_OR_RETURN(auto c_buffer,
                   c_local->buffer->MapMemory<T>(MemoryAccess::kRead));
  ASSIGN_OR_RETURN(auto dst_buffer, dst_local->buffer->MapMemory<T>(
                                        MemoryAccess::kDiscardWrite));
  return KERNEL::Execute(a_buffer.contents(), b_buffer.contents(),
                         c_buffer.contents(), dst_buffer.mutable_contents(),
                         args...);
}

template <typename KERNEL, typename T>
Status ApplyComparisonOp(BufferView* lhs_local, BufferView* rhs_local,
                         BufferView* dst_local) {
  ASSIGN_OR_RETURN(auto lhs_buffer,
                   lhs_local->buffer->MapMemory<T>(MemoryAccess::kRead));
  ASSIGN_OR_RETURN(auto rhs_buffer,
                   rhs_local->buffer->MapMemory<T>(MemoryAccess::kRead));
  ASSIGN_OR_RETURN(auto dst_buffer, dst_local->buffer->MapMemory<uint8_t>(
                                        MemoryAccess::kDiscardWrite));
  return KERNEL::Execute(lhs_buffer.contents(), rhs_buffer.contents(),
                         dst_buffer.mutable_contents());
}

template <typename KERNEL, typename... ARGS>
Status ApplyUnaryOpIS(BufferView* src_local, BufferView* dst_local,
                      ARGS... args) {
  switch (src_local->element_size) {
    case 1:
      return ApplyUnaryOp<KERNEL, int8_t>(src_local, dst_local, args...);
    case 2:
      return ApplyUnaryOp<KERNEL, int16_t>(src_local, dst_local, args...);
    case 4:
      return ApplyUnaryOp<KERNEL, int32_t>(src_local, dst_local, args...);
    case 8:
      return ApplyUnaryOp<KERNEL, int64_t>(src_local, dst_local, args...);
    default:
      return UnimplementedErrorBuilder(IREE_LOC)
             << "Unimplemented element size: " << src_local->element_size;
  }
}

template <typename KERNEL, typename... ARGS>
Status ApplyUnaryOpIU(BufferView* src_local, BufferView* dst_local,
                      ARGS... args) {
  switch (src_local->element_size) {
    case 1:
      return ApplyUnaryOp<KERNEL, uint8_t>(src_local, dst_local, args...);
    case 2:
      return ApplyUnaryOp<KERNEL, uint16_t>(src_local, dst_local, args...);
    case 4:
      return ApplyUnaryOp<KERNEL, uint32_t>(src_local, dst_local, args...);
    case 8:
      return ApplyUnaryOp<KERNEL, uint64_t>(src_local, dst_local, args...);
    default:
      return UnimplementedErrorBuilder(IREE_LOC)
             << "Unimplemented element size: " << src_local->element_size;
  }
}

template <typename KERNEL, typename... ARGS>
Status ApplyUnaryOpF(BufferView* src_local, BufferView* dst_local,
                     ARGS... args) {
  switch (src_local->element_size) {
#if defined(IREE_SUPPORT_F32)
    case 4:
      return ApplyUnaryOp<KERNEL, float>(src_local, dst_local, args...);
#endif  // IREE_SUPPORT_F32
#if defined(IREE_SUPPORT_F64)
    case 8:
      return ApplyUnaryOp<KERNEL, double>(src_local, dst_local, args...);
#endif  // IREE_SUPPORT_F64
    default:
      return UnimplementedErrorBuilder(IREE_LOC)
             << "Unimplemented element size: " << src_local->element_size;
  }
}

template <typename KERNEL, typename... ARGS>
Status ApplyBinaryOpIS(BufferView* lhs_local, BufferView* rhs_local,
                       BufferView* dst_local, ARGS... args) {
  switch (lhs_local->element_size) {
    case 1:
      return ApplyBinaryOp<KERNEL, int8_t>(lhs_local, rhs_local, dst_local,
                                           args...);
    case 2:
      return ApplyBinaryOp<KERNEL, int16_t>(lhs_local, rhs_local, dst_local,
                                            args...);
    case 4:
      return ApplyBinaryOp<KERNEL, int32_t>(lhs_local, rhs_local, dst_local,
                                            args...);
    case 8:
      return ApplyBinaryOp<KERNEL, int64_t>(lhs_local, rhs_local, dst_local,
                                            args...);
    default:
      return UnimplementedErrorBuilder(IREE_LOC)
             << "Unimplemented element size: " << lhs_local->element_size;
  }
}

template <typename KERNEL, typename... ARGS>
Status ApplyBinaryOpIU(BufferView* lhs_local, BufferView* rhs_local,
                       BufferView* dst_local, ARGS... args) {
  switch (lhs_local->element_size) {
    case 1:
      return ApplyBinaryOp<KERNEL, uint8_t>(lhs_local, rhs_local, dst_local,
                                            args...);
    case 2:
      return ApplyBinaryOp<KERNEL, uint16_t>(lhs_local, rhs_local, dst_local,
                                             args...);
    case 4:
      return ApplyBinaryOp<KERNEL, uint32_t>(lhs_local, rhs_local, dst_local,
                                             args...);
    case 8:
      return ApplyBinaryOp<KERNEL, uint64_t>(lhs_local, rhs_local, dst_local,
                                             args...);
    default:
      return UnimplementedErrorBuilder(IREE_LOC)
             << "Unimplemented element size: " << lhs_local->element_size;
  }
}

template <typename KERNEL, typename... ARGS>
Status ApplyBinaryOpF(BufferView* lhs_local, BufferView* rhs_local,
                      BufferView* dst_local, ARGS... args) {
  switch (lhs_local->element_size) {
#if defined(IREE_SUPPORT_F32)
    case 4:
      return ApplyBinaryOp<KERNEL, float>(lhs_local, rhs_local, dst_local,
                                          args...);
#endif  // IREE_SUPPORT_F32
#if defined(IREE_SUPPORT_F64)
    case 8:
      return ApplyBinaryOp<KERNEL, double>(lhs_local, rhs_local, dst_local,
                                           args...);
#endif  // IREE_SUPPORT_F64
    default:
      return UnimplementedErrorBuilder(IREE_LOC)
             << "Unimplemented element size: " << lhs_local->element_size;
  }
}

template <typename KERNEL, typename... ARGS>
Status ApplyTernaryOpIS(BufferView* a_local, BufferView* b_local,
                        BufferView* c_local, BufferView* dst_local,
                        ARGS... args) {
  switch (a_local->element_size) {
    case 1:
      return ApplyTernaryOp<KERNEL, int8_t>(a_local, b_local, c_local,
                                            dst_local, args...);
    case 2:
      return ApplyTernaryOp<KERNEL, int16_t>(a_local, b_local, c_local,
                                             dst_local, args...);
    case 4:
      return ApplyTernaryOp<KERNEL, int32_t>(a_local, b_local, c_local,
                                             dst_local, args...);
    case 8:
      return ApplyTernaryOp<KERNEL, int64_t>(a_local, b_local, c_local,
                                             dst_local, args...);
    default:
      return UnimplementedErrorBuilder(IREE_LOC)
             << "Unimplemented element size: " << a_local->element_size;
  }
}

template <typename KERNEL, typename... ARGS>
Status ApplyTernaryOpIU(BufferView* a_local, BufferView* b_local,
                        BufferView* c_local, BufferView* dst_local,
                        ARGS... args) {
  switch (a_local->element_size) {
    case 1:
      return ApplyTernaryOp<KERNEL, uint8_t>(a_local, b_local, c_local,
                                             dst_local, args...);
    case 2:
      return ApplyTernaryOp<KERNEL, uint16_t>(a_local, b_local, c_local,
                                              dst_local, args...);
    case 4:
      return ApplyTernaryOp<KERNEL, uint32_t>(a_local, b_local, c_local,
                                              dst_local, args...);
    case 8:
      return ApplyTernaryOp<KERNEL, uint64_t>(a_local, b_local, c_local,
                                              dst_local, args...);
    default:
      return UnimplementedErrorBuilder(IREE_LOC)
             << "Unimplemented element size: " << a_local->element_size;
  }
}

template <typename KERNEL, typename... ARGS>
Status ApplyTernaryOpF(BufferView* a_local, BufferView* b_local,
                       BufferView* c_local, BufferView* dst_local,
                       ARGS... args) {
  switch (a_local->element_size) {
#if defined(IREE_SUPPORT_F32)
    case 4:
      return ApplyTernaryOp<KERNEL, float>(a_local, b_local, c_local, dst_local,
                                           args...);
#endif  // IREE_SUPPORT_F32
#if defined(IREE_SUPPORT_F64)
    case 8:
      return ApplyTernaryOp<KERNEL, double>(a_local, b_local, c_local,
                                            dst_local, args...);
#endif  // IREE_SUPPORT_F64
    default:
      return UnimplementedErrorBuilder(IREE_LOC)
             << "Unimplemented element size: " << a_local->element_size;
  }
}

template <typename KERNEL>
Status ApplyComparisonOpIS(BufferView* lhs_local, BufferView* rhs_local,
                           BufferView* dst_local) {
  switch (lhs_local->element_size) {
    case 1:
      return ApplyComparisonOp<KERNEL, int8_t>(lhs_local, rhs_local, dst_local);
    case 2:
      return ApplyComparisonOp<KERNEL, int16_t>(lhs_local, rhs_local,
                                                dst_local);
    case 4:
      return ApplyComparisonOp<KERNEL, int32_t>(lhs_local, rhs_local,
                                                dst_local);
    case 8:
      return ApplyComparisonOp<KERNEL, int64_t>(lhs_local, rhs_local,
                                                dst_local);
    default:
      return UnimplementedErrorBuilder(IREE_LOC)
             << "Unimplemented element size: " << lhs_local->element_size;
  }
}

template <typename KERNEL>
Status ApplyComparisonOpIU(BufferView* lhs_local, BufferView* rhs_local,
                           BufferView* dst_local) {
  switch (lhs_local->element_size) {
    case 1:
      return ApplyComparisonOp<KERNEL, uint8_t>(lhs_local, rhs_local,
                                                dst_local);
    case 2:
      return ApplyComparisonOp<KERNEL, uint16_t>(lhs_local, rhs_local,
                                                 dst_local);
    case 4:
      return ApplyComparisonOp<KERNEL, uint32_t>(lhs_local, rhs_local,
                                                 dst_local);
    case 8:
      return ApplyComparisonOp<KERNEL, uint64_t>(lhs_local, rhs_local,
                                                 dst_local);
    default:
      return UnimplementedErrorBuilder(IREE_LOC)
             << "Unimplemented element size: " << lhs_local->element_size;
  }
}

template <typename KERNEL>
Status ApplyComparisonOpF(BufferView* lhs_local, BufferView* rhs_local,
                          BufferView* dst_local) {
  switch (lhs_local->element_size) {
    case 4:
      return ApplyComparisonOp<KERNEL, float>(lhs_local, rhs_local, dst_local);
    case 8:
      return ApplyComparisonOp<KERNEL, double>(lhs_local, rhs_local, dst_local);
    default:
      return UnimplementedErrorBuilder(IREE_LOC)
             << "Unimplemented element size: " << lhs_local->element_size;
  }
}

template <typename T, typename ACC = int32_t>
Status ApplyMatMulOpI(kernels::MatMul::RuntimeState* runtime_state,
                      BufferView* lhs_local, BufferView* rhs_local,
                      BufferView* bias_local,
                      BufferView* multiplier_mantissa_local,
                      BufferView* multiplier_exponent_local,
                      BufferView* dst_local) {
  kernels::MatMul::Buffers<T, ACC> buffers;
  ASSIGN_OR_RETURN(auto lhs_buffer,
                   lhs_local->buffer->MapMemory<T>(MemoryAccess::kRead));
  buffers.lhs_buffer = lhs_buffer.contents();
  buffers.lhs_shape = lhs_local->shape;
  ASSIGN_OR_RETURN(auto rhs_buffer,
                   rhs_local->buffer->MapMemory<T>(MemoryAccess::kRead));
  buffers.rhs_buffer = rhs_buffer.contents();
  buffers.rhs_shape = rhs_local->shape;
  MappedMemory<ACC> bias_buffer;
  if (bias_local && bias_local->buffer && !bias_local->shape.empty()) {
    if (bias_local->element_size != sizeof(ACC)) {
      return UnimplementedErrorBuilder(IREE_LOC)
             << "Only " << sizeof(ACC) << "b biases are supported right now";
    }
    ASSIGN_OR_RETURN(bias_buffer,
                     bias_local->buffer->MapMemory<ACC>(MemoryAccess::kRead));
    buffers.bias_buffer = bias_buffer.contents();
  }
  ASSIGN_OR_RETURN(
      auto multiplier_mantissa_buffer,
      multiplier_mantissa_local->buffer->MapMemory<ACC>(MemoryAccess::kRead));
  buffers.multiplier_mantissa_buffer = multiplier_mantissa_buffer.contents();
  ASSIGN_OR_RETURN(auto multiplier_exponent_buffer,
                   multiplier_exponent_local->buffer->MapMemory<int32_t>(
                       MemoryAccess::kRead));
  buffers.multiplier_exponent_buffer = multiplier_exponent_buffer.contents();
  ASSIGN_OR_RETURN(auto dst_buffer, dst_local->buffer->MapMemory<T>(
                                        MemoryAccess::kDiscardWrite));
  buffers.dst_buffer = dst_buffer.mutable_contents();
  buffers.dst_shape = dst_local->shape;
  return kernels::MatMul::Execute(runtime_state, buffers);
}

template <typename T>
Status ApplyMatMulOpF(kernels::MatMul::RuntimeState* runtime_state,
                      BufferView* lhs_local, BufferView* rhs_local,
                      BufferView* bias_local, BufferView* dst_local) {
  kernels::MatMul::Buffers<T, T> buffers;
  ASSIGN_OR_RETURN(auto lhs_buffer,
                   lhs_local->buffer->MapMemory<T>(MemoryAccess::kRead));
  buffers.lhs_buffer = lhs_buffer.contents();
  buffers.lhs_shape = lhs_local->shape;
  ASSIGN_OR_RETURN(auto rhs_buffer,
                   rhs_local->buffer->MapMemory<T>(MemoryAccess::kRead));
  buffers.rhs_buffer = rhs_buffer.contents();
  buffers.rhs_shape = rhs_local->shape;
  MappedMemory<T> bias_buffer;
  if (bias_local && bias_local->buffer && !bias_local->shape.empty()) {
    ASSIGN_OR_RETURN(bias_buffer,
                     bias_local->buffer->MapMemory<T>(MemoryAccess::kRead));
    buffers.bias_buffer = bias_buffer.contents();
  }
  ASSIGN_OR_RETURN(auto dst_buffer, dst_local->buffer->MapMemory<T>(
                                        MemoryAccess::kDiscardWrite));
  buffers.dst_buffer = dst_buffer.mutable_contents();
  buffers.dst_shape = dst_local->shape;
  return kernels::MatMul::Execute(runtime_state, buffers);
}

template <typename KERNEL>
Status DispatchElementwiseUnaryOpIS(vm::BytecodeReader* reader) {
  ASSIGN_OR_RETURN(auto* src_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* dst_local, reader->ReadLocal());
  RETURN_IF_ERROR(ValidateElementwiseUnaryOp(src_local, dst_local));
  return ApplyUnaryOpIS<KERNEL>(src_local, dst_local);
}

template <typename KERNEL>
Status DispatchElementwiseUnaryOpIU(vm::BytecodeReader* reader) {
  ASSIGN_OR_RETURN(auto* src_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* dst_local, reader->ReadLocal());
  RETURN_IF_ERROR(ValidateElementwiseUnaryOp(src_local, dst_local));
  return ApplyUnaryOpIU<KERNEL>(src_local, dst_local);
}

template <typename KERNEL>
Status DispatchElementwiseUnaryOpF(vm::BytecodeReader* reader) {
  ASSIGN_OR_RETURN(auto* src_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* dst_local, reader->ReadLocal());
  RETURN_IF_ERROR(ValidateElementwiseUnaryOp(src_local, dst_local));
  return ApplyUnaryOpF<KERNEL>(src_local, dst_local);
}

template <typename KERNEL>
Status DispatchElementwiseBinaryOpIS(vm::BytecodeReader* reader) {
  ASSIGN_OR_RETURN(auto* lhs_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* rhs_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* dst_local, reader->ReadLocal());
  RETURN_IF_ERROR(ValidateElementwiseBinaryOp(lhs_local, rhs_local, dst_local));
  return ApplyBinaryOpIS<KERNEL>(lhs_local, rhs_local, dst_local);
}

template <typename KERNEL>
Status DispatchElementwiseBinaryOpIU(vm::BytecodeReader* reader) {
  ASSIGN_OR_RETURN(auto* lhs_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* rhs_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* dst_local, reader->ReadLocal());
  RETURN_IF_ERROR(ValidateElementwiseBinaryOp(lhs_local, rhs_local, dst_local));
  return ApplyBinaryOpIU<KERNEL>(lhs_local, rhs_local, dst_local);
}

template <typename KERNEL>
Status DispatchElementwiseBinaryOpF(vm::BytecodeReader* reader) {
  ASSIGN_OR_RETURN(auto* lhs_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* rhs_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* dst_local, reader->ReadLocal());
  RETURN_IF_ERROR(ValidateElementwiseBinaryOp(lhs_local, rhs_local, dst_local));
  return ApplyBinaryOpF<KERNEL>(lhs_local, rhs_local, dst_local);
}

template <typename KERNEL>
Status DispatchElementwiseTernaryOpIS(vm::BytecodeReader* reader) {
  ASSIGN_OR_RETURN(auto* a_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* b_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* c_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* dst_local, reader->ReadLocal());
  RETURN_IF_ERROR(
      ValidateElementwiseTernaryOp(a_local, b_local, c_local, dst_local));
  return ApplyTernaryOpIS<KERNEL>(a_local, b_local, c_local, dst_local);
}

template <typename KERNEL>
Status DispatchElementwiseTernaryOpIU(vm::BytecodeReader* reader) {
  ASSIGN_OR_RETURN(auto* a_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* b_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* c_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* dst_local, reader->ReadLocal());
  RETURN_IF_ERROR(
      ValidateElementwiseTernaryOp(a_local, b_local, c_local, dst_local));
  return ApplyTernaryOpIU<KERNEL>(a_local, b_local, c_local, dst_local);
}

template <typename KERNEL>
Status DispatchElementwiseTernaryOpF(vm::BytecodeReader* reader) {
  ASSIGN_OR_RETURN(auto* a_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* b_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* c_local, reader->ReadLocal());
  ASSIGN_OR_RETURN(auto* dst_local, reader->ReadLocal());
  RETURN_IF_ERROR(
      ValidateElementwiseTernaryOp(a_local, b_local, c_local, dst_local));
  return ApplyTernaryOpF<KERNEL>(a_local, b_local, c_local, dst_local);
}

Status ApplyCopy(BufferView* src_local, absl::Span<const int32_t> src_indices,
                 BufferView* dst_local, absl::Span<const int32_t> dst_indices,
                 absl::Span<const int32_t> lengths);

}  // namespace hal
}  // namespace iree

#endif  // IREE_HAL_INTERPRETER_BYTECODE_DISPATCH_UTIL_H_