/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

/*!
 * \file ir_builder.cc
 * \brief IRBuilder for SPIRV block
 */
#include "ir_builder.h"

#include <spirv.hpp>

namespace tvm {
namespace codegen {
namespace spirv {

namespace {

PrimType AsPrimType(const SType& type) { return type.type.as_or_throw<PrimType>(); }

}  // namespace

// implementations

IRBuilder::IRBuilder(const SPIRVSupport& support) : spirv_support_(support) {}

void IRBuilder::InitHeader() {
  TVM_FFI_ICHECK_EQ(header_.size(), 0U);
  header_.push_back(spv::MagicNumber);

  // Target SPIR-V version 1.0.  Additional functionality will be
  // enabled through extensions.
  header_.push_back(0x10000);

  // generator: set to 0, unknown
  header_.push_back(0U);
  // Bound: set during Finalize
  header_.push_back(0U);
  // Schema: reserved
  header_.push_back(0U);

  // Declare CapabilityShader by default.  All other capabilities are
  // determined by the types declared.
  capabilities_used_.insert(spv::CapabilityShader);

#ifdef TVM_SPIRV_KHR_INTEGER_DOT_PRODUCT
  if (spirv_support_.supports_integer_dot_product) {
    capabilities_used_.insert(spv::CapabilityDotProductKHR);
    capabilities_used_.insert(spv::CapabilityDotProductInput4x8BitPackedKHR);
    extensions_used_.insert("SPV_KHR_integer_dot_product");
  }
#endif

  if (spirv_support_.supports_cooperative_matrix) {
    capabilities_used_.insert(spv::CapabilityCooperativeMatrixNV);
    extensions_used_.insert("SPV_NV_cooperative_matrix");
  }

  // memory model
  ib_.Begin(spv::OpMemoryModel)
      .AddSeq(spv::AddressingModelLogical, spv::MemoryModelGLSL450)
      .Commit(&entry_);
  this->InitPreDefs();
}

void IRBuilder::InitPreDefs() {
  ext_glsl450_ = ExtInstImport("GLSL.std.450");
  t_int32_ = DeclareType(PrimType::Int(32));
  t_uint32_ = DeclareType(PrimType::UInt(32));
  t_bool_ = DeclareType(PrimType::Bool());
  t_fp32_ = DeclareType(PrimType::Float(32));
  const_i32_zero_ = IntImm(t_int32_, 0);

  // declare void, and void functions
  t_void_.id = id_counter_++;
  ib_.Begin(spv::OpTypeVoid).Add(t_void_).Commit(&global_);
  t_void_func_.id = id_counter_++;
  ib_.Begin(spv::OpTypeFunction).AddSeq(t_void_func_, t_void_).Commit(&global_);
}

std::vector<uint32_t> IRBuilder::Finalize() {
  std::vector<uint32_t> data;
  // Index for upper bound of id numbers.
  const int kBoundLoc = 3;
  header_[kBoundLoc] = id_counter_;
  data.insert(data.end(), header_.begin(), header_.end());
  for (const auto& capability : capabilities_used_) {
    ib_.Begin(spv::OpCapability).Add(capability).Commit(&data);
  }
  for (const auto& ext_name : extensions_used_) {
    ib_.Begin(spv::OpExtension).Add(ext_name).Commit(&data);
  }
  data.insert(data.end(), extended_instruction_section_.begin(),
              extended_instruction_section_.end());
  data.insert(data.end(), entry_.begin(), entry_.end());
  data.insert(data.end(), exec_mode_.begin(), exec_mode_.end());
  data.insert(data.end(), debug_.begin(), debug_.end());
  data.insert(data.end(), decorate_.begin(), decorate_.end());
  data.insert(data.end(), global_.begin(), global_.end());
  data.insert(data.end(), func_header_.begin(), func_header_.end());
  data.insert(data.end(), function_scope_vars_.begin(), function_scope_vars_.end());
  data.insert(data.end(), function_.begin(), function_.end());
  return data;
}

SType IRBuilder::GetSType(const PrimType& dtype, uint32_t row, uint32_t col) {
  if (dtype == PrimType::Int(32)) {
    return t_int32_;
  } else if (dtype == PrimType::Bool()) {
    return t_bool_;
  } else if (dtype == PrimType::Float(32)) {
    return t_fp32_;
  } else if (dtype == PrimType::UInt(32)) {
    return t_uint32_;
  }
  uint64_t type_key;
  type_key = static_cast<uint32_t>(dtype.code());
  type_key |= static_cast<uint32_t>(dtype.bits()) << 8U;
  if (row * col == 0) {
    TVM_FFI_ICHECK((row == 0) && (col == 0));
    type_key |= static_cast<uint32_t>(dtype.lanes()) << 16U;
  } else {
    type_key |= static_cast<uint64_t>(row) << 32U;
    type_key |= static_cast<uint64_t>(col) << 40U;
  }

  auto it = pod_type_tbl_.find(type_key);
  if (it != pod_type_tbl_.end()) {
    return it->second;
  }
  SType t = DeclareType(dtype, row, col);
  pod_type_tbl_[type_key] = t;
  return t;
}

SType IRBuilder::GetPointerType(const SType& value_type, spv::StorageClass storage_class) {
  TVM_FFI_ICHECK_NE(storage_class, spv::StorageClassMax);
  auto key = std::make_pair(value_type.id, storage_class);
  auto it = pointer_type_tbl_.find(key);
  if (it != pointer_type_tbl_.end()) {
    return it->second;
  }
  SType t;
  t.id = id_counter_++;
  t.type = PointerType::VoidPointerTy();
  t.element_type_id = value_type.id;
  t.storage_class = storage_class;
  ib_.Begin(spv::OpTypePointer).AddSeq(t, storage_class, value_type).Commit(&global_);
  pointer_type_tbl_[key] = t;
  return t;
}

SType IRBuilder::GetStructArrayType(const SType& value_type, uint32_t num_elems,
                                    bool interface_block) {
  auto key = std::make_tuple(value_type.id, num_elems, interface_block);
  auto it = struct_array_type_tbl_.find(key);
  if (it != struct_array_type_tbl_.end()) {
    return it->second;
  }

  SType arr_type;
  arr_type.id = id_counter_++;
  arr_type.type = PointerType::VoidPointerTy();
  arr_type.element_type_id = value_type.id;

  if (num_elems != 0) {
    Value length = UIntImm(GetSType(PrimType::UInt(32)), num_elems);
    ib_.Begin(spv::OpTypeArray).AddSeq(arr_type, value_type, length).Commit(&global_);
  } else {
    ib_.Begin(spv::OpTypeRuntimeArray).AddSeq(arr_type, value_type).Commit(&global_);
  }
  if (interface_block) {
    PrimType primitive_type = AsPrimType(value_type);
    int nbits = primitive_type.bits() * primitive_type.lanes();
    TVM_FFI_ICHECK_EQ(nbits % 8, 0);
    uint32_t nbytes = static_cast<uint32_t>(nbits) / 8;
    // Explicit layout is required for descriptor-backed interface blocks.
    this->Decorate(spv::OpDecorate, arr_type, spv::DecorationArrayStride, nbytes);
  }
  // declare struct of array
  SType struct_type;
  struct_type.id = id_counter_++;
  struct_type.type = PointerType::VoidPointerTy();
  struct_type.element_type_id = value_type.id;
  ib_.Begin(spv::OpTypeStruct).AddSeq(struct_type, arr_type).Commit(&global_);

  if (interface_block) {
    ib_.Begin(spv::OpMemberDecorate)
        .AddSeq(struct_type, 0, spv::DecorationOffset, 0)
        .Commit(&decorate_);
    // Runtime array are always decorated as Block or BufferBlock
    // (shader storage buffer)
    if (spirv_support_.supports_storage_buffer_storage_class) {
      // If SPIRV 1.3+, or with extension
      // SPV_KHR_storage_buffer_storage_class, BufferBlock is
      // deprecated.
      extensions_used_.insert("SPV_KHR_storage_buffer_storage_class");
      this->Decorate(spv::OpDecorate, struct_type, spv::DecorationBlock);
    } else {
      if (num_elems == 0) {
        this->Decorate(spv::OpDecorate, struct_type, spv::DecorationBufferBlock);
      }
    }
  }
  struct_array_type_tbl_[key] = struct_type;
  return struct_type;
}

Value IRBuilder::StructArrayAccess(const SType& res_type, Value buffer, Value index) {
  TVM_FFI_ICHECK(buffer.flag == kStructArrayPtr);
  return MakeValue(spv::OpInBoundsAccessChain, res_type, buffer, const_i32_zero_, index);
}

Value IRBuilder::IntImm(const SType& dtype, int64_t value) {
  return GetConst_(dtype, reinterpret_cast<uint64_t*>(&value));
}

Value IRBuilder::UIntImm(const SType& dtype, uint64_t value) { return GetConst_(dtype, &value); }

Value IRBuilder::FloatImm(const SType& dtype, double value) {
  PrimType primitive_type = AsPrimType(dtype);
  if (primitive_type.bits() == 64) {
    return GetConst_(dtype, reinterpret_cast<uint64_t*>(&value));
  } else if (primitive_type.bits() == 32) {
    float fvalue = static_cast<float>(value);
    uint32_t* ptr = reinterpret_cast<uint32_t*>(&fvalue);
    uint64_t data = ptr[0];
    return GetConst_(dtype, &data);
  } else {
    TVM_FFI_ICHECK_EQ(primitive_type.bits(), 16);
    float fvalue = static_cast<float>(value);
    uint32_t* ptr = reinterpret_cast<uint32_t*>(&fvalue);
    uint64_t data = ptr[0];
    if (data == 0)
      return GetConst_(dtype, &data);
    else
      return Cast(dtype, FloatImm(GetSType(PrimType::Float(32)), value));
  }
}

Value IRBuilder::BufferArgument(const SType& value_type, uint32_t descriptor_set,
                                uint32_t binding) {
  // If SPIRV 1.3+, or with extension SPV_KHR_storage_buffer_storage_class, BufferBlock is
  // deprecated.
  spv::StorageClass storage_class;
  if (spirv_support_.supports_storage_buffer_storage_class) {
    storage_class = spv::StorageClassStorageBuffer;
  } else {
    storage_class = spv::StorageClassUniform;
  }

  SType sarr_type = GetStructArrayType(value_type, 0, true);
  SType ptr_type = GetPointerType(sarr_type, storage_class);
  Value val = NewValue(ptr_type, kStructArrayPtr);

  ib_.Begin(spv::OpVariable).AddSeq(ptr_type, val, storage_class).Commit(&global_);

  this->DecorateBufferArgument(val, descriptor_set, binding);
  return val;
}

Value IRBuilder::DeclareStorageVariable(const std::vector<SType>& value_types,
                                        spv::StorageClass storage_class, ValueKind kind) {
  SType struct_type;
  struct_type.id = id_counter_++;
  struct_type.type = PointerType::VoidPointerTy();
  ib_.Begin(spv::OpTypeStruct).Add(struct_type);
  for (const SType& vtype : value_types) {
    ib_.Add(vtype);
  }
  ib_.Commit(&global_);

  uint32_t offset = 0;
  for (uint32_t i = 0; i < value_types.size(); ++i) {
    ib_.Begin(spv::OpMemberDecorate)
        .AddSeq(struct_type, i, spv::DecorationOffset, offset)
        .Commit(&decorate_);
    PrimType t = AsPrimType(value_types[i]);
    uint32_t nbits = t.bits() * t.lanes();
    TVM_FFI_ICHECK_EQ(nbits % 8, 0);
    uint32_t bytes = (nbits / 8);
    if (t.bits() == 32) {
      // In our Vulkan runtime, each scalar argument always occupies 64 bit.
      offset += bytes * 2;
    } else {
      TVM_FFI_ICHECK_EQ(t.bits(), 64);
      offset += bytes;
    }
  }
  this->Decorate(spv::OpDecorate, struct_type, spv::DecorationBlock);

  SType ptr_type = GetPointerType(struct_type, storage_class);
  Value val = NewValue(ptr_type, kind);
  ib_.Begin(spv::OpVariable).AddSeq(ptr_type, val, storage_class).Commit(&global_);
  return val;
}

Value IRBuilder::DeclarePushConstant(const std::vector<SType>& value_types) {
  TVM_FFI_ICHECK_EQ(push_const_.id, 0);
  return DeclareStorageVariable(value_types, spv::StorageClassPushConstant, kPushConstantPtr);
}

Value IRBuilder::GetPushConstant(Value ptr_push_const, const SType& v_type, uint32_t index) {
  SType ptr_vtype = this->GetPointerType(v_type, spv::StorageClassPushConstant);
  Value ptr = this->MakeValue(spv::OpAccessChain, ptr_vtype, ptr_push_const,
                              IntImm(t_int32_, static_cast<int64_t>(index)));
  return this->MakeValue(spv::OpLoad, v_type, ptr);
}

Value IRBuilder::DeclareUniformBuffer(const std::vector<SType>& value_types,
                                      uint32_t descriptor_set, uint32_t binding) {
  Value val = DeclareStorageVariable(value_types, spv::StorageClassUniform, kUniformPtr);
  this->DecorateBufferArgument(val, descriptor_set, binding);
  return val;
}

void IRBuilder::DecorateBufferArgument(Value val, uint32_t descriptor_set, uint32_t binding) {
  this->Decorate(spv::OpDecorate, val, spv::DecorationDescriptorSet, descriptor_set);
  this->Decorate(spv::OpDecorate, val, spv::DecorationBinding, binding);
}

Value IRBuilder::GetUniform(Value ptr_push_const, const SType& v_type, uint32_t index) {
  SType ptr_vtype = this->GetPointerType(v_type, spv::StorageClassUniform);
  Value ptr = this->MakeValue(spv::OpAccessChain, ptr_vtype, ptr_push_const,
                              IntImm(t_int32_, static_cast<int64_t>(index)));
  return this->MakeValue(spv::OpLoad, v_type, ptr);
}

Value IRBuilder::NewFunction() { return NewValue(t_void_func_, kFunction); }

void IRBuilder::CommitKernelFunction(const Value& func, const std::string& name) {
  TVM_FFI_ICHECK_EQ(func.flag, kFunction);
  ib_.Begin(spv::OpEntryPoint).AddSeq(spv::ExecutionModelGLCompute, func, name);
  for (auto& it : built_in_tbl_) {
    ib_.Add(it.second);
  }
  ib_.Commit(&entry_);
}

void IRBuilder::StartFunction(const Value& func) {
  TVM_FFI_ICHECK_EQ(func.flag, kFunction);
  // add function declaration to the header.
  ib_.Begin(spv::OpFunction).AddSeq(t_void_, func, 0, t_void_func_).Commit(&func_header_);

  spirv::Label start_label = this->NewLabel();
  ib_.Begin(spv::OpLabel).AddSeq(start_label).Commit(&func_header_);
  curr_label_ = start_label;
}

void IRBuilder::SetLocalSize(const Value& func, uint32_t local_size[3]) {
  TVM_FFI_ICHECK_EQ(func.flag, kFunction);
  ib_.Begin(spv::OpExecutionMode)
      .AddSeq(func, spv::ExecutionModeLocalSize, local_size[0], local_size[1], local_size[2])
      .Commit(&exec_mode_);
}

Value IRBuilder::Allocate(const SType& value_type, uint32_t num_elems,
                          spv::StorageClass storage_class) {
  TVM_FFI_ICHECK_NE(num_elems, 0U);
  SType sarr_type = GetStructArrayType(value_type, num_elems, false);
  SType ptr_type = GetPointerType(sarr_type, storage_class);
  Value val = NewValue(ptr_type, kStructArrayPtr);
  if (storage_class == spv::StorageClassFunction) {
    ib_.Begin(spv::OpVariable).AddSeq(ptr_type, val, storage_class).Commit(&func_header_);
  } else {
    ib_.Begin(spv::OpVariable).AddSeq(ptr_type, val, storage_class).Commit(&global_);
  }
  return val;
}

Value IRBuilder::GetWorkgroupID(uint32_t dim_index) {
  std::string name = "blockIdx." + std::string(1, 'x' + dim_index);
  return GetBuiltInValue(spv::BuiltInWorkgroupId, dim_index, name);
}

Value IRBuilder::GetLocalID(uint32_t dim_index) {
  std::string name = "threadIdx." + std::string(1, 'x' + dim_index);
  return GetBuiltInValue(spv::BuiltInLocalInvocationId, dim_index, name);
}

Value IRBuilder::GetBuiltInValue(spv::BuiltIn built_in, uint32_t index, const std::string& name) {
  // Returned cached value if it exists
  {
    auto it = built_in_values_tbl_.find({built_in, index});
    if (it != built_in_values_tbl_.end()) {
      return it->second;
    }
  }

  PrimType data_type = PrimType::Int(32);
  PrimType global_arr_type = data_type.WithLanes(3);
  switch (built_in) {
    case spv::BuiltInLocalInvocationId:
    case spv::BuiltInWorkgroupId:
      break;

    default:
      TVM_FFI_THROW(InternalError) << "No data type defined for SPIR-V Built-In " << built_in;
  }

  // Look up the decorated array value at global scope.  If it doesn't
  // exist already, declare it.
  Value global_array;
  {
    auto it = built_in_tbl_.find(built_in);
    if (it != built_in_tbl_.end()) {
      global_array = it->second;
    } else {
      SType ptr_arr_type = this->GetPointerType(GetSType(global_arr_type), spv::StorageClassInput);
      global_array = NewValue(ptr_arr_type, kVectorPtr);

      ib_.Begin(spv::OpVariable)
          .AddSeq(ptr_arr_type, global_array, spv::StorageClassInput)
          .Commit(&global_);
      this->Decorate(spv::OpDecorate, global_array, spv::DecorationBuiltIn, built_in);

      switch (built_in) {
        case spv::BuiltInLocalInvocationId:
          SetName(global_array, "BuiltInLocalInvocationId");
          break;
        case spv::BuiltInWorkgroupId:
          SetName(global_array, "BuiltInWorkgroupId");
          break;

        default:
          break;
      }

      built_in_tbl_[built_in] = global_array;
    }
  }

  // Declare the dereferenced value
  SType data_stype = GetSType(data_type);
  SType ptr_type = this->GetPointerType(data_stype, spv::StorageClassInput);
  Value global_const_index = UIntImm(t_int32_, static_cast<int64_t>(index));

  Value ptr = NewValue(ptr_type, kNormal);
  ib_.Begin(spv::OpAccessChain)
      .AddSeq(ptr_type, ptr, global_array, global_const_index)
      .Commit(&function_scope_vars_);

  Value output = NewValue(data_stype, kNormal);
  ib_.Begin(spv::OpLoad).AddSeq(data_stype, output, ptr).Commit(&function_scope_vars_);
  if (name.size()) {
    SetName(output, name);
  }

  // Store to cache and return
  built_in_values_tbl_[{built_in, index}] = output;
  return output;
}

Value IRBuilder::GetConst_(const SType& dtype, const uint64_t* pvalue) {
  PrimType primitive_type = AsPrimType(dtype);
  auto key = std::make_pair(dtype.id, pvalue[0]);
  auto it = const_tbl_.find(key);
  if (it != const_tbl_.end()) {
    return it->second;
  }
  TVM_FFI_ICHECK_LE(primitive_type.bits(), 64);
  Value ret = NewValue(dtype, kConstant);
  if (primitive_type == PrimType::Bool()) {
    // bool types.
    if (*pvalue) {
      ib_.Begin(spv::OpConstantTrue).AddSeq(dtype, ret);
    } else {
      ib_.Begin(spv::OpConstantFalse).AddSeq(dtype, ret);
    }
  } else {
    // Integral/floating-point types.
    ib_.Begin(spv::OpConstant).AddSeq(dtype, ret);
    uint64_t mask = 0xFFFFFFFFUL;
    ib_.Add(static_cast<uint32_t>(pvalue[0] & mask));
    if (primitive_type.bits() > 32) {
      if (primitive_type.MatchesCode(DLDataTypeCode::kDLInt)) {
        int64_t sign_mask = 0xFFFFFFFFL;
        const int64_t* sign_ptr = reinterpret_cast<const int64_t*>(pvalue);
        ib_.Add(static_cast<uint32_t>((sign_ptr[0] >> 32L) & sign_mask));
      } else {
        ib_.Add(static_cast<uint32_t>((pvalue[0] >> 32UL) & mask));
      }
    }
  }
  ib_.Commit(&global_);
  const_tbl_[key] = ret;
  return ret;
}

SType IRBuilder::DeclareType(const PrimType& dtype, uint32_t row, uint32_t col) {
  AddCapabilityFor(dtype);

  if (dtype.lanes() == 1) {
    SType t;
    t.id = id_counter_++;
    t.type = dtype;
    if (dtype.MatchesCode(DLDataTypeCode::kDLBool)) {
      ib_.Begin(spv::OpTypeBool).Add(t).Commit(&global_);
    } else if (dtype.MatchesCode(DLDataTypeCode::kDLInt)) {
      ib_.Begin(spv::OpTypeInt).AddSeq(t, dtype.bits(), 1).Commit(&global_);
    } else if (dtype.MatchesCode(DLDataTypeCode::kDLUInt)) {
      ib_.Begin(spv::OpTypeInt).AddSeq(t, dtype.bits(), 0).Commit(&global_);
    } else if (dtype.MatchesCode(DLDataTypeCode::kDLFloat)) {
      ib_.Begin(spv::OpTypeFloat).AddSeq(t, dtype.bits()).Commit(&global_);
    } else {
      TVM_FFI_THROW(InternalError) << "Unsupported primitive type " << dtype;
    }
    return t;
  } else {
    SType t;
    t.id = id_counter_++;
    t.type = dtype;
    SType base_type = GetSType(dtype.WithLanes(1));

    if (row * col == 0) {
      TVM_FFI_ICHECK((row == 0) && (col == 0));
      ib_.Begin(spv::OpTypeVector).AddSeq(t, base_type, dtype.lanes()).Commit(&global_);
    } else {
      Value v_row = GetSpecConst(GetSType(PrimType::UInt(32)), row);
      Value v_col = GetSpecConst(GetSType(PrimType::UInt(32)), col);
      Value scope = UIntImm(GetSType(PrimType::UInt(32)), spv::ScopeSubgroup);
      ib_.Begin(spv::OpTypeCooperativeMatrixNV)
          .AddSeq(t, base_type, scope, v_row, v_col)
          .Commit(&global_);
    }
    return t;
  }
}

void IRBuilder::AddCapabilityFor(const PrimType& dtype) {
  // Declare appropriate capabilities for int/float types
  if (dtype.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {
    if (dtype.bits() == 8) {
      TVM_FFI_ICHECK(spirv_support_.supports_int8)
          << "Vulkan target does not support Int8 capability.  "
          << "If your device supports 8-bit int operations, "
          << "please either add -supports_int8=1 to the target, "
          << "or query all device parameters by adding -from_device=0.";
      capabilities_used_.insert(spv::CapabilityInt8);
    } else if (dtype.bits() == 16) {
      TVM_FFI_ICHECK(spirv_support_.supports_int16)
          << "Vulkan target does not support Int16 capability.  "
          << "If your device supports 16-bit int operations, "
          << "please either add -supports_int16=1 to the target, "
          << "or query all device parameters by adding -from_device=0.";
      capabilities_used_.insert(spv::CapabilityInt16);
    } else if (dtype.bits() == 64) {
      TVM_FFI_ICHECK(spirv_support_.supports_int64)
          << "Vulkan target does not support Int64 capability.  "
          << "If your device supports 64-bit int operations, "
          << "please either add -supports_int64=1 to the target, "
          << "or query all device parameters by adding -from_device=0.";
      capabilities_used_.insert(spv::CapabilityInt64);
    }

  } else if (dtype.MatchesCode(DLDataTypeCode::kDLFloat)) {
    if (dtype.bits() == 16) {
      TVM_FFI_ICHECK(spirv_support_.supports_float16)
          << "Vulkan target does not support Float16 capability.  "
          << "If your device supports 16-bit float operations, "
          << "please either add -supports_float16=1 to the target, "
          << "or query all device parameters by adding -from_device=0.";
      capabilities_used_.insert(spv::CapabilityFloat16);
    } else if (dtype.bits() == 64) {
      TVM_FFI_ICHECK(spirv_support_.supports_float64)
          << "Vulkan target does not support Float64 capability.  "
          << "If your device supports 64-bit float operations, "
          << "please either add -supports_float64=1 to the target, "
          << "or query all device parameters by adding -from_device=0.";
      capabilities_used_.insert(spv::CapabilityFloat64);
    }
  }

  // Declare ability to read type to/from storage buffers.  Doing so
  // here is a little bit overzealous, should be relaxed in the
  // future.  Requiring StorageBuffer8BitAccess in order to declare an
  // Int8 prevents use of an 8-bit loop iterator on a device that
  // supports Int8 but doesn't support 8-bit buffer access.
  if (dtype.bits() == 8 && !dtype.MatchesCode(DLDataTypeCode::kDLBool)) {
    TVM_FFI_ICHECK(spirv_support_.supports_storage_buffer_8bit_access)
        << "Vulkan target does not support StorageBuffer8BitAccess.  "
        << "If your device supports 8-bit buffer access, "
        << "please either add -supports_8bit_buffer=1 to the target, "
        << "or query all device parameters by adding -from_device=0.";
    capabilities_used_.insert(spv::CapabilityStorageBuffer8BitAccess);
    extensions_used_.insert("SPV_KHR_8bit_storage");

    TVM_FFI_ICHECK(spirv_support_.supports_storage_buffer_storage_class)
        << "Illegal Vulkan target description.  "
        << "Vulkan spec requires extension VK_KHR_storage_buffer_storage_class "
        << "if VK_KHR_8bit_storage is supported.  "
        << "Please either add -supports_storage_buffer_storage_class=1 to the target, "
        << "or query all device parameters by adding -from_device=0.";
  } else if (dtype.bits() == 16) {
    TVM_FFI_ICHECK(spirv_support_.supports_storage_buffer_16bit_access)
        << "Vulkan target does not support StorageBuffer16BitAccess.  "
        << "If your device supports 16-bit buffer access, "
        << "please either add -supports_16bit_buffer=1 to the target, "
        << "or query all device parameters by adding -from_device=0.";

    extensions_used_.insert("SPV_KHR_16bit_storage");
    if (spirv_support_.supports_storage_buffer_storage_class) {
      capabilities_used_.insert(spv::CapabilityStorageBuffer16BitAccess);
    } else {
      capabilities_used_.insert(spv::CapabilityStorageUniformBufferBlock16);
    }
  }
}

PhiValue IRBuilder::MakePhi(const SType& out_type, uint32_t num_incoming) {
  Value val = NewValue(out_type, kNormal);
  ib_.Begin(spv::OpPhi).AddSeq(out_type, val);
  for (uint32_t i = 0; i < 2 * num_incoming; ++i) {
    ib_.Add(0);
  }
  PhiValue phi;
  phi.id = val.id;
  phi.stype = out_type;
  phi.flag = kNormal;
  phi.instr = ib_.Commit(&function_);
  TVM_FFI_ICHECK_EQ(phi.instr.WordCount(), 2 * num_incoming + 3);
  return phi;
}

Value IRBuilder::CallGLSL450(const SType& ret_type, uint32_t inst_id,
                             const std::vector<Value>& args) {
  Value val = NewValue(ret_type, kNormal);
  ib_.Begin(spv::OpExtInst).AddSeq(ret_type, val, ext_glsl450_, inst_id);
  for (const Value& v : args) {
    ib_.Add(v);
  }
  ib_.Commit(&function_);
  return val;
}

Value IRBuilder::CallKHRIntegerDotProduct(const SType& ret_type, const std::vector<Value>& args,
                                          const PrimType& dtype) {
  if (args.size() != 3) {
    TVM_FFI_THROW(InternalError) << "Unresolved arguments in SPIRV_KHR_integer_dot_product";
  }
  Value val = NewValue(ret_type, kNormal);
#ifdef TVM_SPIRV_KHR_INTEGER_DOT_PRODUCT
  TVM_FFI_ICHECK(spirv_support_.supports_integer_dot_product)
      << "Vulkan target does not support integer dot product capability.  "
      << "If your device supports integer dot product operations, "
      << "please either add -mattr=+dotprod to the target, "
      << "or query all device parameters by adding -from_device=0.";
  if (dtype.MatchesCode(DLDataTypeCode::kDLInt)) {
    ib_.Begin(spv::OpSDotAccSatKHR).AddSeq(ret_type, val);
  } else if (dtype.MatchesCode(DLDataTypeCode::kDLUInt)) {
    ib_.Begin(spv::OpUDotAccSatKHR).AddSeq(ret_type, val);
  } else {
    TVM_FFI_THROW(InternalError) << "Unsupported type";
  }
#else
  TVM_FFI_THROW(InternalError)
      << "Please turn on USE_SPIRV_KHR_INTEGER_DOT_PRODUCT in config.cmake";
#endif

  for (const Value& v : args) {
    ib_.Add(v);
  }
  ib_.Commit(&function_);
  return val;
}

Value IRBuilder::Concat(const std::vector<Value>& vec) {
  bool is_const = vec[0].flag == kConstant;
  PrimType etype = AsPrimType(vec[0].stype);
  int lanes = etype.lanes();
  for (size_t i = 1; i < vec.size(); ++i) {
    PrimType vector_type = AsPrimType(vec[i].stype);
    TVM_FFI_ICHECK_EQ(etype, vector_type.WithLanes(1))
        << "Cannot concat vector of different element type";
    lanes += vector_type.lanes();
    is_const = is_const && (vec[i].flag == kConstant);
  }
  Value ret = NewValue(GetSType(etype.WithLanes(lanes)), kNormal);
  if (is_const && vec.size() == static_cast<size_t>(lanes)) {
    ib_.Begin(spv::OpConstantComposite);
    ib_.AddSeq(ret.stype, ret);
    for (const Value& v : vec) {
      ib_.Add(v);
    }
    ib_.Commit(&global_);
  } else {
    ib_.Begin(spv::OpCompositeConstruct);
    ib_.AddSeq(ret.stype, ret);
    for (const Value& v : vec) {
      ib_.Add(v);
    }
    ib_.Commit(&function_);
  }
  return ret;
}

Value IRBuilder::Cast(const SType& dst_type, spirv::Value value) {
  TVM_FFI_ICHECK_NE(value.stype.id, 0U);
  if (value.stype.id == dst_type.id) return value;
  PrimType from = AsPrimType(value.stype);
  PrimType to = AsPrimType(dst_type);
  TVM_FFI_ICHECK_EQ(from.lanes(), to.lanes());
  if (from == PrimType::Bool()) {
    if (to.MatchesCode(DLDataTypeCode::kDLInt)) {
      return Select(value, IntImm(dst_type, 1), IntImm(dst_type, 0));
    } else if (to.MatchesCode(DLDataTypeCode::kDLUInt)) {
      return Select(value, UIntImm(dst_type, 1), UIntImm(dst_type, 0));
    } else if (to.MatchesCode(DLDataTypeCode::kDLFloat)) {
      return MakeValue(spv::OpConvertUToF, dst_type,
                       Select(value, UIntImm(t_uint32_, 1), UIntImm(t_uint32_, 0)));
    } else {
      TVM_FFI_THROW(InternalError) << "cannot cast from " << from << " to " << to;
      return Value();
    }
  } else if (to == PrimType::Bool()) {
    if (from.MatchesCode(DLDataTypeCode::kDLInt)) {
      return NE(value, IntImm(value.stype, 0));
    } else if (from.MatchesCode(DLDataTypeCode::kDLUInt)) {
      return NE(value, UIntImm(value.stype, 0));
    } else {
      TVM_FFI_THROW(InternalError) << "cannot cast from " << from << " to " << to;
      return Value();
    }
  } else if (from.MatchesCode(DLDataTypeCode::kDLInt) && to.MatchesCode(DLDataTypeCode::kDLInt)) {
    return MakeValue(spv::OpSConvert, dst_type, value);
  } else if (from.MatchesCode(DLDataTypeCode::kDLUInt) && to.MatchesCode(DLDataTypeCode::kDLUInt)) {
    return MakeValue(spv::OpUConvert, dst_type, value);
  } else if (from.MatchesCode(DLDataTypeCode::kDLUInt) && to.MatchesCode(DLDataTypeCode::kDLInt)) {
    if (from.bits() != to.bits()) {
      value = MakeValue(spv::OpUConvert, GetSType(from.WithBits(to.bits())), value);
    }
    return MakeValue(spv::OpBitcast, dst_type, value);
  } else if (from.MatchesCode(DLDataTypeCode::kDLInt) && to.MatchesCode(DLDataTypeCode::kDLUInt)) {
    if (from.bits() != to.bits()) {
      value = MakeValue(spv::OpSConvert, GetSType(from.WithBits(to.bits())), value);
    }
    return MakeValue(spv::OpBitcast, dst_type, value);
  } else if (from.MatchesCode(DLDataTypeCode::kDLFloat) && to.MatchesCode(DLDataTypeCode::kDLInt)) {
    return MakeValue(spv::OpConvertFToS, dst_type, value);
  } else if (from.MatchesCode(DLDataTypeCode::kDLFloat) &&
             to.MatchesCode(DLDataTypeCode::kDLUInt)) {
    return MakeValue(spv::OpConvertFToU, dst_type, value);
  } else if (from.MatchesCode(DLDataTypeCode::kDLInt) && to.MatchesCode(DLDataTypeCode::kDLFloat)) {
    return MakeValue(spv::OpConvertSToF, dst_type, value);
  } else if (from.MatchesCode(DLDataTypeCode::kDLUInt) &&
             to.MatchesCode(DLDataTypeCode::kDLFloat)) {
    return MakeValue(spv::OpConvertUToF, dst_type, value);
  } else if (from.MatchesCode(DLDataTypeCode::kDLFloat) &&
             to.MatchesCode(DLDataTypeCode::kDLFloat)) {
    return MakeValue(spv::OpFConvert, dst_type, value);
  } else {
    TVM_FFI_THROW(InternalError) << "do not support type cast from " << from << " to " << to;
    return Value();
  }
}

Value IRBuilder::GetCompositeConst(const SType& ele_stype, const SType& composite_stype,
                                   const double dval) {
  auto key = std::make_pair(composite_stype.id, dval);
  auto it = composite_const_tbl_.find(key);
  if (it != composite_const_tbl_.end()) {
    return it->second;
  }
  spirv::Value const_val = FloatImm(ele_stype, dval);
  Value new_val = NewValue(composite_stype, kNormal);
  ib_.Begin(spv::OpConstantComposite).AddSeq(composite_stype, new_val, const_val);
  ib_.Commit(&global_);
  composite_const_tbl_[key] = new_val;
  return new_val;
}

Value IRBuilder::GetSpecConst(const SType& dtype, uint64_t value) {
  TVM_FFI_ICHECK_LE(AsPrimType(dtype).bits(), 32);
  Value ret = NewValue(dtype, kSpecConst);
  ib_.Begin(spv::OpSpecConstant).AddSeq(dtype, ret);
  ib_.Add(static_cast<uint32_t>(value));
  ib_.Commit(&global_);
  return ret;
}

#define DEFINE_BUILDER_BINARY_USIGN_OP(_OpName, _Op)                         \
  Value IRBuilder::_OpName(Value a, Value b) {                               \
    TVM_FFI_ICHECK_EQ(a.stype.id, b.stype.id);                               \
    PrimType type = AsPrimType(a.stype);                                     \
    if (type.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) { \
      return MakeValue(spv::OpI##_Op, a.stype, a, b);                        \
    } else {                                                                 \
      TVM_FFI_ICHECK(type.MatchesCode(DLDataTypeCode::kDLFloat));            \
      return MakeValue(spv::OpF##_Op, a.stype, a, b);                        \
    }                                                                        \
  }

#define DEFINE_BUILDER_BINARY_SIGN_OP(_OpName, _Op)               \
  Value IRBuilder::_OpName(Value a, Value b) {                    \
    TVM_FFI_ICHECK_EQ(a.stype.id, b.stype.id);                    \
    PrimType type = AsPrimType(a.stype);                          \
    if (type.MatchesCode(DLDataTypeCode::kDLInt)) {               \
      return MakeValue(spv::OpS##_Op, a.stype, a, b);             \
    } else if (type.MatchesCode(DLDataTypeCode::kDLUInt)) {       \
      return MakeValue(spv::OpU##_Op, a.stype, a, b);             \
    } else {                                                      \
      TVM_FFI_ICHECK(type.MatchesCode(DLDataTypeCode::kDLFloat)); \
      return MakeValue(spv::OpF##_Op, a.stype, a, b);             \
    }                                                             \
  }

DEFINE_BUILDER_BINARY_USIGN_OP(Add, Add);
DEFINE_BUILDER_BINARY_USIGN_OP(Sub, Sub);
DEFINE_BUILDER_BINARY_USIGN_OP(Mul, Mul);
DEFINE_BUILDER_BINARY_SIGN_OP(Div, Div);

Value IRBuilder::Mod(Value a, Value b) {
  TVM_FFI_ICHECK_EQ(a.stype.id, b.stype.id);
  PrimType type = AsPrimType(a.stype);
  if (type.MatchesCode(DLDataTypeCode::kDLInt)) {
    return MakeValue(spv::OpSRem, a.stype, a, b);
  } else if (type.MatchesCode(DLDataTypeCode::kDLUInt)) {
    return MakeValue(spv::OpUMod, a.stype, a, b);
  } else {
    TVM_FFI_ICHECK(type.MatchesCode(DLDataTypeCode::kDLFloat));
    return MakeValue(spv::OpFRem, a.stype, a, b);
  }
}

#define DEFINE_BUILDER_CMP_OP(_OpName, _Op)                                             \
  Value IRBuilder::_OpName(Value a, Value b) {                                          \
    TVM_FFI_ICHECK_EQ(a.stype.id, b.stype.id);                                          \
    PrimType a_type = AsPrimType(a.stype);                                              \
    PrimType b_type = AsPrimType(b.stype);                                              \
    TVM_FFI_ICHECK_EQ(a_type.lanes(), b_type.lanes());                                  \
    const auto& bool_type = this->GetSType(PrimType::Bool().WithLanes(a_type.lanes())); \
    if (a_type.MatchesCode(DLDataTypeCode::kDLInt)) {                                   \
      return MakeValue(spv::OpS##_Op, bool_type, a, b);                                 \
    } else if (a_type.MatchesCode(DLDataTypeCode::kDLUInt)) {                           \
      return MakeValue(spv::OpU##_Op, bool_type, a, b);                                 \
    } else {                                                                            \
      TVM_FFI_ICHECK(a_type.MatchesCode(DLDataTypeCode::kDLFloat));                     \
      return MakeValue(spv::OpFOrd##_Op, bool_type, a, b);                              \
    }                                                                                   \
  }

DEFINE_BUILDER_CMP_OP(LT, LessThan);
DEFINE_BUILDER_CMP_OP(LE, LessThanEqual);
DEFINE_BUILDER_CMP_OP(GT, GreaterThan);
DEFINE_BUILDER_CMP_OP(GE, GreaterThanEqual);

#define DEFINE_BUILDER_CMP_UOP(_OpName, _Op)                                            \
  Value IRBuilder::_OpName(Value a, Value b) {                                          \
    TVM_FFI_ICHECK_EQ(a.stype.id, b.stype.id);                                          \
    PrimType a_type = AsPrimType(a.stype);                                              \
    PrimType b_type = AsPrimType(b.stype);                                              \
    TVM_FFI_ICHECK_EQ(a_type.lanes(), b_type.lanes());                                  \
    const auto& bool_type = this->GetSType(PrimType::Bool().WithLanes(a_type.lanes())); \
    if (a_type.MatchesCode(DLDataTypeCode::kDLInt, DLDataTypeCode::kDLUInt)) {          \
      return MakeValue(spv::OpI##_Op, bool_type, a, b);                                 \
    } else {                                                                            \
      TVM_FFI_ICHECK(a_type.MatchesCode(DLDataTypeCode::kDLFloat));                     \
      return MakeValue(spv::OpFOrd##_Op, bool_type, a, b);                              \
    }                                                                                   \
  }

DEFINE_BUILDER_CMP_UOP(EQ, Equal);
DEFINE_BUILDER_CMP_UOP(NE, NotEqual);

Value IRBuilder::Select(Value cond, Value a, Value b) {
  TVM_FFI_ICHECK_EQ(a.stype.id, b.stype.id);
  TVM_FFI_ICHECK_EQ(AsPrimType(cond.stype).WithLanes(1), PrimType::Bool());
  return MakeValue(spv::OpSelect, a.stype, cond, a, b);
}

}  // namespace spirv
}  // namespace codegen
}  // namespace tvm
