// 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.

// Implements a full bytecode dispatch system.
// Currently this is verbose and object oriented, but future revisions
// (once we have interesting benchmarks) will likely simplify and inline
// a lot of the checks to make things faster. Consider this to be as
// experimental an implementation as the entire rest of the project :)

#include "third_party/mlir_edge/iree/hal/interpreter/bytecode_dispatch.h"

#include <algorithm>

#include "third_party/absl/base/attributes.h"
#include "third_party/absl/container/inlined_vector.h"
#include "third_party/absl/types/span.h"
#include "third_party/mlir_edge/iree/base/logging.h"
#include "third_party/mlir_edge/iree/base/memory.h"
#include "third_party/mlir_edge/iree/base/status.h"
#include "third_party/mlir_edge/iree/hal/buffer_view.h"
#include "third_party/mlir_edge/iree/hal/heap_buffer.h"
#include "third_party/mlir_edge/iree/hal/interpreter/bytecode_dispatch_conversion.h"
#include "third_party/mlir_edge/iree/hal/interpreter/bytecode_dispatch_util.h"
#include "third_party/mlir_edge/iree/hal/interpreter/bytecode_kernels.h"
#include "third_party/mlir_edge/iree/schemas/bytecode/interpreter_bytecode_v0.h"
#include "third_party/mlir_edge/iree/vm/bytecode_reader.h"
#include "third_party/mlir_edge/iree/vm/bytecode_tables_interpreter.h"
#include "third_party/mlir_edge/iree/vm/bytecode_util.h"
#include "third_party/mlir_edge/iree/vm/function.h"
#include "third_party/mlir_edge/iree/vm/opcode_info.h"

