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

#include "iree/hal/vmla/vmla_module.h"

#include <cstdint>

#include "absl/types/span.h"
#include "iree/base/tracing.h"
#include "iree/hal/vmla/op_kernels.h"
#include "iree/vm/module_abi_packing.h"

//===----------------------------------------------------------------------===//
// Type registration
//===----------------------------------------------------------------------===//

static iree_vm_ref_type_descriptor_t Buffer_descriptor = {0};
static iree_vm_ref_type_descriptor_t Interface_descriptor = {0};

IREE_VM_DEFINE_TYPE_ADAPTERS(Buffer, iree::hal::vmla::Buffer);
IREE_VM_DEFINE_TYPE_ADAPTERS(Interface, iree::hal::vmla::Interface);

#define IREE_VMLA_REGISTER_CC_TYPE(type, name, descriptor)     \
  descriptor.type_name = iree_make_cstring_view(name);         \
  descriptor.offsetof_counter = type::offsetof_counter();      \
  descriptor.destroy = type::DirectDestroy;                    \
  IREE_RETURN_IF_ERROR(iree_vm_ref_register_type(&descriptor)) \
      << "Failed to register type " << name;

namespace iree {
namespace hal {
namespace vmla {

Status ModuleRegisterTypes() {
  static bool has_registered = false;
  if (has_registered) return OkStatus();

  IREE_VMLA_REGISTER_CC_TYPE(Buffer, "vmla.buffer", Buffer_descriptor);
  IREE_VMLA_REGISTER_CC_TYPE(Interface, "vmla.interface", Interface_descriptor);

  has_registered = true;
  return OkStatus();
}

//===----------------------------------------------------------------------===//
// API type implementations
//===----------------------------------------------------------------------===//

// static
StatusOr<vm::ref<Buffer>> Buffer::Allocate(size_t byte_length,
                                           iree_allocator_t allocator) {
  void* data = nullptr;
  IREE_RETURN_IF_ERROR(iree_allocator_malloc(allocator, byte_length, &data))
      << "Failed to allocate buffer of size " << byte_length;

  auto buffer = vm::assign_ref(new Buffer());
  buffer->data_ = data;
  buffer->data_length_ = byte_length;
  buffer->allocator_ = allocator;
  return std::move(buffer);
}

// static
StatusOr<vm::ref<Buffer>> Buffer::Wrap(const void* data, size_t data_length,
                                       iree_allocator_t allocator) {
  auto buffer = vm::assign_ref(new Buffer());
  buffer->data_ = const_cast<void*>(data);
  buffer->data_length_ = data_length;
  buffer->allocator_ = allocator;
  return std::move(buffer);
}

// static
StatusOr<vm::ref<Buffer>> Buffer::WrapMutable(void* data, size_t data_length,
                                              iree_allocator_t allocator) {
  auto buffer = vm::assign_ref(new Buffer());
  buffer->data_ = data;
  buffer->data_length_ = data_length;
  buffer->allocator_ = allocator;
  return std::move(buffer);
}

Buffer::~Buffer() {
  if (!parent_) {
    iree_allocator_free(allocator_, data_);
    data_ = nullptr;
  }
  parent_.reset();
}

StatusOr<absl::Span<uint8_t>> Buffer::MakeRange(
    iree_vmla_size_t byte_offset, iree_vmla_size_t byte_length) const {
  if (byte_length == kVMLAWholeBuffer) {
    byte_length = size() - byte_offset;
  }
  if (byte_offset > size()) {
    return OutOfRangeErrorBuilder(IREE_LOC)
           << "Attempted to access an address off the end of the valid "
              "buffer range (offset="
           << byte_offset << ", length=" << byte_length
           << ", buffer byte_length=" << size() << ")";
  }
  size_t end = byte_offset + byte_length - 1;
  if (end >= size()) {
    return OutOfRangeErrorBuilder(IREE_LOC)
           << "Attempted to access an address outside of the valid buffer "
              "range (offset="
           << byte_offset << ", length=" << byte_length << ", end=" << end
           << ", buffer byte_length=" << size() << ")";
  }
  uint8_t* data = reinterpret_cast<uint8_t*>(data_) + byte_offset;
  size_t data_length = byte_length;
  return absl::MakeSpan(data, data_length);
}

constexpr int Interface::kMaxConstants;
constexpr int Interface::kMaxSets;
constexpr int Interface::kMaxBindings;

void Interface::Reset() {
  for (int i = 0; i < bindings_.size(); ++i) {
    for (int j = 0; j < bindings_[i].size(); ++j) {
      bindings_[i][j] = {};
    }
  }
}

StatusOr<uint32_t> Interface::GetConstant(uint32_t offset) const {
  if (offset >= kMaxConstants) {
    return InvalidArgumentErrorBuilder(IREE_LOC)
           << "Invalid constant offset=" << offset;
  }
  return constants_[offset];
}

Status Interface::SetConstants(absl::Span<const uint32_t> values) {
  if (values.size() > kMaxConstants) {
    return InvalidArgumentErrorBuilder(IREE_LOC)
           << "Constant value overflow; have " << values.size()
           << " but max is " << kMaxConstants;
  }
  for (int i = 0; i < values.size(); ++i) {
    constants_[i] = values[i];
  }
  return OkStatus();
}

StatusOr<const Interface::Binding> Interface::GetBinding(
    int32_t set, int32_t binding) const {
  if (set < 0 || set > kMaxSets || binding < 0 || binding > kMaxBindings) {
    return InvalidArgumentErrorBuilder(IREE_LOC)
           << "Invalid binding set=" << set << ", binding=" << binding;
  }
  return bindings_[set][binding];
}

Status Interface::SetBinding(int32_t set, int32_t binding, Binding value) {
  if (set < 0 || set > kMaxSets || binding < 0 || binding > kMaxBindings) {
    return InvalidArgumentErrorBuilder(IREE_LOC)
           << "Invalid binding set=" << set << ", binding=" << binding;
  }
  bindings_[set][binding] = std::move(value);
  return OkStatus();
}

//===----------------------------------------------------------------------===//
// Module state and method implementation
//===----------------------------------------------------------------------===//

namespace {

// Per-executable VMLA module state.
// This provides the exported kernel functions to the VM and is instantiated
// one or more times per executable used within a device. Any state here can be
// treated as workgroup-local memory.
//
// Thread-compatible.
class VMLAModuleState final {
 public:
  VMLAModuleState(iree_allocator_t allocator,
                  kernels::RuntimeState* kernel_state)
      : allocator_(allocator), kernel_state_(kernel_state) {}

  ~VMLAModuleState() = default;

  //===--------------------------------------------------------------------===//
  // vmla.interface.*
  //===--------------------------------------------------------------------===//

  StatusOr<uint32_t> InterfaceConst(vm::ref<Interface> interface,
                                    uint32_t offset) {
    IREE_TRACE_SCOPE0("VMLAModuleState::InterfaceConst");
    return interface->GetConstant(offset);
  }

  StatusOr<vm::ref<Buffer>> InterfaceBinding(vm::ref<Interface> interface,
                                             int32_t set, int32_t binding) {
    IREE_TRACE_SCOPE0("VMLAModuleState::InterfaceBinding");
    IREE_ASSIGN_OR_RETURN(const auto& value,
                          interface->GetBinding(set, binding));
    return vm::retain_ref(value.buffer);
  }

  //===--------------------------------------------------------------------===//
  // vmla.buffer.*
  //===--------------------------------------------------------------------===//