namespace iree {
namespace hal {

namespace {

using ::iree::vm::BytecodeReader;
using ::iree::vm::ImportFunction;
using ::iree::vm::NativeFunction;
using ::iree::vm::Stack;
using ::iree::vm::StackFrame;

}  // namespace

Status Dispatch(hal::Allocator* allocator,
                kernels::RuntimeState* kernel_runtime_state, Stack* stack,
                StackFrame* entry_stack_frame,
                absl::Span<BufferView> entry_results) {
  // Dispatch table mapping 1:1 with bytecode ops.
  // Each entry is a label within this function that can be used for computed
  // goto. You can find more information on computed goto here:
  // https://eli.thegreenplace.net/2012/07/12/computed-goto-for-efficient-dispatch-tables
  //
  // Note that we ensure the table is 256 elements long exactly to make sure
  // that unused opcodes are handled gracefully.
  static const void* kDispatchTable[256] = {
#define DECLARE_DISPATCH(ordinal, name, ...) &&_dispatch_##name,
#define DECLARE_DISPATCH_RESERVED(ordinal, name, ...) &&_dispatch_unhandled,
      IREE_INTERPRETER_OPCODE_LIST(DECLARE_DISPATCH, DECLARE_DISPATCH_RESERVED)
#undef DECLARE_DISPATCH
#undef DECLARE_DISPATCH_RESERVED
  };

  // Primary dispatch state. This is our 'native stack frame' and really just
  // enough to make dereferencing common addresses (like the current offset)
  // faster. You can think of this like CPU state (like PC).
  //
  // We hope that LLVM decides to keep these in registers (as they are touched
  // for every instruction executed). The stack_frame will change as we call
  // into different functions.
  BytecodeReader reader(stack);
  RETURN_IF_ERROR(reader.SwitchStackFrame(entry_stack_frame));

#define DISPATCH_NEXT()                                                    \
  {                                                                        \
    uint8_t opcode = *reader.AdvanceOffset().ValueOrDie();                 \
    DVLOG(1)                                                               \
        << "Interpreter dispatching op code: "                             \
        << GetOpcodeInfo(vm::interpreter_opcode_table(), opcode).mnemonic; \
    goto* kDispatchTable[opcode];                                          \
  }

#define DISPATCH_CORE_OPCODE(opcode, body) \
  _dispatch_##opcode : {body} DISPATCH_NEXT()
#if defined(IREE_SUPPORT_F32) || defined(IREE_SUPPORT_F64)
#define DISPATCH_FLOAT_OPCODE(opcode, body) \
  _dispatch_##opcode : {body} DISPATCH_NEXT()
#else
#define DISPATCH_FLOAT_OPCODE(...)
#endif  // IREE_SUPPORT_F32 || IREE_SUPPORT_F64

  DISPATCH_NEXT();

  DISPATCH_CORE_OPCODE(kConstant, {
    ASSIGN_OR_RETURN(auto value, reader.ReadConstant());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    *dst_local = std::move(value);
  });

  DISPATCH_CORE_OPCODE(kCall, {
    auto* old_stack_frame = stack->current_frame();
    ASSIGN_OR_RETURN(const auto& target_function, reader.ReadFunction());
    ASSIGN_OR_RETURN(auto* new_stack_frame, stack->PushFrame(target_function));
    RETURN_IF_ERROR(
        reader.CopyInputsAndSwitchStackFrame(old_stack_frame, new_stack_frame));
    DVLOG(1) << "Call; stack now: " << stack->DebugString();
  });

  DISPATCH_CORE_OPCODE(kCallImport, {
    auto* old_stack_frame = stack->current_frame();
    ASSIGN_OR_RETURN(const auto* target_function, reader.ReadImportFunction());
    switch (target_function->link_type()) {
      case ImportFunction::LinkType::kModule: {
        ASSIGN_OR_RETURN(auto* new_stack_frame,
                         stack->PushFrame(target_function->linked_function()));
        RETURN_IF_ERROR(reader.CopyInputsAndSwitchStackFrame(old_stack_frame,
                                                             new_stack_frame));
        DVLOG(1) << "Call module import; stack now: " << stack->DebugString();
        break;
      }
      case ImportFunction::LinkType::kNativeFunction: {
        ASSIGN_OR_RETURN(auto* new_stack_frame,
                         stack->PushFrame(*target_function));
        RETURN_IF_ERROR(reader.CopyInputsAndSwitchStackFrame(old_stack_frame,
                                                             new_stack_frame));
        DVLOG(1) << "Call native import; stack now: " << stack->DebugString();
        RETURN_IF_ERROR(CallNativeFunction(stack, *target_function));
        RETURN_IF_ERROR(reader.CopyResultsAndSwitchStackFrame(old_stack_frame,
                                                              new_stack_frame));
        RETURN_IF_ERROR(stack->PopFrame());
        DVLOG(1) << "Return from native; stack now: " << stack->DebugString();
        break;
      }
    }
  });

  DISPATCH_CORE_OPCODE(kCallIndirect, {
    return UnimplementedErrorBuilder(ABSL_LOC) << "Unimplemented call_indirect";
  });

  DISPATCH_CORE_OPCODE(kReturn, {
    auto* old_stack_frame = stack->current_frame();
    auto* new_stack_frame = stack->caller_frame();
    if (old_stack_frame == entry_stack_frame) {
      // Returning from entry function. Marshal results from the return stmt.
      ASSIGN_OR_RETURN(int32_t src_count, reader.ReadCount());
      for (int i = 0; i < src_count; ++i) {
        ASSIGN_OR_RETURN(auto* src_local,
                         reader.ReadLocal(old_stack_frame->mutable_locals()));
        entry_results[i] = std::move(*src_local);
      }
      DVLOG(1) << "Returning to entry";
      return OkStatus();
    } else if (!new_stack_frame) {
      return FailedPreconditionErrorBuilder(ABSL_LOC) << "Stack underflow";
    }
    RETURN_IF_ERROR(reader.CopyResultsAndSwitchStackFrame(old_stack_frame,
                                                          new_stack_frame));
    RETURN_IF_ERROR(stack->PopFrame());
    DVLOG(1) << "Return; stack now: " << stack->DebugString();
  });

  DISPATCH_CORE_OPCODE(kBranch, {
    ASSIGN_OR_RETURN(int32_t offset, reader.ReadBlockOffset());
    RETURN_IF_ERROR(reader.CopySlots());
    RETURN_IF_ERROR(reader.BranchToOffset(offset));
  });

  DISPATCH_CORE_OPCODE(kCondBranch, {
    // Evaluate condition first so we can do the copies as we read them for
    // which side of the branch we take.
    ASSIGN_OR_RETURN(auto* cond_local, reader.ReadLocal());
    bool cond_value = BufferViewIsTrue(*cond_local);
    ASSIGN_OR_RETURN(int32_t true_offset, reader.ReadBlockOffset());
    if (cond_value) {
      RETURN_IF_ERROR(reader.CopySlots());
      RETURN_IF_ERROR(reader.BranchToOffset(true_offset));
    } else {
      ASSIGN_OR_RETURN(int32_t true_op_count, reader.ReadCount());
      RETURN_IF_ERROR(reader.SkipLocals(2 * true_op_count));
      ASSIGN_OR_RETURN(int32_t false_offset, reader.ReadBlockOffset());
      RETURN_IF_ERROR(reader.CopySlots());
      RETURN_IF_ERROR(reader.BranchToOffset(false_offset));
    }
  });

  DISPATCH_CORE_OPCODE(kCmpI, {
    ASSIGN_OR_RETURN(uint8_t predicate, reader.ReadUint8_t());
    ASSIGN_OR_RETURN(auto* lhs_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* rhs_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());

    switch (static_cast<CmpIPredicate>(predicate)) {
      case CmpIPredicate::kEq:
        RETURN_IF_ERROR(ApplyComparisonOpIS<kernels::CompareEQ>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpIPredicate::kNe:
        RETURN_IF_ERROR(ApplyComparisonOpIS<kernels::CompareNE>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpIPredicate::kSlt:
        RETURN_IF_ERROR(ApplyComparisonOpIS<kernels::CompareLT>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpIPredicate::kSle:
        RETURN_IF_ERROR(ApplyComparisonOpIS<kernels::CompareLE>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpIPredicate::kSgt:
        RETURN_IF_ERROR(ApplyComparisonOpIS<kernels::CompareGT>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpIPredicate::kSge:
        RETURN_IF_ERROR(ApplyComparisonOpIS<kernels::CompareGE>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpIPredicate::kUlt:
        RETURN_IF_ERROR(ApplyComparisonOpIU<kernels::CompareLT>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpIPredicate::kUle:
        RETURN_IF_ERROR(ApplyComparisonOpIU<kernels::CompareLE>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpIPredicate::kUgt:
        RETURN_IF_ERROR(ApplyComparisonOpIU<kernels::CompareGT>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpIPredicate::kUge:
        RETURN_IF_ERROR(ApplyComparisonOpIU<kernels::CompareGE>(
            lhs_local, rhs_local, dst_local));
        break;
    }
  });

  DISPATCH_FLOAT_OPCODE(kCmpF, {
    ASSIGN_OR_RETURN(uint8_t p, reader.ReadUint8_t());
    ASSIGN_OR_RETURN(auto* lhs_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* rhs_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());

    auto predicate = static_cast<CmpFPredicate>(p);
    switch (predicate) {
      case CmpFPredicate::kOeq:
        RETURN_IF_ERROR(ApplyComparisonOpF<kernels::CompareEQ>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpFPredicate::kUne:
        RETURN_IF_ERROR(ApplyComparisonOpF<kernels::CompareNE>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpFPredicate::kOlt:
        RETURN_IF_ERROR(ApplyComparisonOpF<kernels::CompareLT>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpFPredicate::kOle:
        RETURN_IF_ERROR(ApplyComparisonOpF<kernels::CompareLE>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpFPredicate::kOgt:
        RETURN_IF_ERROR(ApplyComparisonOpF<kernels::CompareGT>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpFPredicate::kOge:
        RETURN_IF_ERROR(ApplyComparisonOpF<kernels::CompareGE>(
            lhs_local, rhs_local, dst_local));
        break;
      case CmpFPredicate::kFalse:
      case CmpFPredicate::kOne:
      case CmpFPredicate::kOrd:
      case CmpFPredicate::kUeq:
      case CmpFPredicate::kUgt:
      case CmpFPredicate::kUge:
      case CmpFPredicate::kUlt:
      case CmpFPredicate::kUle:
      case CmpFPredicate::kUno:
      case CmpFPredicate::kTrue:
        // TODO(b/132183250) support these if we ever need them.
        return UnimplementedErrorBuilder(ABSL_LOC)
               << "Unsupported comparison predicate value "
               << static_cast<int>(p) << " ("
               << vm::PredicateToString(predicate) << ")";
    }
  });

  DISPATCH_CORE_OPCODE(kAllocStatic, {
    return UnimplementedErrorBuilder(ABSL_LOC) << "Unimplemented alloc_static";
  });

  DISPATCH_CORE_OPCODE(kAllocStack, {
    return UnimplementedErrorBuilder(ABSL_LOC) << "Unimplemented alloc_stack";
  });

  DISPATCH_CORE_OPCODE(kAllocStackInit, {
    return UnimplementedErrorBuilder(ABSL_LOC)
           << "Unimplemented alloc_stack_init";
  });

  DISPATCH_CORE_OPCODE(kAllocHeap, {
    ASSIGN_OR_RETURN(auto heap_type, reader.ReadInt32());
    ASSIGN_OR_RETURN(auto type, reader.ReadType());
    size_t element_size = type.element_size();

    // TODO(benvanik): more efficient reading and storage.
    size_t element_count = 0;
    ASSIGN_OR_RETURN(auto shape, reader.ReadShapePieces(&element_count));
    size_t allocation_size = element_size * element_count;

    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    dst_local->element_size = element_size;
    dst_local->shape = shape;

    // TODO(benvanik): properly allocate with attributes from op.
    CHECK_EQ(heap_type, 0);
    ASSIGN_OR_RETURN(
        dst_local->buffer,
        allocator->Allocate(MemoryType::kHostLocal | MemoryType::kDeviceVisible,
                            BufferUsage::kAll, allocation_size));
  });

  DISPATCH_CORE_OPCODE(kDiscard, {
    // NOTE: if we were an encoder we would actually discard the buffer.
    ASSIGN_OR_RETURN(auto* local, reader.ReadLocal());
    *local = {};
  });

  DISPATCH_CORE_OPCODE(kRank, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    int32_t rank = src_local->shape.size();
    RETURN_IF_ERROR(dst_local->buffer->WriteData(0, &rank, sizeof(int32_t)));
  });

  DISPATCH_CORE_OPCODE(kDim, {
    ASSIGN_OR_RETURN(int32_t axis, reader.ReadUint8_t());
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(int32_t dim, src_local->shape.ResolveAxis(axis));
    RETURN_IF_ERROR(dst_local->buffer->WriteData(0, &dim, sizeof(int32_t)));
  });

  DISPATCH_CORE_OPCODE(kShape, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    RETURN_IF_ERROR(dst_local->buffer->WriteData(
        0, src_local->shape.subspan().data(),
        src_local->shape.subspan().size() * sizeof(int32_t)));
  });

  DISPATCH_CORE_OPCODE(kLength, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    int32_t length = src_local->shape.element_count();
    RETURN_IF_ERROR(dst_local->buffer->WriteData(0, &length, sizeof(int32_t)));
  });

  DISPATCH_CORE_OPCODE(kDynamicSlice, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto indices, reader.ReadSlotElements<int32_t>());
    ASSIGN_OR_RETURN(auto lengths, reader.ReadSlotElements<int32_t>());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(*dst_local, src_local->Slice(indices, lengths));
  });

  DISPATCH_CORE_OPCODE(kStaticSlice, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto indices, reader.ReadIndexList());
    ASSIGN_OR_RETURN(auto lengths, reader.ReadIndexList());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(*dst_local, src_local->Slice(indices, lengths));
  });

  DISPATCH_CORE_OPCODE(kDynamicCopy, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto src_indices, reader.ReadSlotElements<int32_t>());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto dst_indices, reader.ReadSlotElements<int32_t>());
    ASSIGN_OR_RETURN(auto lengths, reader.ReadSlotElements<int32_t>());
    RETURN_IF_ERROR(
        ApplyCopy(src_local, src_indices, dst_local, dst_indices, lengths));
  });

  DISPATCH_CORE_OPCODE(kStaticCopy, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto src_indices, reader.ReadIndexList());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto dst_indices, reader.ReadIndexList());
    ASSIGN_OR_RETURN(auto lengths, reader.ReadIndexList());
    RETURN_IF_ERROR(
        ApplyCopy(src_local, src_indices, dst_local, dst_indices, lengths));
  });

  DISPATCH_CORE_OPCODE(kClone, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    dst_local->element_size = src_local->element_size;
    dst_local->shape = src_local->shape;
    dst_local->buffer = HeapBuffer::Allocate(src_local->buffer->usage(),
                                             src_local->buffer->byte_length());
    RETURN_IF_ERROR(dst_local->buffer->CopyData(0, src_local->buffer.get()));
  });

  DISPATCH_CORE_OPCODE(kSplit, {
    return UnimplementedErrorBuilder(ABSL_LOC) << "Unimplemented split";
  });

  DISPATCH_CORE_OPCODE(kAssign, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    *dst_local = *src_local;
  });

  DISPATCH_CORE_OPCODE(kCondAssign, {
    ASSIGN_OR_RETURN(auto* cond_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* lhs_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* rhs_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    *dst_local = BufferViewIsTrue(*cond_local) ? *lhs_local : *rhs_local;
  });

  DISPATCH_CORE_OPCODE(kReshape, {
    // TODO(benvanik): more logic required if strides differ.
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto shape_data, reader.ReadSlotElements<int32_t>());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    Shape new_shape = Shape{shape_data};
    if (src_local->shape.element_count() != new_shape.element_count()) {
      return InvalidArgumentErrorBuilder(ABSL_LOC)
             << "New element count " << new_shape.element_count()
             << " != source element count " << src_local->shape.element_count();
    }
    dst_local->shape = new_shape;
    dst_local->buffer = add_ref(src_local->buffer);
    dst_local->element_size = src_local->element_size;
  });

  DISPATCH_CORE_OPCODE(kSelect, {
    ASSIGN_OR_RETURN(auto* cond_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* lhs_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* rhs_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto cond_buffer, cond_local->buffer->MapMemory<uint8_t>(
                                           MemoryAccess::kRead));
    ASSIGN_OR_RETURN(auto lhs_buffer, lhs_local->buffer->MapMemory<uint8_t>(
                                          MemoryAccess::kRead));
    ASSIGN_OR_RETURN(auto rhs_buffer, rhs_local->buffer->MapMemory<uint8_t>(
                                          MemoryAccess::kRead));
    ASSIGN_OR_RETURN(auto dst_buffer, dst_local->buffer->MapMemory<uint8_t>(
                                          MemoryAccess::kDiscardWrite));
    if (cond_local->element_size != 1) {
      return InvalidArgumentErrorBuilder(ABSL_LOC) << "Select cond must be i8";
    } else if (lhs_buffer.size() != rhs_buffer.size()) {
      return InvalidArgumentErrorBuilder(ABSL_LOC)
             << "LHS " << lhs_buffer.size() << "b != RHS " << rhs_buffer.size()
             << "b; both arguments must match";
    } else if (lhs_buffer.size() != dst_buffer.size()) {
      return InvalidArgumentErrorBuilder(ABSL_LOC)
             << "Dest " << dst_buffer.size() << "b != LHS/RHS "
             << lhs_buffer.size() << "b; dest must match inputs";
    }
    switch (lhs_local->element_size) {
      case 1:
        RETURN_IF_ERROR(kernels::Select::Execute<uint8_t>(
            cond_buffer.contents(), lhs_buffer.contents(),
            rhs_buffer.contents(), dst_buffer.mutable_contents()));
        break;
      case 2:
        RETURN_IF_ERROR(kernels::Select::Execute<uint16_t>(
            cond_buffer.contents(),
            ReinterpretSpan<uint16_t>(lhs_buffer.contents()),
            ReinterpretSpan<uint16_t>(rhs_buffer.contents()),
            ReinterpretSpan<uint16_t>(dst_buffer.mutable_contents())));
        break;
      case 4:
        RETURN_IF_ERROR(kernels::Select::Execute<uint32_t>(
            cond_buffer.contents(),
            ReinterpretSpan<uint32_t>(lhs_buffer.contents()),
            ReinterpretSpan<uint32_t>(rhs_buffer.contents()),
            ReinterpretSpan<uint32_t>(dst_buffer.mutable_contents())));
        break;
      case 8:
        RETURN_IF_ERROR(kernels::Select::Execute<uint64_t>(
            cond_buffer.contents(),
            ReinterpretSpan<uint64_t>(lhs_buffer.contents()),
            ReinterpretSpan<uint64_t>(rhs_buffer.contents()),
            ReinterpretSpan<uint64_t>(dst_buffer.mutable_contents())));
        break;
      default:
        return UnimplementedErrorBuilder(ABSL_LOC)
               << "Unimplemented element size: " << lhs_local->element_size;
    }
  });