  StatusOr<vm::ref<Buffer>> BufferConst(
      vm::ref<iree_vm_ro_byte_buffer_t> value) {
    IREE_TRACE_SCOPE0("VMLAModuleState::BufferConst");
    iree_allocator_t external_allocator = {0};
    external_allocator.self = vm::retain_ref(value).release();
    external_allocator.free = +[](void* self, void* ptr) {
      vm::assign_ref(reinterpret_cast<iree_vm_ro_byte_buffer_t*>(self)).reset();
    };
    return Buffer::Wrap(value->data.data, value->data.data_length,
                        external_allocator);
  }

  StatusOr<vm::ref<Buffer>> BufferAlloc(iree_vmla_size_t byte_length) {
    IREE_TRACE_SCOPE0("VMLAModuleState::BufferAlloc");
    return Buffer::Allocate(byte_length, allocator_);
  }

  StatusOr<vm::ref<Buffer>> BufferClone(vm::ref<Buffer> src) {
    IREE_TRACE_SCOPE0("VMLAModuleState::BufferClone");
    IREE_ASSIGN_OR_RETURN(auto dst, Buffer::Allocate(src->size(), allocator_));
    std::memcpy(dst->data(), src->data(), dst->size());
    return std::move(dst);
  }

  StatusOr<iree_vmla_size_t> BufferByteLength(vm::ref<Buffer> buffer) {
    IREE_TRACE_SCOPE0("VMLAModuleState::BufferByteLength");
    return buffer->size();
  }

  StatusOr<vm::ref<Buffer>> BufferView(vm::ref<Buffer> src,
                                       iree_vmla_size_t byte_offset,
                                       iree_vmla_size_t byte_length) {
    IREE_TRACE_SCOPE0("VMLAModuleState::BufferView");

    if (byte_length == kVMLAWholeBuffer) {
      byte_length = src->size() - byte_offset;
    }

    if (byte_offset == 0 && byte_length == src->size()) {
      // Asking for the same buffer.
      return vm::retain_ref(src);
    } else if (byte_offset > src->size()) {
      return OutOfRangeErrorBuilder(IREE_LOC)
             << "Attempted to access an address off the end of the valid "
                "buffer range (offset="
             << byte_offset << ", length=" << byte_length
             << ", buffer byte_length=" << src->size() << ")";
    }
    size_t end = byte_offset + byte_length - 1;
    if (end >= src->size()) {
      return OutOfRangeErrorBuilder(IREE_LOC)
             << "Attempted to access an address outside of the valid buffer "
                "range (offset="
             << byte_offset << ", length=" << byte_length << ", end=" << end
             << ", buffer byte_length=" << src->size() << ")";
    }
    uint8_t* data = reinterpret_cast<uint8_t*>(src->data()) + byte_offset;
    size_t data_length = byte_length;

    iree_allocator_t external_allocator = {0};
    external_allocator.self = vm::retain_ref(src).release();
    external_allocator.free = +[](void* self, void* ptr) {
      vm::assign_ref(reinterpret_cast<Buffer*>(self)).reset();
    };
    return Buffer::Wrap(data, data_length, external_allocator);
  }

  Status BufferCopy(vm::ref<Buffer> src, iree_vmla_size_t src_byte_offset,
                    vm::ref<Buffer> dst, iree_vmla_size_t dst_byte_offset,
                    iree_vmla_size_t byte_length) {
    IREE_TRACE_SCOPE0("VMLAModuleState::BufferCopy");
    if (byte_length == kVMLAWholeBuffer) {
      byte_length = src->size() - src_byte_offset;
    }
    IREE_ASSIGN_OR_RETURN(auto src_bytes, src->RangeAs<const uint8_t>(
                                              src_byte_offset, byte_length));
    IREE_ASSIGN_OR_RETURN(auto dst_bytes,
                          dst->RangeAs<uint8_t>(dst_byte_offset, byte_length));
    std::memcpy(dst_bytes.data(), src_bytes.data(), dst_bytes.size());
    return OkStatus();
  }

  Status BufferFill(vm::ref<Buffer> value, vm::ref<Buffer> dst) {
    IREE_TRACE_SCOPE0("VMLAModuleState::BufferFill");
    if (value->size() == 1) {
      // Fast-path for single-byte memset values.
      std::memset(dst->data(), value->As<uint8_t>()[0], dst->size());
      return OkStatus();
    } else if (dst->size() % value->size() != 0) {
      return InvalidArgumentErrorBuilder(IREE_LOC)
             << "Fill value length (" << value->size()
             << ") must divide evenly into buffer length (" << dst->size()
             << ")";
    }
    auto value_bytes = value->As<uint8_t>();
    auto dst_bytes = dst->As<uint8_t>();
    for (size_t i = 0; i < dst_bytes.size(); i += value_bytes.size()) {
      std::memcpy(dst_bytes.data() + i, value_bytes.data(), value_bytes.size());
    }
    return OkStatus();
  }

  StatusOr<int32_t> BufferLoadI32(vm::ref<Buffer> src,
                                  iree_vmla_size_t byte_offset) {
    IREE_TRACE_SCOPE0("VMLAModuleState::BufferLoadI32");
    IREE_ASSIGN_OR_RETURN(auto data,
                          src->RangeAs<int32_t>(byte_offset, sizeof(int32_t)));
    return data[0];
  }

  //===--------------------------------------------------------------------===//
  // Common helpers for defining ops
  //===--------------------------------------------------------------------===//

#define IREE_VMLA_NONARY_OP(name, kernel, type)    \
  Status name(vm::ref<Buffer> dst) {               \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);  \
    return kernel::Execute<type>(dst->As<type>()); \
  }

#define IREE_VMLA_UNARY_OP(name, kernel, type)                      \
  Status name(vm::ref<Buffer> src, vm::ref<Buffer> dst) {           \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                   \
    return kernel::Execute<type>(src->As<type>(), dst->As<type>()); \
  }

#define IREE_VMLA_BINARY_OP(name, kernel, type)                                \
  Status name(vm::ref<Buffer> lhs, vm::ref<Buffer> rhs, vm::ref<Buffer> dst) { \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                              \
    return kernel::Execute<type>(lhs->As<type>(), rhs->As<type>(),             \
                                 dst->As<type>());                             \
  }

#define IREE_VMLA_TERNARY_OP(name, kernel, type)                              \
  Status name(vm::ref<Buffer> a, vm::ref<Buffer> b, vm::ref<Buffer> c,        \
              vm::ref<Buffer> dst) {                                          \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                             \
    return kernel::Execute<type>(a->As<type>(), b->As<type>(), c->As<type>(), \
                                 dst->As<type>());                            \
  }

  //===--------------------------------------------------------------------===//
  // VMLA Ops: comparison
  //===--------------------------------------------------------------------===//

  enum class CmpPredicate : uint32_t {
    kEQ = 0,
    kNE = 1,
    kLT = 2,
    kLE = 3,
    kGT = 4,
    kGE = 5,
  };

#define IREE_VMLA_COMPARE_OP(name, type)                                   \
  Status name(int32_t predicate, vm::ref<Buffer> lhs, vm::ref<Buffer> rhs, \
              vm::ref<Buffer> dst) {                                       \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                          \
    switch (static_cast<CmpPredicate>(predicate)) {                        \
      case CmpPredicate::kEQ:                                              \
        return kernels::CompareEQ::Execute<type>(                          \
            lhs->As<type>(), rhs->As<type>(), dst->As<uint8_t>());         \
      case CmpPredicate::kNE:                                              \
        return kernels::CompareNE::Execute<type>(                          \
            lhs->As<type>(), rhs->As<type>(), dst->As<uint8_t>());         \
      case CmpPredicate::kLT:                                              \
        return kernels::CompareLT::Execute<type>(                          \
            lhs->As<type>(), rhs->As<type>(), dst->As<uint8_t>());         \
      case CmpPredicate::kLE:                                              \
        return kernels::CompareLE::Execute<type>(                          \
            lhs->As<type>(), rhs->As<type>(), dst->As<uint8_t>());         \
      case CmpPredicate::kGT:                                              \
        return kernels::CompareGT::Execute<type>(                          \
            lhs->As<type>(), rhs->As<type>(), dst->As<uint8_t>());         \
      case CmpPredicate::kGE:                                              \
        return kernels::CompareGE::Execute<type>(                          \
            lhs->As<type>(), rhs->As<type>(), dst->As<uint8_t>());         \
      default:                                                             \
        return InvalidArgumentErrorBuilder(IREE_LOC)                       \
               << "Unsupported predicate " << predicate;                   \
    }                                                                      \
  }
  IREE_VMLA_COMPARE_OP(CmpI8, int8_t);
  IREE_VMLA_COMPARE_OP(CmpI16, int16_t);
  IREE_VMLA_COMPARE_OP(CmpI32, int32_t);
  IREE_VMLA_COMPARE_OP(CmpF32, float);

#define IREE_VMLA_SELECT_OP(name, type)                                       \
  Status name(vm::ref<Buffer> cond, vm::ref<Buffer> lhs, vm::ref<Buffer> rhs, \
              vm::ref<Buffer> dst) {                                          \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                             \
    return kernels::Select::Execute<type>(cond->As<uint8_t>(),                \
                                          lhs->As<type>(), rhs->As<type>(),   \
                                          dst->As<type>());                   \
  }
  IREE_VMLA_SELECT_OP(SelectX8, uint8_t);
  IREE_VMLA_SELECT_OP(SelectX16, uint16_t);
  IREE_VMLA_SELECT_OP(SelectX32, uint32_t);

  //===--------------------------------------------------------------------===//
  // VMLA Ops: shape/structure
  //===--------------------------------------------------------------------===//

#define IREE_VMLA_COPY_OP(name, size)                                     \
  Status name(vm::ref<Buffer> src, iree_vmla_shape_t src_shape,           \
              absl::Span<const int32_t> src_indices, vm::ref<Buffer> dst, \
              iree_vmla_shape_t dst_shape,                                \
              absl::Span<const int32_t> dst_indices,                      \
              absl::Span<const int32_t> lengths) {                        \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                         \
    return kernels::Copy::Execute<size>(src->As<uint8_t>(), src_shape,    \
                                        src_indices, dst->As<uint8_t>(),  \
                                        dst_shape, dst_indices, lengths); \
  }
  IREE_VMLA_COPY_OP(CopyX8, sizeof(uint8_t));
  IREE_VMLA_COPY_OP(CopyX16, sizeof(uint16_t));
  IREE_VMLA_COPY_OP(CopyX32, sizeof(uint32_t));

#define IREE_VMLA_TRANSPOSE_OP(name, type)                                     \
  Status name(vm::ref<Buffer> src, iree_vmla_shape_t src_shape,                \
              absl::Span<const int32_t> permutation, vm::ref<Buffer> dst,      \
              iree_vmla_shape_t dst_shape) {                                   \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                              \
    return kernels::Transpose::Execute<type>(src->As<type>(), dst->As<type>(), \
                                             src_shape, permutation);          \
  }
  IREE_VMLA_TRANSPOSE_OP(TransposeX8, uint8_t);
  IREE_VMLA_TRANSPOSE_OP(TransposeX16, uint16_t);
  IREE_VMLA_TRANSPOSE_OP(TransposeX32, uint32_t);

#define IREE_VMLA_REVERSE_OP(name, type)                                     \
  Status name(vm::ref<Buffer> src, iree_vmla_shape_t src_shape,              \
              absl::Span<const int32_t> dims, vm::ref<Buffer> dst,           \
              iree_vmla_shape_t dst_shape) {                                 \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                            \
    return kernels::Reverse::Execute<type>(src->As<type>(), dst->As<type>(), \
                                           src_shape, dims);                 \
  }
  IREE_VMLA_REVERSE_OP(ReverseX8, uint8_t);
  IREE_VMLA_REVERSE_OP(ReverseX16, uint16_t);
  IREE_VMLA_REVERSE_OP(ReverseX32, uint32_t);

#define IREE_VMLA_PAD_OP(name, type)                                       \
  Status name(vm::ref<Buffer> src, iree_vmla_shape_t src_shape,            \
              vm::ref<Buffer> value, iree_vmla_shape_t value_shape,        \
              vm::ref<Buffer> dst, iree_vmla_shape_t dst_shape,            \
              absl::Span<const int32_t> edge_padding_low,                  \
              absl::Span<const int32_t> edge_padding_high,                 \
              absl::Span<const int32_t> interior_padding) {                \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                          \
    return kernels::Pad::Execute<type>(                                    \
        src->As<type>(), value->As<type>(), dst->As<type>(), src_shape,    \
        dst_shape, edge_padding_low, edge_padding_high, interior_padding); \
  }
  IREE_VMLA_PAD_OP(PadX8, uint8_t);
  IREE_VMLA_PAD_OP(PadX16, uint16_t);
  IREE_VMLA_PAD_OP(PadX32, uint32_t);

#define IREE_VMLA_GATHER_OP(name, type)                                  \
  Status name(vm::ref<Buffer> src, iree_vmla_shape_t src_shape,          \
              vm::ref<Buffer> indices, iree_vmla_shape_t indices_shape,  \
              vm::ref<Buffer> dst, iree_vmla_shape_t dst_shape,          \
              const int32_t dim, const int32_t batch_dims) {             \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                        \
    return kernels::Gather::Execute<type>(                               \
        src->As<type>(), indices->As<int>(), dst->As<type>(), src_shape, \
        indices_shape, dst_shape, dim, batch_dims);                      \
  }
  IREE_VMLA_GATHER_OP(GatherX8, uint8_t);
  IREE_VMLA_GATHER_OP(GatherX16, uint16_t);
  IREE_VMLA_GATHER_OP(GatherX32, uint32_t);

#define IREE_VMLA_SCATTER_OP(name, type)                                 \
  Status name(vm::ref<Buffer> src, iree_vmla_shape_t src_shape,          \
              vm::ref<Buffer> indices, iree_vmla_shape_t indices_shape,  \
              vm::ref<Buffer> dst, iree_vmla_shape_t dst_shape) {        \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                        \
    return kernels::Scatter::Execute<type>(                              \
        src->As<type>(), indices->As<int>(), dst->As<type>(), src_shape, \
        indices_shape, dst_shape);                                       \
  }
  IREE_VMLA_SCATTER_OP(ScatterX8, uint8_t);
  IREE_VMLA_SCATTER_OP(ScatterX16, uint16_t);
  IREE_VMLA_SCATTER_OP(ScatterX32, uint32_t);

#define IREE_VMLA_BROADCAST_OP(name, type)                        \
  Status name(vm::ref<Buffer> src, iree_vmla_shape_t src_shape,   \
              vm::ref<Buffer> dst, iree_vmla_shape_t dst_shape) { \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                 \
    return kernels::Broadcast::Execute<type>(src->As<type>(),     \
                                             dst->As<type>());    \
  }
  IREE_VMLA_BROADCAST_OP(BroadcastX8, uint8_t);
  IREE_VMLA_BROADCAST_OP(BroadcastX16, uint16_t);
  IREE_VMLA_BROADCAST_OP(BroadcastX32, uint32_t);