  DISPATCH_CORE_OPCODE(kTranspose, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto perm_data, reader.ReadSlotElements<int32_t>());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    RETURN_IF_ERROR(ApplyUnaryOpIU<kernels::Transpose>(
        src_local, dst_local, src_local->shape,
        absl::MakeConstSpan(perm_data)));
  });

  DISPATCH_CORE_OPCODE(kReverse, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto perm_data, reader.ReadSlotElements<int32_t>());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    RETURN_IF_ERROR(
        ApplyUnaryOpIU<kernels::Reverse>(src_local, dst_local, src_local->shape,
                                         absl::MakeConstSpan(perm_data)));
  });

  DISPATCH_CORE_OPCODE(kPad, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* padding_value, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto edge_padding_low, reader.ReadSlotElements<int32_t>());
    ASSIGN_OR_RETURN(auto edge_padding_high,
                     reader.ReadSlotElements<int32_t>());
    ASSIGN_OR_RETURN(auto interior_padding, reader.ReadSlotElements<int32_t>());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());

    RETURN_IF_ERROR(ApplyBinaryOpIU<kernels::Pad>(
        src_local, padding_value, dst_local, src_local->shape, dst_local->shape,
        absl::MakeConstSpan(edge_padding_low),
        absl::MakeConstSpan(edge_padding_high),
        absl::MakeConstSpan(interior_padding)));
  });

  DISPATCH_CORE_OPCODE(kBroadcast, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto shape_data, reader.ReadSlotElements<int32_t>());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    dst_local->shape = Shape{shape_data};
    RETURN_IF_ERROR(ApplyUnaryOpIU<kernels::Broadcast>(src_local, dst_local));
  });

  DISPATCH_CORE_OPCODE(kTile, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto shape_data, reader.ReadSlotElements<int32_t>());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    dst_local->shape = Shape{shape_data};
    RETURN_IF_ERROR(ApplyUnaryOpIU<kernels::Tile>(
        src_local, dst_local, src_local->shape, dst_local->shape));
  });

  DISPATCH_CORE_OPCODE(kNot, {
    RETURN_IF_ERROR(DispatchElementwiseUnaryOpIU<kernels::Not>(&reader));
  });
  DISPATCH_CORE_OPCODE(kAnd, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpIU<kernels::And>(&reader));
  });
  DISPATCH_CORE_OPCODE(kOr, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpIU<kernels::Or>(&reader));
  });
  DISPATCH_CORE_OPCODE(kXor, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpIU<kernels::Xor>(&reader));
  });
  DISPATCH_CORE_OPCODE(kShiftLeft, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpIU<kernels::ShiftLeft>(&reader));
  });
  DISPATCH_CORE_OPCODE(kShiftRightLogical, {
    RETURN_IF_ERROR(
        DispatchElementwiseBinaryOpIU<kernels::ShiftRight>(&reader));
  });
  DISPATCH_CORE_OPCODE(kShiftRightArithmetic, {
    RETURN_IF_ERROR(
        DispatchElementwiseBinaryOpIS<kernels::ShiftRight>(&reader));
  });

  DISPATCH_CORE_OPCODE(kAddI, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpIU<kernels::Add>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kAddF, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpF<kernels::Add>(&reader));
  });

  DISPATCH_CORE_OPCODE(kSubI, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpIU<kernels::Sub>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kSubF, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpF<kernels::Sub>(&reader));
  });

  DISPATCH_CORE_OPCODE(kAbsI, {
    RETURN_IF_ERROR(DispatchElementwiseUnaryOpIS<kernels::Abs>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kAbsF, {
    RETURN_IF_ERROR(DispatchElementwiseUnaryOpF<kernels::Abs>(&reader));
  });

  DISPATCH_CORE_OPCODE(kMulI, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpIU<kernels::Mul>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kMulF, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpF<kernels::Mul>(&reader));
  });

  DISPATCH_CORE_OPCODE(kDivIS, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpIS<kernels::Div>(&reader));
  });
  DISPATCH_CORE_OPCODE(kDivIU, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpIU<kernels::Div>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kDivF, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpF<kernels::Div>(&reader));
  });

  DISPATCH_CORE_OPCODE(kMulAddI, {
    RETURN_IF_ERROR(DispatchElementwiseTernaryOpIU<kernels::MulAdd>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kMulAddF, {
    RETURN_IF_ERROR(DispatchElementwiseTernaryOpF<kernels::MulAdd>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kExpF, {
    RETURN_IF_ERROR(DispatchElementwiseUnaryOpF<kernels::Exp>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kLogF, {
    RETURN_IF_ERROR(DispatchElementwiseUnaryOpF<kernels::Log>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kRsqrtF, {
    RETURN_IF_ERROR(DispatchElementwiseUnaryOpF<kernels::Rsqrt>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kCosF, {
    RETURN_IF_ERROR(DispatchElementwiseUnaryOpF<kernels::Cos>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kSinF, {
    RETURN_IF_ERROR(DispatchElementwiseUnaryOpF<kernels::Sin>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kTanhF, {
    RETURN_IF_ERROR(DispatchElementwiseUnaryOpF<kernels::Tanh>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kAtan2F, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpF<kernels::Atan2>(&reader));
  });

  DISPATCH_CORE_OPCODE(kMinIS, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpIS<kernels::Min>(&reader));
  });
  DISPATCH_CORE_OPCODE(kMinIU, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpIU<kernels::Min>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kMinF, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpF<kernels::Min>(&reader));
  });

  DISPATCH_CORE_OPCODE(kMaxIS, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpIS<kernels::Max>(&reader));
  });
  DISPATCH_CORE_OPCODE(kMaxIU, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpIU<kernels::Max>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kMaxF, {
    RETURN_IF_ERROR(DispatchElementwiseBinaryOpF<kernels::Max>(&reader));
  });

  DISPATCH_CORE_OPCODE(kClampIS, {
    RETURN_IF_ERROR(DispatchElementwiseTernaryOpIS<kernels::Clamp>(&reader));
  });
  DISPATCH_CORE_OPCODE(kClampIU, {
    RETURN_IF_ERROR(DispatchElementwiseTernaryOpIS<kernels::Clamp>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kClampF, {
    RETURN_IF_ERROR(DispatchElementwiseTernaryOpF<kernels::Clamp>(&reader));
  });

  DISPATCH_FLOAT_OPCODE(kFloorF, {
    RETURN_IF_ERROR(DispatchElementwiseUnaryOpF<kernels::Floor>(&reader));
  });
  DISPATCH_FLOAT_OPCODE(kCeilF, {
    RETURN_IF_ERROR(DispatchElementwiseUnaryOpF<kernels::Ceil>(&reader));
  });

  DISPATCH_CORE_OPCODE(kConvertSS, {
    ASSIGN_OR_RETURN(auto src_type, reader.ReadType());
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto dst_type, reader.ReadType());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    RETURN_IF_ERROR(
        ApplyConvertSS::Apply(src_type, src_local, dst_type, dst_local));
  });
  DISPATCH_CORE_OPCODE(kConvertUU, {
    ASSIGN_OR_RETURN(auto src_type, reader.ReadType());
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto dst_type, reader.ReadType());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    RETURN_IF_ERROR(
        ApplyConvertUU::Apply(src_type, src_local, dst_type, dst_local));
  });
  DISPATCH_CORE_OPCODE(kConvertSU, {
    ASSIGN_OR_RETURN(auto src_type, reader.ReadType());
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto dst_type, reader.ReadType());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    RETURN_IF_ERROR(
        ApplyConvertSU::Apply(src_type, src_local, dst_type, dst_local));
  });
  DISPATCH_CORE_OPCODE(kConvertUS, {
    ASSIGN_OR_RETURN(auto src_type, reader.ReadType());
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto dst_type, reader.ReadType());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    RETURN_IF_ERROR(
        ApplyConvertUS::Apply(src_type, src_local, dst_type, dst_local));
  });

  DISPATCH_CORE_OPCODE(kMatMulI, {
    ASSIGN_OR_RETURN(auto* lhs_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* rhs_local, reader.ReadLocal());
    // TODO(benvanik): add fused matmul-with-bias op in MLIR and lower to this.
    BufferView* bias_local = nullptr;
    ASSIGN_OR_RETURN(auto* multiplier_mantissa_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* multiplier_exponent_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    RETURN_IF_ERROR(ValidateMatMulOpI(lhs_local, rhs_local, bias_local,
                                      multiplier_mantissa_local,
                                      multiplier_exponent_local, dst_local));
    auto* mat_mul_state = kernel_runtime_state->mat_mul_state.get();
    // TODO(benvanik): define as a matrix of supported types to enable 8*8=16,
    // accumulator options, and other precision modes.
    switch (lhs_local->element_size) {
      case 1:
        RETURN_IF_ERROR(ApplyMatMulOpI<int8_t>(
            mat_mul_state, lhs_local, rhs_local, bias_local,
            multiplier_mantissa_local, multiplier_exponent_local, dst_local));
        break;
      case 2:
        RETURN_IF_ERROR(ApplyMatMulOpI<int16_t>(
            mat_mul_state, lhs_local, rhs_local, bias_local,
            multiplier_mantissa_local, multiplier_exponent_local, dst_local));
        break;
      case 4:
        RETURN_IF_ERROR(ApplyMatMulOpI<int32_t>(
            mat_mul_state, lhs_local, rhs_local, bias_local,
            multiplier_mantissa_local, multiplier_exponent_local, dst_local));
        break;
      case 8:
        RETURN_IF_ERROR(ApplyMatMulOpI<int64_t>(
            mat_mul_state, lhs_local, rhs_local, bias_local,
            multiplier_mantissa_local, multiplier_exponent_local, dst_local));
        break;
      default:
        return UnimplementedErrorBuilder(ABSL_LOC)
               << "Unimplemented element size: " << lhs_local->element_size;
    }
  });

  DISPATCH_FLOAT_OPCODE(kMatMulF, {
    ASSIGN_OR_RETURN(auto* lhs_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* rhs_local, reader.ReadLocal());
    BufferView* bias_local = nullptr;
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    RETURN_IF_ERROR(
        ValidateMatMulOpF(lhs_local, rhs_local, bias_local, dst_local));
    auto* mat_mul_state = kernel_runtime_state->mat_mul_state.get();
    switch (lhs_local->element_size) {
      case 4:
        RETURN_IF_ERROR(ApplyMatMulOpF<float>(
            mat_mul_state, lhs_local, rhs_local, bias_local, dst_local));
        break;
      case 8:
        RETURN_IF_ERROR(ApplyMatMulOpF<double>(
            mat_mul_state, lhs_local, rhs_local, bias_local, dst_local));
        break;
      default:
        return UnimplementedErrorBuilder(ABSL_LOC)
               << "Unimplemented element size: " << lhs_local->element_size;
    }
  });

  DISPATCH_CORE_OPCODE(kReduceSumI, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* init_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto dimension, reader.ReadInt32());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    // TODO(scotttodd): validate
    RETURN_IF_ERROR(ApplyBinaryOpIS<kernels::ReduceSum>(
        src_local, init_local, dst_local, dimension, src_local->shape,
        dst_local->shape));
  });

  DISPATCH_FLOAT_OPCODE(kReduceSumF, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* init_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto dimension, reader.ReadInt32());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    // TODO(scotttodd): validate
    RETURN_IF_ERROR(ApplyBinaryOpF<kernels::ReduceSum>(
        src_local, init_local, dst_local, dimension, src_local->shape,
        dst_local->shape));
  });

  DISPATCH_CORE_OPCODE(kReduceMinI, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* init_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto dimension, reader.ReadInt32());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    // TODO(scotttodd): validate
    RETURN_IF_ERROR(ApplyBinaryOpIS<kernels::ReduceMin>(
        src_local, init_local, dst_local, dimension, src_local->shape,
        dst_local->shape));
  });

  DISPATCH_FLOAT_OPCODE(kReduceMinF, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* init_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto dimension, reader.ReadInt32());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    // TODO(scotttodd): validate
    RETURN_IF_ERROR(ApplyBinaryOpF<kernels::ReduceMin>(
        src_local, init_local, dst_local, dimension, src_local->shape,
        dst_local->shape));
  });

  DISPATCH_CORE_OPCODE(kReduceMaxI, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* init_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto dimension, reader.ReadInt32());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    // TODO(scotttodd): validate
    RETURN_IF_ERROR(ApplyBinaryOpIS<kernels::ReduceMax>(
        src_local, init_local, dst_local, dimension, src_local->shape,
        dst_local->shape));
  });

  DISPATCH_FLOAT_OPCODE(kReduceMaxF, {
    ASSIGN_OR_RETURN(auto* src_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto* init_local, reader.ReadLocal());
    ASSIGN_OR_RETURN(auto dimension, reader.ReadInt32());
    ASSIGN_OR_RETURN(auto* dst_local, reader.ReadLocal());
    // TODO(scotttodd): validate
    RETURN_IF_ERROR(ApplyBinaryOpF<kernels::ReduceMax>(
        src_local, init_local, dst_local, dimension, src_local->shape,
        dst_local->shape));
  });

  DISPATCH_CORE_OPCODE(kTrace, {
    return UnimplementedErrorBuilder(ABSL_LOC) << "Unimplemented trace";
  });

  DISPATCH_CORE_OPCODE(kBreak, {
    return UnimplementedErrorBuilder(ABSL_LOC) << "Unimplemented break";
  });

  DISPATCH_CORE_OPCODE(kCondBreak, {
    return UnimplementedErrorBuilder(ABSL_LOC) << "Unimplemented cond_break";
  });

_dispatch_unhandled:
  // TODO(benvanik): better tracing.
  return UnimplementedErrorBuilder(ABSL_LOC) << "Unknown dispatch opcode";
}  // NOLINT(readability/fn_size)

}  // namespace hal
}  // namespace iree