  IREE_VMLA_NONARY_OP(IotaI8, kernels::Iota, int8_t);
  IREE_VMLA_NONARY_OP(IotaI16, kernels::Iota, int16_t);
  IREE_VMLA_NONARY_OP(IotaI32, kernels::Iota, int32_t);
  IREE_VMLA_NONARY_OP(IotaF32, kernels::Iota, float_t);

#define IREE_VMLA_TILE_OP(name, type)                                     \
  Status name(vm::ref<Buffer> src, iree_vmla_shape_t src_shape,           \
              vm::ref<Buffer> dst, iree_vmla_shape_t dst_shape) {         \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                         \
    return kernels::Tile::Execute<type>(src->As<type>(), dst->As<type>(), \
                                        src_shape, dst_shape);            \
  }
  IREE_VMLA_TILE_OP(TileX8, uint8_t);
  IREE_VMLA_TILE_OP(TileX16, uint16_t);
  IREE_VMLA_TILE_OP(TileX32, uint32_t);

  //===--------------------------------------------------------------------===//
  // VMLA Ops: bit manipulation
  //===--------------------------------------------------------------------===//

  IREE_VMLA_UNARY_OP(NotX8, kernels::Not, uint8_t);
  IREE_VMLA_UNARY_OP(NotX16, kernels::Not, uint16_t);
  IREE_VMLA_UNARY_OP(NotX32, kernels::Not, uint32_t);
  IREE_VMLA_BINARY_OP(AndX8, kernels::And, uint8_t);
  IREE_VMLA_BINARY_OP(AndX16, kernels::And, uint16_t);
  IREE_VMLA_BINARY_OP(AndX32, kernels::And, uint32_t);
  IREE_VMLA_BINARY_OP(OrX8, kernels::Or, uint8_t);
  IREE_VMLA_BINARY_OP(OrX16, kernels::Or, uint16_t);
  IREE_VMLA_BINARY_OP(OrX32, kernels::Or, uint32_t);
  IREE_VMLA_BINARY_OP(XorX8, kernels::Xor, uint8_t);
  IREE_VMLA_BINARY_OP(XorX16, kernels::Xor, uint16_t);
  IREE_VMLA_BINARY_OP(XorX32, kernels::Xor, uint32_t);
  IREE_VMLA_BINARY_OP(ShlX8, kernels::ShiftLeft, uint8_t);
  IREE_VMLA_BINARY_OP(ShlX16, kernels::ShiftLeft, uint16_t);
  IREE_VMLA_BINARY_OP(ShlX32, kernels::ShiftLeft, uint32_t);
  IREE_VMLA_BINARY_OP(ShrU8, kernels::ShiftRight, uint8_t);
  IREE_VMLA_BINARY_OP(ShrU16, kernels::ShiftRight, uint16_t);
  IREE_VMLA_BINARY_OP(ShrU32, kernels::ShiftRight, uint32_t);
  IREE_VMLA_BINARY_OP(ShrI8, kernels::ShiftRight, int8_t);
  IREE_VMLA_BINARY_OP(ShrI16, kernels::ShiftRight, int16_t);
  IREE_VMLA_BINARY_OP(ShrI32, kernels::ShiftRight, int32_t);

  //===--------------------------------------------------------------------===//
  // VMLA Ops: arithmetic
  //===--------------------------------------------------------------------===//

  IREE_VMLA_BINARY_OP(AddI8, kernels::Add, int8_t);
  IREE_VMLA_BINARY_OP(AddI16, kernels::Add, int16_t);
  IREE_VMLA_BINARY_OP(AddI32, kernels::Add, int32_t);
  IREE_VMLA_BINARY_OP(AddF32, kernels::Add, float);
  IREE_VMLA_BINARY_OP(SubI8, kernels::Sub, int8_t);
  IREE_VMLA_BINARY_OP(SubI16, kernels::Sub, int16_t);
  IREE_VMLA_BINARY_OP(SubI32, kernels::Sub, int32_t);
  IREE_VMLA_BINARY_OP(SubF32, kernels::Sub, float);
  IREE_VMLA_UNARY_OP(AbsI8, kernels::Abs, int8_t);
  IREE_VMLA_UNARY_OP(AbsI16, kernels::Abs, int16_t);
  IREE_VMLA_UNARY_OP(AbsI32, kernels::Abs, int32_t);
  IREE_VMLA_UNARY_OP(AbsF32, kernels::Abs, float);
  IREE_VMLA_UNARY_OP(NegI8, kernels::Neg, int8_t);
  IREE_VMLA_UNARY_OP(NegI16, kernels::Neg, int16_t);
  IREE_VMLA_UNARY_OP(NegI32, kernels::Neg, int32_t);
  IREE_VMLA_UNARY_OP(NegF32, kernels::Neg, float);
  IREE_VMLA_BINARY_OP(MulI8, kernels::Mul, int8_t);
  IREE_VMLA_BINARY_OP(MulI16, kernels::Mul, int16_t);
  IREE_VMLA_BINARY_OP(MulI32, kernels::Mul, int32_t);
  IREE_VMLA_BINARY_OP(MulF32, kernels::Mul, float);
  IREE_VMLA_BINARY_OP(DivI8, kernels::Div, int8_t);
  IREE_VMLA_BINARY_OP(DivI16, kernels::Div, int16_t);
  IREE_VMLA_BINARY_OP(DivI32, kernels::Div, int32_t);
  IREE_VMLA_BINARY_OP(DivU8, kernels::Div, uint8_t);
  IREE_VMLA_BINARY_OP(DivU16, kernels::Div, uint16_t);
  IREE_VMLA_BINARY_OP(DivU32, kernels::Div, uint32_t);
  IREE_VMLA_BINARY_OP(DivF32, kernels::Div, float);
  IREE_VMLA_BINARY_OP(RemI8, kernels::Rem, int8_t);
  IREE_VMLA_BINARY_OP(RemI16, kernels::Rem, int16_t);
  IREE_VMLA_BINARY_OP(RemI32, kernels::Rem, int32_t);
  IREE_VMLA_BINARY_OP(RemU8, kernels::Rem, uint8_t);
  IREE_VMLA_BINARY_OP(RemU16, kernels::Rem, uint16_t);
  IREE_VMLA_BINARY_OP(RemU32, kernels::Rem, uint32_t);
  IREE_VMLA_BINARY_OP(RemF32, kernels::Rem, float);
  IREE_VMLA_BINARY_OP(PowF32, kernels::Pow, float);
  IREE_VMLA_UNARY_OP(ExpF32, kernels::Exp, float);
  IREE_VMLA_UNARY_OP(LogF32, kernels::Log, float);
  IREE_VMLA_UNARY_OP(RsqrtF32, kernels::Rsqrt, float);
  IREE_VMLA_UNARY_OP(SqrtF32, kernels::Sqrt, float);
  IREE_VMLA_UNARY_OP(CosF32, kernels::Cos, float);
  IREE_VMLA_UNARY_OP(SinF32, kernels::Sin, float);
  IREE_VMLA_UNARY_OP(TanhF32, kernels::Tanh, float);
  IREE_VMLA_BINARY_OP(Atan2F32, kernels::Atan2, float);

  IREE_VMLA_BINARY_OP(MinI8, kernels::Min, int8_t);
  IREE_VMLA_BINARY_OP(MinI16, kernels::Min, int16_t);
  IREE_VMLA_BINARY_OP(MinI32, kernels::Min, int32_t);
  IREE_VMLA_BINARY_OP(MinF32, kernels::Min, float);
  IREE_VMLA_BINARY_OP(MaxI8, kernels::Max, int8_t);
  IREE_VMLA_BINARY_OP(MaxI16, kernels::Max, int16_t);
  IREE_VMLA_BINARY_OP(MaxI32, kernels::Max, int32_t);
  IREE_VMLA_BINARY_OP(MaxF32, kernels::Max, float);
  IREE_VMLA_TERNARY_OP(ClampI8, kernels::Clamp, int8_t);
  IREE_VMLA_TERNARY_OP(ClampI16, kernels::Clamp, int16_t);
  IREE_VMLA_TERNARY_OP(ClampI32, kernels::Clamp, int32_t);
  IREE_VMLA_TERNARY_OP(ClampF32, kernels::Clamp, float);
  IREE_VMLA_UNARY_OP(FloorF32, kernels::Floor, float);
  IREE_VMLA_UNARY_OP(CeilF32, kernels::Ceil, float);

  //===--------------------------------------------------------------------===//
  // VMLA Ops: conversion
  //===--------------------------------------------------------------------===//

#define IREE_VMLA_CONVERSION_OP(name, src_type, dst_type)                      \
  Status name(vm::ref<Buffer> src, vm::ref<Buffer> dst) {                      \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                              \
    return kernels::Convert::Execute<src_type, dst_type>(src->As<src_type>(),  \
                                                         dst->As<dst_type>()); \
  }
  IREE_VMLA_CONVERSION_OP(ConvertI8I16, int8_t, int16_t);
  IREE_VMLA_CONVERSION_OP(ConvertI8I32, int8_t, int32_t);
  IREE_VMLA_CONVERSION_OP(ConvertI8F32, int8_t, float);
  IREE_VMLA_CONVERSION_OP(ConvertI16I8, int16_t, int8_t);
  IREE_VMLA_CONVERSION_OP(ConvertI16I32, int16_t, int32_t);
  IREE_VMLA_CONVERSION_OP(ConvertI16F32, int16_t, float);
  IREE_VMLA_CONVERSION_OP(ConvertI32I8, int32_t, int8_t);
  IREE_VMLA_CONVERSION_OP(ConvertI32I16, int32_t, int16_t);
  IREE_VMLA_CONVERSION_OP(ConvertI32F32, int32_t, float);
  IREE_VMLA_CONVERSION_OP(ConvertF32I8, float, int8_t);
  IREE_VMLA_CONVERSION_OP(ConvertF32I16, float, int16_t);
  IREE_VMLA_CONVERSION_OP(ConvertF32I32, float, int32_t);

  //===--------------------------------------------------------------------===//
  // VMLA Ops: Convolution
  //===--------------------------------------------------------------------===//

  Status ConvF32F32F32(vm::ref<Buffer> input, iree_vmla_shape_t input_shape,
                       vm::ref<Buffer> filter, iree_vmla_shape_t filter_shape,
                       vm::ref<Buffer> dst, iree_vmla_shape_t dst_shape,
                       absl::Span<const int32_t> window_strides,
                       absl::Span<const int32_t> padding,
                       absl::Span<const int32_t> lhs_dilation,
                       absl::Span<const int32_t> rhs_dilation,
                       const int32_t feature_group_count,
                       const int32_t batch_group_count) {
    IREE_TRACE_SCOPE0("VMLAModuleState::ConvF32F32F32");
    if (input_shape.size() != 4 || filter_shape.size() != 4 ||
        dst_shape.size() != 4) {
      return InvalidArgumentErrorBuilder(IREE_LOC)
             << "Expecting 4-d tensors for Conv2D kernel";
    }

    const int32_t batch_size = input_shape[0];
    const auto input_example_shape = input_shape.subspan(1, 3);
    const auto output_example_shape = dst_shape.subspan(1, 3);
    const auto filter_shape_4d = filter_shape.subspan(0, 4);
    const auto dilation = lhs_dilation.subspan(0, 2);
    const auto pad_h = padding.subspan(0, 2);
    const auto pad_w = padding.subspan(2, 2);
    const auto window_strides_2d = window_strides.subspan(0, 2);

    const float* raw_inputs_data = input->As<float>().data();
    const float* raw_filter_data = filter->As<float>().data();
    float* raw_dst_data = dst->As<float>().data();
    auto filter_buffer = absl::MakeConstSpan(
        raw_filter_data, kernels::GetElementCount(filter_shape_4d));

    const size_t input_stride = kernels::GetElementCount(input_example_shape);
    const size_t output_stride = kernels::GetElementCount(output_example_shape);

    for (int i = 0; i < batch_size; ++i) {
      auto input_example =
          absl::MakeConstSpan(raw_inputs_data + i * input_stride, input_stride);
      auto output_example =
          absl::MakeSpan(raw_dst_data + i * output_stride, output_stride);
      IREE_RETURN_IF_ERROR(kernels::Conv2D::Execute(
          input_example, input_example_shape, filter_buffer, filter_shape_4d,
          output_example, output_example_shape, window_strides_2d, pad_h, pad_w,
          dilation, feature_group_count));
    }
    return OkStatus();
  }

  //===--------------------------------------------------------------------===//
  // VMLA Ops: GEMM/GEMV
  //===--------------------------------------------------------------------===//

  Status BatchMatMulF32F32F32(vm::ref<Buffer> lhs, iree_vmla_shape_t lhs_shape,
                              vm::ref<Buffer> rhs, iree_vmla_shape_t rhs_shape,
                              vm::ref<Buffer> dst,
                              iree_vmla_shape_t dst_shape) {
    IREE_TRACE_SCOPE0("VMLAModuleState::BatchMatMulF32F32F32");
    // Compiler guarantees. Here for documentation purposes.
    assert(lhs_shape.size() == 3 && rhs_shape.size() == 3 &&
           dst_shape.size() == 3);
    assert(lhs_shape[0] == rhs_shape[0] && rhs_shape[0] == dst_shape[0]);

    iree_vmla_shape_t lhs_batch_element_shape = lhs_shape.subspan(1);
    iree_vmla_shape_t rhs_batch_element_shape = rhs_shape.subspan(1);
    iree_vmla_shape_t dst_batch_element_shape = dst_shape.subspan(1);
    auto lhs_batch_element_shape2 = lhs_batch_element_shape.subspan(0, 2);
    auto rhs_batch_element_shape2 = rhs_batch_element_shape.subspan(0, 2);
    auto dst_batch_element_shape2 = dst_batch_element_shape.subspan(0, 2);
    size_t lhs_batch_stride = kernels::GetElementCount(lhs_batch_element_shape);
    size_t rhs_batch_stride = kernels::GetElementCount(rhs_batch_element_shape);
    size_t dst_batch_stride = kernels::GetElementCount(dst_batch_element_shape);
    float* lhs_batch_base = lhs->As<float>().data();
    float* rhs_batch_base = rhs->As<float>().data();
    float* dst_batch_base = dst->As<float>().data();
    int32_t batch_dim = lhs_shape[0];
    for (int i = 0; i < batch_dim; i++) {
      kernels::MatMul::Buffers<float, float> buffers;
      buffers.lhs_buffer = absl::MakeSpan(lhs_batch_base + i * lhs_batch_stride,
                                          lhs_batch_stride);
      buffers.lhs_shape = lhs_batch_element_shape2;
      buffers.rhs_buffer = absl::MakeSpan(rhs_batch_base + i * rhs_batch_stride,
                                          rhs_batch_stride);
      buffers.rhs_shape = rhs_batch_element_shape2;
      buffers.dst_buffer = absl::MakeSpan(dst_batch_base + i * dst_batch_stride,
                                          dst_batch_stride);
      buffers.dst_shape = dst_batch_element_shape2;

      IREE_RETURN_IF_ERROR(kernels::MatMul::Execute(
          kernel_state_->mat_mul_state.get(), buffers));
    }
    return OkStatus();
  }

  //===--------------------------------------------------------------------===//
  // VMLA Ops: reduction
  //===--------------------------------------------------------------------===//

#define IREE_VMLA_REDUCTION_OP(name, kernel, type)                      \
  Status name(vm::ref<Buffer> src, iree_vmla_shape_t src_shape,         \
              vm::ref<Buffer> init, iree_vmla_shape_t init_shape,       \
              int32_t dimension, vm::ref<Buffer> dst,                   \
              iree_vmla_shape_t dst_shape) {                            \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                       \
    return kernel::Execute<type>(src->As<type>(), init->As<type>(),     \
                                 dst->As<type>(), dimension, src_shape, \
                                 dst_shape);                            \
  }
  IREE_VMLA_REDUCTION_OP(ReduceSumI8, kernels::ReduceSum, int8_t);
  IREE_VMLA_REDUCTION_OP(ReduceSumI16, kernels::ReduceSum, int16_t);
  IREE_VMLA_REDUCTION_OP(ReduceSumI32, kernels::ReduceSum, int32_t);
  IREE_VMLA_REDUCTION_OP(ReduceSumF32, kernels::ReduceSum, float);
  IREE_VMLA_REDUCTION_OP(ReduceMinI8, kernels::ReduceMin, int8_t);
  IREE_VMLA_REDUCTION_OP(ReduceMinI16, kernels::ReduceMin, int16_t);
  IREE_VMLA_REDUCTION_OP(ReduceMinI32, kernels::ReduceMin, int32_t);
  IREE_VMLA_REDUCTION_OP(ReduceMinF32, kernels::ReduceMin, float);
  IREE_VMLA_REDUCTION_OP(ReduceMaxI8, kernels::ReduceMax, int8_t);
  IREE_VMLA_REDUCTION_OP(ReduceMaxI16, kernels::ReduceMax, int16_t);
  IREE_VMLA_REDUCTION_OP(ReduceMaxI32, kernels::ReduceMax, int32_t);
  IREE_VMLA_REDUCTION_OP(ReduceMaxF32, kernels::ReduceMax, float);

#define IREE_VMLA_POOLING_OP(name, kernel, type)                              \
  Status name(vm::ref<Buffer> src, iree_vmla_shape_t src_shape,               \
              vm::ref<Buffer> init, iree_vmla_shape_t init_shape,             \
              vm::ref<Buffer> dst, iree_vmla_shape_t dst_shape,               \
              iree_vmla_shape_t window_dimensions, iree_vmla_shape_t strides, \
              iree_vmla_shape_t pad_low) {                                    \
    IREE_TRACE_SCOPE0("VMLAModuleState::" #name);                             \
    return kernel::Execute<type>(src->As<type>(), init->As<type>(),           \
                                 dst->As<type>(), src_shape, dst_shape,       \
                                 window_dimensions, strides, pad_low);        \
  }
  IREE_VMLA_POOLING_OP(PoolingSumI8, kernels::PoolingSum, int8_t);
  IREE_VMLA_POOLING_OP(PoolingSumI16, kernels::PoolingSum, int16_t);
  IREE_VMLA_POOLING_OP(PoolingSumI32, kernels::PoolingSum, int32_t);
  IREE_VMLA_POOLING_OP(PoolingSumF32, kernels::PoolingSum, float);
  IREE_VMLA_POOLING_OP(PoolingMinI8, kernels::PoolingMin, int8_t);
  IREE_VMLA_POOLING_OP(PoolingMinI16, kernels::PoolingMin, int16_t);
  IREE_VMLA_POOLING_OP(PoolingMinI32, kernels::PoolingMin, int32_t);
  IREE_VMLA_POOLING_OP(PoolingMinF32, kernels::PoolingMin, float);
  IREE_VMLA_POOLING_OP(PoolingMaxI8, kernels::PoolingMax, int8_t);
  IREE_VMLA_POOLING_OP(PoolingMaxI16, kernels::PoolingMax, int16_t);
  IREE_VMLA_POOLING_OP(PoolingMaxI32, kernels::PoolingMax, int32_t);
  IREE_VMLA_POOLING_OP(PoolingMaxF32, kernels::PoolingMax, float);

 private:
  iree_allocator_t allocator_;

  // NOTE: kernel state must be externally synchronized as it is shared across
  // all contexts using the VMLA module. This is fine in our current design as
  // we only ever execute a single context at a time but if we start to allow
  // concurrency across contexts we'll need to introduce locks.
  kernels::RuntimeState* kernel_state_ = nullptr;
};

//===----------------------------------------------------------------------===//
// VM module interface implementation
//===----------------------------------------------------------------------===//

static const vm::NativeFunction<VMLAModuleState> kVMLAModuleFunctions[] = {
    vm::MakeNativeFunction("interface.const", &VMLAModuleState::InterfaceConst),
    vm::MakeNativeFunction("interface.binding",
                           &VMLAModuleState::InterfaceBinding),

    vm::MakeNativeFunction("buffer.const", &VMLAModuleState::BufferConst),
    vm::MakeNativeFunction("buffer.alloc", &VMLAModuleState::BufferAlloc),
    vm::MakeNativeFunction("buffer.clone", &VMLAModuleState::BufferClone),
    vm::MakeNativeFunction("buffer.view", &VMLAModuleState::BufferView),
    vm::MakeNativeFunction("buffer.copy", &VMLAModuleState::BufferCopy),
    vm::MakeNativeFunction("buffer.fill", &VMLAModuleState::BufferFill),
    vm::MakeNativeFunction("buffer.load.i32", &VMLAModuleState::BufferLoadI32),

    vm::MakeNativeFunction("cmp.i8", &VMLAModuleState::CmpI8),
    vm::MakeNativeFunction("cmp.i16", &VMLAModuleState::CmpI16),
    vm::MakeNativeFunction("cmp.i32", &VMLAModuleState::CmpI32),
    vm::MakeNativeFunction("cmp.f32", &VMLAModuleState::CmpF32),
    vm::MakeNativeFunction("select.x8", &VMLAModuleState::SelectX8),
    vm::MakeNativeFunction("select.x16", &VMLAModuleState::SelectX16),
    vm::MakeNativeFunction("select.x32", &VMLAModuleState::SelectX32),

    vm::MakeNativeFunction("broadcast.x8", &VMLAModuleState::BroadcastX8),
    vm::MakeNativeFunction("broadcast.x16", &VMLAModuleState::BroadcastX16),
    vm::MakeNativeFunction("broadcast.x32", &VMLAModuleState::BroadcastX32),
    vm::MakeNativeFunction("copy.x8", &VMLAModuleState::CopyX8),
    vm::MakeNativeFunction("copy.x16", &VMLAModuleState::CopyX16),
    vm::MakeNativeFunction("copy.x32", &VMLAModuleState::CopyX32),
    vm::MakeNativeFunction("transpose.x8", &VMLAModuleState::TransposeX8),
    vm::MakeNativeFunction("transpose.x16", &VMLAModuleState::TransposeX16),
    vm::MakeNativeFunction("transpose.x32", &VMLAModuleState::TransposeX32),
    vm::MakeNativeFunction("reverse.x8", &VMLAModuleState::ReverseX8),
    vm::MakeNativeFunction("reverse.x16", &VMLAModuleState::ReverseX16),
    vm::MakeNativeFunction("reverse.x32", &VMLAModuleState::ReverseX32),
    vm::MakeNativeFunction("pad.x8", &VMLAModuleState::PadX8),
    vm::MakeNativeFunction("pad.x16", &VMLAModuleState::PadX16),
    vm::MakeNativeFunction("pad.x32", &VMLAModuleState::PadX32),
    vm::MakeNativeFunction("gather.x8", &VMLAModuleState::GatherX8),
    vm::MakeNativeFunction("gather.x16", &VMLAModuleState::GatherX16),
    vm::MakeNativeFunction("gather.x32", &VMLAModuleState::GatherX32),
    vm::MakeNativeFunction("scatter.x8", &VMLAModuleState::ScatterX8),
    vm::MakeNativeFunction("scatter.x16", &VMLAModuleState::ScatterX16),
    vm::MakeNativeFunction("scatter.x32", &VMLAModuleState::ScatterX32),
    vm::MakeNativeFunction("iota.i8", &VMLAModuleState::IotaI8),
    vm::MakeNativeFunction("iota.i16", &VMLAModuleState::IotaI16),
    vm::MakeNativeFunction("iota.i32", &VMLAModuleState::IotaI32),
    vm::MakeNativeFunction("iota.f32", &VMLAModuleState::IotaF32),
    vm::MakeNativeFunction("tile.x8", &VMLAModuleState::TileX8),
    vm::MakeNativeFunction("tile.x16", &VMLAModuleState::TileX16),
    vm::MakeNativeFunction("tile.x32", &VMLAModuleState::TileX32),

    vm::MakeNativeFunction("not.x8", &VMLAModuleState::NotX8),
    vm::MakeNativeFunction("not.x16", &VMLAModuleState::NotX16),
    vm::MakeNativeFunction("not.x32", &VMLAModuleState::NotX32),
    vm::MakeNativeFunction("and.x8", &VMLAModuleState::AndX8),
    vm::MakeNativeFunction("and.x16", &VMLAModuleState::AndX16),
    vm::MakeNativeFunction("and.x32", &VMLAModuleState::AndX32),
    vm::MakeNativeFunction("or.x8", &VMLAModuleState::OrX8),
    vm::MakeNativeFunction("or.x16", &VMLAModuleState::OrX16),
    vm::MakeNativeFunction("or.x32", &VMLAModuleState::OrX32),
    vm::MakeNativeFunction("xor.x8", &VMLAModuleState::XorX8),
    vm::MakeNativeFunction("xor.x16", &VMLAModuleState::XorX16),
    vm::MakeNativeFunction("xor.x32", &VMLAModuleState::XorX32),
    vm::MakeNativeFunction("shl.x8", &VMLAModuleState::ShlX8),
    vm::MakeNativeFunction("shl.x16", &VMLAModuleState::ShlX16),
    vm::MakeNativeFunction("shl.x32", &VMLAModuleState::ShlX32),
    vm::MakeNativeFunction("shr.u8", &VMLAModuleState::ShrU8),
    vm::MakeNativeFunction("shr.u16", &VMLAModuleState::ShrU16),
    vm::MakeNativeFunction("shr.u32", &VMLAModuleState::ShrU32),
    vm::MakeNativeFunction("shr.i8", &VMLAModuleState::ShrI8),
    vm::MakeNativeFunction("shr.i16", &VMLAModuleState::ShrI16),
    vm::MakeNativeFunction("shr.i32", &VMLAModuleState::ShrI32),

    vm::MakeNativeFunction("add.i8", &VMLAModuleState::AddI8),
    vm::MakeNativeFunction("add.i16", &VMLAModuleState::AddI16),
    vm::MakeNativeFunction("add.i32", &VMLAModuleState::AddI32),
    vm::MakeNativeFunction("add.f32", &VMLAModuleState::AddF32),
    vm::MakeNativeFunction("sub.i8", &VMLAModuleState::SubI8),
    vm::MakeNativeFunction("sub.i16", &VMLAModuleState::SubI16),
    vm::MakeNativeFunction("sub.i32", &VMLAModuleState::SubI32),
    vm::MakeNativeFunction("sub.f32", &VMLAModuleState::SubF32),
    vm::MakeNativeFunction("abs.i8", &VMLAModuleState::AbsI8),
    vm::MakeNativeFunction("abs.i16", &VMLAModuleState::AbsI16),
    vm::MakeNativeFunction("abs.i32", &VMLAModuleState::AbsI32),
    vm::MakeNativeFunction("abs.f32", &VMLAModuleState::AbsF32),
    vm::MakeNativeFunction("neg.i8", &VMLAModuleState::NegI8),
    vm::MakeNativeFunction("neg.i16", &VMLAModuleState::NegI16),
    vm::MakeNativeFunction("neg.i32", &VMLAModuleState::NegI32),
    vm::MakeNativeFunction("neg.f32", &VMLAModuleState::NegF32),
    vm::MakeNativeFunction("mul.i8", &VMLAModuleState::MulI8),
    vm::MakeNativeFunction("mul.i16", &VMLAModuleState::MulI16),
    vm::MakeNativeFunction("mul.i32", &VMLAModuleState::MulI32),
    vm::MakeNativeFunction("mul.f32", &VMLAModuleState::MulF32),
    vm::MakeNativeFunction("div.i8", &VMLAModuleState::DivI8),
    vm::MakeNativeFunction("div.i16", &VMLAModuleState::DivI16),
    vm::MakeNativeFunction("div.i32", &VMLAModuleState::DivI32),
    vm::MakeNativeFunction("div.u8", &VMLAModuleState::DivU8),
    vm::MakeNativeFunction("div.u16", &VMLAModuleState::DivU16),
    vm::MakeNativeFunction("div.u32", &VMLAModuleState::DivU32),
    vm::MakeNativeFunction("div.f32", &VMLAModuleState::DivF32),
    vm::MakeNativeFunction("rem.i8", &VMLAModuleState::RemI8),
    vm::MakeNativeFunction("rem.i16", &VMLAModuleState::RemI16),
    vm::MakeNativeFunction("rem.i32", &VMLAModuleState::RemI32),
    vm::MakeNativeFunction("rem.u8", &VMLAModuleState::RemU8),
    vm::MakeNativeFunction("rem.u16", &VMLAModuleState::RemU16),
    vm::MakeNativeFunction("rem.u32", &VMLAModuleState::RemU32),
    vm::MakeNativeFunction("rem.f32", &VMLAModuleState::RemF32),
    vm::MakeNativeFunction("pow.f32", &VMLAModuleState::PowF32),
    vm::MakeNativeFunction("exp.f32", &VMLAModuleState::ExpF32),
    vm::MakeNativeFunction("log.f32", &VMLAModuleState::LogF32),
    vm::MakeNativeFunction("rsqrt.f32", &VMLAModuleState::RsqrtF32),
    vm::MakeNativeFunction("sqrt.f32", &VMLAModuleState::SqrtF32),
    vm::MakeNativeFunction("cos.f32", &VMLAModuleState::CosF32),
    vm::MakeNativeFunction("sin.f32", &VMLAModuleState::SinF32),
    vm::MakeNativeFunction("tanh.f32", &VMLAModuleState::TanhF32),
    vm::MakeNativeFunction("atan2.f32", &VMLAModuleState::Atan2F32),

    vm::MakeNativeFunction("min.i8", &VMLAModuleState::MinI8),
    vm::MakeNativeFunction("min.i16", &VMLAModuleState::MinI16),
    vm::MakeNativeFunction("min.i32", &VMLAModuleState::MinI32),
    vm::MakeNativeFunction("min.f32", &VMLAModuleState::MinF32),
    vm::MakeNativeFunction("max.i8", &VMLAModuleState::MaxI8),
    vm::MakeNativeFunction("max.i16", &VMLAModuleState::MaxI16),
    vm::MakeNativeFunction("max.i32", &VMLAModuleState::MaxI32),
    vm::MakeNativeFunction("max.f32", &VMLAModuleState::MaxF32),
    vm::MakeNativeFunction("clamp.i8", &VMLAModuleState::ClampI8),
    vm::MakeNativeFunction("clamp.i16", &VMLAModuleState::ClampI16),
    vm::MakeNativeFunction("clamp.i32", &VMLAModuleState::ClampI32),
    vm::MakeNativeFunction("clamp.f32", &VMLAModuleState::ClampF32),
    vm::MakeNativeFunction("floor.f32", &VMLAModuleState::FloorF32),
    vm::MakeNativeFunction("ceil.f32", &VMLAModuleState::CeilF32),

    vm::MakeNativeFunction("convert.i8.i16", &VMLAModuleState::ConvertI8I16),
    vm::MakeNativeFunction("convert.i8.i32", &VMLAModuleState::ConvertI8I32),
    vm::MakeNativeFunction("convert.i8.f32", &VMLAModuleState::ConvertI8F32),
    vm::MakeNativeFunction("convert.i16.i8", &VMLAModuleState::ConvertI16I8),
    vm::MakeNativeFunction("convert.i16.i32", &VMLAModuleState::ConvertI16I32),
    vm::MakeNativeFunction("convert.i16.f32", &VMLAModuleState::ConvertI16F32),
    vm::MakeNativeFunction("convert.i32.i8", &VMLAModuleState::ConvertI32I8),
    vm::MakeNativeFunction("convert.i32.i16", &VMLAModuleState::ConvertI32I16),
    vm::MakeNativeFunction("convert.i32.f32", &VMLAModuleState::ConvertI32F32),
    vm::MakeNativeFunction("convert.f32.i8", &VMLAModuleState::ConvertF32I8),
    vm::MakeNativeFunction("convert.f32.i16", &VMLAModuleState::ConvertF32I16),
    vm::MakeNativeFunction("convert.f32.i32", &VMLAModuleState::ConvertF32I32),

    vm::MakeNativeFunction("reduce.sum.i8", &VMLAModuleState::ReduceSumI8),
    vm::MakeNativeFunction("reduce.sum.i16", &VMLAModuleState::ReduceSumI16),
    vm::MakeNativeFunction("reduce.sum.i32", &VMLAModuleState::ReduceSumI32),
    vm::MakeNativeFunction("reduce.sum.f32", &VMLAModuleState::ReduceSumF32),
    vm::MakeNativeFunction("reduce.min.i8", &VMLAModuleState::ReduceMinI8),
    vm::MakeNativeFunction("reduce.min.i16", &VMLAModuleState::ReduceMinI16),
    vm::MakeNativeFunction("reduce.min.i32", &VMLAModuleState::ReduceMinI32),
    vm::MakeNativeFunction("reduce.min.f32", &VMLAModuleState::ReduceMinF32),
    vm::MakeNativeFunction("reduce.max.i8", &VMLAModuleState::ReduceMaxI8),
    vm::MakeNativeFunction("reduce.max.i16", &VMLAModuleState::ReduceMaxI16),
    vm::MakeNativeFunction("reduce.max.i32", &VMLAModuleState::ReduceMaxI32),
    vm::MakeNativeFunction("reduce.max.f32", &VMLAModuleState::ReduceMaxF32),

    vm::MakeNativeFunction("pooling.sum.i8", &VMLAModuleState::PoolingSumI8),
    vm::MakeNativeFunction("pooling.sum.i16", &VMLAModuleState::PoolingSumI16),
    vm::MakeNativeFunction("pooling.sum.i32", &VMLAModuleState::PoolingSumI32),
    vm::MakeNativeFunction("pooling.sum.f32", &VMLAModuleState::PoolingSumF32),
    vm::MakeNativeFunction("pooling.min.i8", &VMLAModuleState::PoolingMinI8),
    vm::MakeNativeFunction("pooling.min.i16", &VMLAModuleState::PoolingMinI16),
    vm::MakeNativeFunction("pooling.min.i32", &VMLAModuleState::PoolingMinI32),
    vm::MakeNativeFunction("pooling.min.f32", &VMLAModuleState::PoolingMinF32),
    vm::MakeNativeFunction("pooling.max.i8", &VMLAModuleState::PoolingMaxI8),
    vm::MakeNativeFunction("pooling.max.i16", &VMLAModuleState::PoolingMaxI16),
    vm::MakeNativeFunction("pooling.max.i32", &VMLAModuleState::PoolingMaxI32),
    vm::MakeNativeFunction("pooling.max.f32", &VMLAModuleState::PoolingMaxF32),

    vm::MakeNativeFunction("batch.matmul.f32f32.f32",
                           &VMLAModuleState::BatchMatMulF32F32F32),

    vm::MakeNativeFunction("conv.f32f32.f32", &VMLAModuleState::ConvF32F32F32)};

// Per-device VMLA module.
// One of these will be created per device and be shared across all executables
// that are created within that device. Large shared kernel state can go here
// (such as thread pools/caches/etc), though note that they must be either
// thread-safe or internally synchronized.
//
// Thread-safe.
class VMLAModule final : public vm::NativeModule<VMLAModuleState> {
 public:
  explicit VMLAModule(iree_allocator_t allocator)
      : vm::NativeModule<VMLAModuleState>(
            "vmla", allocator, absl::MakeConstSpan(kVMLAModuleFunctions)) {}
  ~VMLAModule() = default;

  Status Initialize() {
    IREE_TRACE_SCOPE0("VMLAModule::Initialize");
    return OkStatus();
  }

  StatusOr<std::unique_ptr<VMLAModuleState>> CreateState(
      iree_allocator_t allocator) override {
    IREE_TRACE_SCOPE0("VMLAModule::CreateState");
    auto state = std::make_unique<VMLAModuleState>(allocator, &kernel_state_);
    return state;
  }

 private:
  // NOTE: shared across all contexts with the VMLA module loaded. See
  // VMLAModuleState::kernel_state_ for more information.
  kernels::RuntimeState kernel_state_;
};

}  // namespace

Status ModuleCreate(iree_allocator_t allocator, iree_vm_module_t** out_module) {
  if (!out_module) {
    return InvalidArgumentErrorBuilder(IREE_LOC)
           << "out_module must not be null";
  }
  *out_module = nullptr;
  auto module = std::make_unique<VMLAModule>(allocator);
  IREE_RETURN_IF_ERROR(module->Initialize());
  *out_module = module.release()->interface();
  return OkStatus();
}

}  // namespace vmla
}  // namespace hal
}  // namespace iree