From b8cd4e1370629d773f1960b9a8698a2d264d00d9 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Sun, 6 Mar 2022 06:29:35 +0900 Subject: [PATCH 01/10] [Vulkan] Add cmake change for spirv dot product --- CMakeLists.txt | 10 +++++++ cmake/config.cmake | 11 ++++++++ cmake/modules/Vulkan.cmake | 6 ++++- cmake/utils/FindVulkan.cmake | 26 +++++++++++++++---- .../deploy_models/deploy_prequantized.py | 3 ++- 5 files changed, 49 insertions(+), 7 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 9c6a7dddfdf6..8e66d364eade 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -26,6 +26,16 @@ endif() tvm_option(USE_CUDA "Build with CUDA" OFF) tvm_option(USE_OPENCL "Build with OpenCL" OFF) tvm_option(USE_VULKAN "Build with Vulkan" OFF) + + +# Whether to use spirv-tools.and SPIRV-Headers from Khronos github or gitlab. +# +# Possible values: +# - OFF: not to use +# - /path/to/install: path to your khronis spirv-tools and SPIRV-Headers installation directory +# +tvm_option(USE_KHRONOS_SPIRV "Whether to use spirv-tools.and SPIRV-Headers from Khronos github or gitlab" OFF) +tvm_option(USE_SPIRV_KHR_INTEGER_DOT_PRODUCT "whether enable SPIRV_KHR_DOT_PRODUCT" OFF) tvm_option(USE_METAL "Build with Metal" OFF) tvm_option(USE_ROCM "Build with ROCM" OFF) tvm_option(ROCM_PATH "The path to rocm" /opt/rocm) diff --git a/cmake/config.cmake b/cmake/config.cmake index 62eeb34fead7..3adac49db8fe 100644 --- a/cmake/config.cmake +++ b/cmake/config.cmake @@ -81,6 +81,17 @@ set(USE_METAL OFF) # - /path/to/vulkan-sdk: use specific path to vulkan-sdk set(USE_VULKAN OFF) +# Whether to use spirv-tools.and SPIRV-Headers from Khronos github or gitlab. +# +# Possible values: +# - OFF: not to use +# - /path/to/install: path to your khronis spirv-tools and SPIRV-Headers installation directory +# +set(USE_KHRONOS_SPIRV OFF) + +# whether enable SPIRV_KHR_DOT_PRODUCT +set(USE_SPIRV_KHR_INTEGER_DOT_PRODUCT OFF) + # Whether enable OpenGL runtime set(USE_OPENGL OFF) diff --git a/cmake/modules/Vulkan.cmake b/cmake/modules/Vulkan.cmake index 17af5bafde75..7470fb6125a4 100644 --- a/cmake/modules/Vulkan.cmake +++ b/cmake/modules/Vulkan.cmake @@ -16,12 +16,16 @@ # under the License. # Be compatible with older version of CMake -find_vulkan(${USE_VULKAN}) +find_vulkan(${USE_VULKAN} ${USE_KHRONOS_SPIRV}) if(USE_VULKAN) if(NOT Vulkan_FOUND) message(FATAL_ERROR "Cannot find Vulkan, USE_VULKAN=" ${USE_VULKAN}) endif() + if (USE_SPIRV_KHR_INTEGER_DOT_PRODUCT) + add_definitions(-DTVM_SPIRV_KHR_INTEGER_DOT_PRODUCT=1) + message(STATUS "Enable SPIRV_KHR_INTEGER_DOT_PRODUCT") + endif() include_directories(SYSTEM ${Vulkan_INCLUDE_DIRS}) message(STATUS "Build with Vulkan support") tvm_file_glob(GLOB RUNTIME_VULKAN_SRCS src/runtime/vulkan/*.cc) diff --git a/cmake/utils/FindVulkan.cmake b/cmake/utils/FindVulkan.cmake index feb5eec74d89..4349125d3f37 100644 --- a/cmake/utils/FindVulkan.cmake +++ b/cmake/utils/FindVulkan.cmake @@ -32,7 +32,7 @@ # - Vulkan_SPIRV_TOOLS_LIBRARY # -macro(find_vulkan use_vulkan) +macro(find_vulkan use_vulkan use_khronos_spirv) set(__use_vulkan ${use_vulkan}) if(IS_DIRECTORY ${__use_vulkan}) set(__vulkan_sdk ${__use_vulkan}) @@ -43,6 +43,15 @@ macro(find_vulkan use_vulkan) set(__vulkan_sdk "") endif() + + if(IS_DIRECTORY ${use_khronos_spirv}) + set(__use_khronos_spirv ${use_khronos_spirv}) + message(STATUS "Custom khronos spirv PATH=" ${__use_khronos_spirv}) + else() + set(__use_khronos_spirv "") + endif() + + if(__vulkan_sdk) set(Vulkan_INCLUDE_DIRS ${__vulkan_sdk}/include) find_library(Vulkan_LIBRARY NAMES vulkan vulkan-1 PATHS ${__vulkan_sdk}/lib) @@ -61,11 +70,18 @@ macro(find_vulkan use_vulkan) if(Vulkan_FOUND) get_filename_component(VULKAN_LIBRARY_PATH ${Vulkan_LIBRARY} DIRECTORY) - find_library(Vulkan_SPIRV_TOOLS_LIBRARY SPIRV-Tools - HINTS ${VULKAN_LIBRARY_PATH} ${VULKAN_LIBRARY_PATH}/spirv-tools ${VULKAN_SDK}/lib) + if (WIN32) + find_library(Vulkan_SPIRV_TOOLS_LIBRARY SPIRV-Tools + HINTS ${__use_khronos_spirv}/spirv-tools/lib ${VULKAN_LIBRARY_PATH} ${VULKAN_LIBRARY_PATH}/spirv-tools ${VULKAN_SDK}/lib) + find_path(_libspirv libspirv.h HINTS ${__use_khronos_spirv}/spirv-tools/include ${Vulkan_INCLUDE_DIRS} PATH_SUFFIXES vulkan spirv-tools) + find_path(_spirv spirv.hpp HINTS ${__use_khronos_spirv}/SPIRV-Headers/include ${Vulkan_INCLUDE_DIRS} PATH_SUFFIXES vulkan SPIRV spirv/unified1 spirv-headers) + else() + find_library(Vulkan_SPIRV_TOOLS_LIBRARY SPIRV-Tools + HINTS ${__use_khronos_spirv}/lib ${VULKAN_LIBRARY_PATH} ${VULKAN_LIBRARY_PATH}/spirv-tools ${VULKAN_SDK}/lib) + find_path(_libspirv libspirv.h HINTS ${__use_khronos_spirv}/include ${Vulkan_INCLUDE_DIRS} PATH_SUFFIXES vulkan spirv-tools) + find_path(_spirv spirv.hpp HINTS ${__use_khronos_spirv}/include ${Vulkan_INCLUDE_DIRS} PATH_SUFFIXES vulkan SPIRV spirv/unified1 spirv-headers) + endif() - find_path(_libspirv libspirv.h HINTS ${Vulkan_INCLUDE_DIRS} PATH_SUFFIXES vulkan spirv-tools) - find_path(_spirv spirv.hpp HINTS ${Vulkan_INCLUDE_DIRS} PATH_SUFFIXES vulkan SPIRV spirv/unified1 spirv-headers) find_path(_glsl_std GLSL.std.450.h HINTS ${Vulkan_INCLUDE_DIRS} PATH_SUFFIXES vulkan SPIRV spirv/unified1 spirv-headers) list(APPEND Vulkan_INCLUDE_DIRS ${_libspirv} ${_spirv} ${_glsl_std}) message(STATUS "Vulkan_INCLUDE_DIRS=" ${Vulkan_INCLUDE_DIRS}) diff --git a/gallery/how_to/deploy_models/deploy_prequantized.py b/gallery/how_to/deploy_models/deploy_prequantized.py index 11a9e3e3eee8..caee2b3b415a 100644 --- a/gallery/how_to/deploy_models/deploy_prequantized.py +++ b/gallery/how_to/deploy_models/deploy_prequantized.py @@ -175,7 +175,8 @@ def quantize_model(model, inp): # # Under the hood, quantization specific operators are lowered to a sequence of # standard Relay operators before compilation. -tvm_result, rt_mod = run_tvm_model(mod, params, input_name, inp, target="llvm") +target = "llvm" +tvm_result, rt_mod = run_tvm_model(mod, params, input_name, inp, target=target) ########################################################################## # Compare the output labels From 4dde9b286029ffeeb8fef7395a30541cbfbfc039 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Sun, 6 Mar 2022 10:11:51 +0900 Subject: [PATCH 02/10] add spirv and vk runtime change --- python/tvm/target/target.py | 6 ++++ src/runtime/vulkan/vulkan_device.cc | 25 ++++++++-------- src/runtime/vulkan/vulkan_device.h | 1 + src/runtime/vulkan/vulkan_device_api.cc | 5 ++++ src/target/spirv/codegen_spirv.cc | 23 ++++++++++++--- src/target/spirv/ir_builder.cc | 38 +++++++++++++++++++++++++ src/target/spirv/ir_builder.h | 11 +++++++ src/target/spirv/spirv_support.cc | 13 +++++++++ src/target/spirv/spirv_support.h | 14 +++++++++ src/target/target_kind.cc | 2 ++ 10 files changed, 122 insertions(+), 16 deletions(-) diff --git a/python/tvm/target/target.py b/python/tvm/target/target.py index d08ef3d3b92c..72d5b97a6498 100644 --- a/python/tvm/target/target.py +++ b/python/tvm/target/target.py @@ -188,6 +188,12 @@ def mattr(self): """Returns the mattr from the target if it exists.""" return list(self.attrs.get("mattr", [])) + @property + def supports_integer_dot_product(self): + if self.attrs.get("supports_integer_dot_product", []): + return bool(self.attrs["supports_integer_dot_product"]) + return False + @property def libs(self): return list(self.attrs.get("libs", [])) diff --git a/src/runtime/vulkan/vulkan_device.cc b/src/runtime/vulkan/vulkan_device.cc index f5d26ace50d9..29908bed8189 100644 --- a/src/runtime/vulkan/vulkan_device.cc +++ b/src/runtime/vulkan/vulkan_device.cc @@ -132,6 +132,8 @@ VulkanDeviceProperties::VulkanDeviceProperties(const VulkanInstance& instance, device.HasExtension("VK_KHR_dedicated_allocation") && !support::BoolEnvironmentVar("TVM_VULKAN_DISABLE_DEDICATED_ALLOCATION"); + supports_integer_dot_product = device.HasExtension("VK_KHR_shader_integer_dot_product"); + // The check of VK_SHADER_STAGE_COMPUTE_BIT isn't technically // needed, since it will be set so long at least one queue has // VK_QUEUE_COMPUTE_BIT. Including it to avoid potential future @@ -410,18 +412,17 @@ uint32_t VulkanDevice::SelectComputeQueueFamily() const { std::vector VulkanDevice::SelectEnabledExtensions() const { std::vector required_extensions{}; - std::vector optional_extensions{ - "VK_KHR_driver_properties", - "VK_KHR_storage_buffer_storage_class", - "VK_KHR_8bit_storage", - "VK_KHR_16bit_storage", - "VK_KHR_shader_float16_int8", - "VK_KHR_push_descriptor", - "VK_KHR_descriptor_update_template", - "VK_KHR_get_memory_requirements2", - "VK_KHR_dedicated_allocation", - "VK_KHR_spirv_1_4", - }; + std::vector optional_extensions{"VK_KHR_driver_properties", + "VK_KHR_storage_buffer_storage_class", + "VK_KHR_8bit_storage", + "VK_KHR_16bit_storage", + "VK_KHR_shader_float16_int8", + "VK_KHR_push_descriptor", + "VK_KHR_descriptor_update_template", + "VK_KHR_get_memory_requirements2", + "VK_KHR_dedicated_allocation", + "VK_KHR_spirv_1_4", + "VK_KHR_shader_integer_dot_product"}; uint32_t device_extension_prop_count; VULKAN_CALL(vkEnumerateDeviceExtensionProperties(physical_device_, nullptr, diff --git a/src/runtime/vulkan/vulkan_device.h b/src/runtime/vulkan/vulkan_device.h index d7788ef5df29..3ca2d093bf1d 100644 --- a/src/runtime/vulkan/vulkan_device.h +++ b/src/runtime/vulkan/vulkan_device.h @@ -81,6 +81,7 @@ struct VulkanDeviceProperties { bool supports_storage_buffer_storage_class{false}; bool supports_push_descriptor{false}; bool supports_dedicated_allocation{false}; + bool supports_integer_dot_product{false}; uint32_t supported_subgroup_operations{0}; uint32_t max_num_threads{1}; uint32_t thread_warp_size{1}; diff --git a/src/runtime/vulkan/vulkan_device_api.cc b/src/runtime/vulkan/vulkan_device_api.cc index cf0b16c6c471..676f14667d70 100644 --- a/src/runtime/vulkan/vulkan_device_api.cc +++ b/src/runtime/vulkan/vulkan_device_api.cc @@ -236,6 +236,11 @@ void VulkanDeviceAPI::GetTargetProperty(Device dev, const std::string& property, if (property == "max_shared_memory_per_block") { *rv = int64_t(prop.max_shared_memory_per_block); } + + if (property == "supports_integer_dot_product") { + *rv = prop.supports_integer_dot_product; + } + if (property == "device_name") { *rv = prop.device_name; } diff --git a/src/target/spirv/codegen_spirv.cc b/src/target/spirv/codegen_spirv.cc index 0427d8cd5853..1bb457a8b913 100644 --- a/src/target/spirv/codegen_spirv.cc +++ b/src/target/spirv/codegen_spirv.cc @@ -377,8 +377,21 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const CallNode* op) { } else if (op->op.same_as(builtin::popcount())) { return builder_->MakeValue(spv::OpBitCount, builder_->GetSType(op->dtype), MakeValue(op->args[0])); - } else if (op->op.same_as(builtin::call_extern()) || - op->op.same_as(builtin::call_pure_extern())) { + } else if (op->op.same_as(builtin::call_pure_extern())) { + ICHECK_GE(op->args.size(), 1U); + const std::string& func_name = op->args[0].as()->value; + if (func_name == "__dp4a") { + std::vector values; + for (size_t i = 1; i < op->args.size(); ++i) { + values.push_back(MakeValue(op->args[i])); + } + return builder_->CallKHRIntegerDotProduct(builder_->GetSType(op->dtype), values, op->dtype); + } else { + LOG(FATAL) << "SPIR-V shader cannot make extern calls. Graph contains extern \"" + << Downcast(op->args[0]) << "\""; + return spirv::Value(); + } + } else if (op->op.same_as(builtin::call_extern())) { ICHECK_GE(op->args.size(), 1U); LOG(FATAL) << "SPIR-V shader cannot make extern calls. Graph contains extern \"" << Downcast(op->args[0]) << "\""; @@ -653,8 +666,10 @@ void CodeGenSPIRV::VisitStmt_(const AllocateNode* op) { builder_->Allocate(etype, static_cast(constant_size), spv::StorageClassFunction); } else if (storage_scope.rank == runtime::StorageRank::kShared) { // Shared memory - buf = - builder_->Allocate(etype, static_cast(constant_size), spv::StorageClassWorkgroup); + // Aligned on 4-byte boundary + int32_t aligned_constant_size = ((constant_size + 3) & ~0x3); + buf = builder_->Allocate(etype, static_cast(aligned_constant_size), + spv::StorageClassWorkgroup); size_t num_bytes = op->dtype.bytes() * op->dtype.lanes() * static_cast(constant_size); shared_memory_bytes_used_ += num_bytes; diff --git a/src/target/spirv/ir_builder.cc b/src/target/spirv/ir_builder.cc index 9696043a244d..46c9c5869c79 100644 --- a/src/target/spirv/ir_builder.cc +++ b/src/target/spirv/ir_builder.cc @@ -52,6 +52,14 @@ void IRBuilder::InitHeader() { // 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 + // memory model ib_.Begin(spv::OpMemoryModel) .AddSeq(spv::AddressingModelLogical, spv::MemoryModelGLSL450) @@ -603,6 +611,36 @@ Value IRBuilder::CallGLSL450(const SType& ret_type, uint32_t inst_id, return val; } +Value IRBuilder::CallKHRIntegerDotProduct(const SType& ret_type, const std::vector& args, + const DataType& dtype) { + if (args.size() != 3) { + LOG(FATAL) << "Unresolved arguments in SPIRV_KHR_integer_dot_product"; + } + Value val = NewValue(ret_type, kNormal); +#ifdef TVM_SPIRV_KHR_INTEGER_DOT_PRODUCT + 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.is_int()) { + ib_.Begin(spv::OpSDotAccSatKHR).AddSeq(ret_type, val); + } else if (dtype.is_uint()) { + ib_.Begin(spv::OpUDotAccSatKHR).AddSeq(ret_type, val); + } else { + LOG(FATAL) << "Unsupported type"; + } +#else + LOG(FATAL) << "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& vec) { bool is_const = vec[0].flag == kConstant; DataType etype = vec[0].stype.type; diff --git a/src/target/spirv/ir_builder.h b/src/target/spirv/ir_builder.h index 3e19b98100c0..f1b5397b3757 100644 --- a/src/target/spirv/ir_builder.h +++ b/src/target/spirv/ir_builder.h @@ -420,6 +420,17 @@ class IRBuilder { * \return The result value. */ Value CallGLSL450(const SType& ret_type, uint32_t inst_id, const std::vector& args); + + /*! + * \brief Create a SPIRV_KHR_integer_dot_product call + * + * \param ret_type The result type. + * \param args The arguments + * \return The result value. + */ + Value CallKHRIntegerDotProduct(const SType& ret_type, const std::vector& args, + const DataType& dtype); + /*! * \brief Build vector by concatenating components * diff --git a/src/target/spirv/spirv_support.cc b/src/target/spirv/spirv_support.cc index 1ef56198df7f..33055e7399d5 100644 --- a/src/target/spirv/spirv_support.cc +++ b/src/target/spirv/spirv_support.cc @@ -84,6 +84,19 @@ SPIRVSupport::SPIRVSupport(tvm::Target target) { if (target->GetAttr("supports_int64")) { supports_int64 = target->GetAttr("supports_int64").value(); } + // Check whether integer dot product is enabled in the target string. + if (target->GetAttr("supports_integer_dot_product")) { + supports_integer_dot_product = target->GetAttr("supports_integer_dot_product").value(); + } + // Check whether integer dot product is enabled in mattr. + if (const Optional>& v = target->GetAttr>("mattr")) { + for (const String& s : v.value()) { + if (s.compare("+dotprod") == 0) { + supports_integer_dot_product = true; + break; + } + } + } } } // namespace codegen diff --git a/src/target/spirv/spirv_support.h b/src/target/spirv/spirv_support.h index 04d13cca5031..6365e576b8cf 100644 --- a/src/target/spirv/spirv_support.h +++ b/src/target/spirv/spirv_support.h @@ -262,6 +262,20 @@ struct SPIRVSupport { * attempting to create a 64-bit int. */ bool supports_int64{false}; + + /*! + * \brief Whether the driver supports operations involving integer dot product. + * + * Vulkan extension: VK_KHR_shader_integer_dot_product + * SPV Extension name: SPV_KHR_integer_dot_product + * SPV Capability: spv::CapabilityDotProductKHR, + * spv::CapabilityDotProductInput4x8BitPackedKHR); + * + * If support is present, can perform integer dot product operations. If + * support is not present, codegen will throw exception on + * attempting to perform integer dot product. + */ + bool supports_integer_dot_product{false}; }; } // namespace codegen diff --git a/src/target/target_kind.cc b/src/target/target_kind.cc index 1131e6e7d2a8..6fef8b48c396 100644 --- a/src/target/target_kind.cc +++ b/src/target/target_kind.cc @@ -330,6 +330,7 @@ TVM_REGISTER_TARGET_KIND("metal", kDLMetal) .set_default_keys({"metal", "gpu"}); TVM_REGISTER_TARGET_KIND("vulkan", kDLVulkan) + .add_attr_option>("mattr") .add_attr_option("system-lib") // Feature support .add_attr_option("supports_float16") @@ -344,6 +345,7 @@ TVM_REGISTER_TARGET_KIND("vulkan", kDLVulkan) .add_attr_option("supports_storage_buffer_storage_class") .add_attr_option("supports_push_descriptor") .add_attr_option("supports_dedicated_allocation") + .add_attr_option("supports_integer_dot_product") .add_attr_option("supported_subgroup_operations") // Physical device limits .add_attr_option("max_num_threads", Integer(256)) From 4940215e0360f44112f3884d8fc737f2dbf912af Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Sun, 6 Mar 2022 15:44:25 +0900 Subject: [PATCH 03/10] add back conv2d int8 related change --- python/tvm/relay/op/strategy/cuda.py | 10 +++++-- python/tvm/relay/qnn/op/legalizations.py | 37 +++++++++++++++++++----- python/tvm/topi/cuda/conv2d_alter_op.py | 17 +++++++++-- python/tvm/topi/cuda/conv2d_int8.py | 17 +++++++---- python/tvm/topi/cuda/tensor_intrin.py | 26 ++++++++++++----- 5 files changed, 80 insertions(+), 27 deletions(-) diff --git a/python/tvm/relay/op/strategy/cuda.py b/python/tvm/relay/op/strategy/cuda.py index c059cf18f2a2..ec0d6e3a903e 100644 --- a/python/tvm/relay/op/strategy/cuda.py +++ b/python/tvm/relay/op/strategy/cuda.py @@ -145,7 +145,7 @@ def conv2d_strategy_cuda(attrs, inputs, out_type, target): if layout == "NCHW": assert kernel_layout == "OIHW" if ( - target.kind.name == "cuda" + (target.kind.name in ["cuda", "vulkan"]) and data.dtype in ("int8", "uint8") and kernel.dtype in ("int8", "uint8") ): @@ -296,7 +296,11 @@ def conv2d_strategy_cuda(attrs, inputs, out_type, target): "Unsupported shape for conv2d HWNC.\ Need to satisfy tensor core schedule." ) - elif target.kind.name == "cuda" and layout == "NCHW4c" and data.dtype in ["int8", "uint8"]: + elif ( + (target.kind.name in ["cuda", "vulkan"]) + and layout == "NCHW4c" + and data.dtype in ["int8", "uint8"] + ): assert kernel_layout == "OIHW4o4i" strategy.add_implementation( wrap_compute_conv2d(topi.cuda.conv2d_NCHWc_int8, True), @@ -372,7 +376,7 @@ def conv2d_strategy_cuda(attrs, inputs, out_type, target): ic_chunk = in_channels // 4 if ( - target.kind.name == "cuda" + (target.kind.name in ["cuda", "vulkan"]) and data.dtype in ["int8", "uint8"] and kernel.dtype in ["int8", "uint8"] and channels % groups == 0 diff --git a/python/tvm/relay/qnn/op/legalizations.py b/python/tvm/relay/qnn/op/legalizations.py index 11677315ac86..93b1ad7a44a8 100644 --- a/python/tvm/relay/qnn/op/legalizations.py +++ b/python/tvm/relay/qnn/op/legalizations.py @@ -387,6 +387,18 @@ def is_aarch64_arm(): return "aarch64" in target.attrs.get("mtriple", "") +def is_vulkan(): + """Checks whether we are compiling for a vulkan/spirv target.""" + target = tvm.target.Target.current(allow_none=False) + return "vulkan" in target.keys + + +def is_cuda(): + """Checks whether we are compiling for a cuda target.""" + target = tvm.target.Target.current(allow_none=False) + return "cuda" in target.keys + + ######################## # ARM CPU legalizations. ######################## @@ -438,17 +450,28 @@ def _qnn_dense_legalize_intel_cpu(attrs, inputs, types): ##################### -# CUDA legalizations. +# CUDA and vulkan legalizations. ##################### -@qnn_conv2d_legalize.register("cuda") +@qnn_conv2d_legalize.register(["cuda", "gpu"]) def _qnn_conv2d_legalize_cuda(attrs, inputs, types): - # CUDA prefers both datatypes to be int8. - return helper_change_dtypes_to_int8(attrs, inputs, types, relay.qnn.op.conv2d) + if is_vulkan(): + # prefers the dtypes to be same. Mixed type is not yet supported. + return helper_change_dtypes_to_be_same(attrs, inputs, types, relay.qnn.op.conv2d) + if is_cuda(): + # CUDA prefers both datatypes to be int8. + return helper_change_dtypes_to_int8(attrs, inputs, types, relay.qnn.op.conv2d) + return None -@qnn_dense_legalize.register("cuda") +@qnn_dense_legalize.register(["cuda", "gpu"]) def _qnn_dense_legalize_cuda(attrs, inputs, types): - # CUDA prefers both datatypes to be the int8. - return helper_change_dtypes_to_int8(attrs, inputs, types, relay.qnn.op.dense) + if is_vulkan(): + # prefers the dtypes to be same. Mixed type is not yet supported. + return helper_change_dtypes_to_be_same(attrs, inputs, types, relay.qnn.op.dense) + if is_cuda(): + # CUDA prefers both datatypes to be the int8. + return helper_change_dtypes_to_int8(attrs, inputs, types, relay.qnn.op.dense) + + return None diff --git a/python/tvm/topi/cuda/conv2d_alter_op.py b/python/tvm/topi/cuda/conv2d_alter_op.py index e6631d57b29e..eaafe15e9600 100644 --- a/python/tvm/topi/cuda/conv2d_alter_op.py +++ b/python/tvm/topi/cuda/conv2d_alter_op.py @@ -34,6 +34,9 @@ @nn.conv2d_alter_layout.register(["cuda", "gpu"]) def _alter_conv2d_layout(attrs, inputs, tinfos, out_type): target = tvm.target.Target.current(allow_none=False) + doit = "vulkan" in target.keys or "cuda" in target.keys + if not doit: + return None dispatch_ctx = autotvm.task.DispatchContext.current new_attrs = {k: attrs[k] for k in attrs.keys()} @@ -83,14 +86,18 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type): cfg = dispatch_ctx.query(target, workload) if cfg.is_fallback: # if is fallback, clear query cache and return None autotvm.task.clear_fallback_cache(target, workload) - return None + do_new_layout = False + if "vulkan" in target.keys: + do_new_layout = "+dotprod" in target.mattr or target.supports_integer_dot_product + if not do_new_layout: + return None topi_tmpl = workload[0] if topi_tmpl == "conv2d_NCHWc_int8.cuda": assert data_layout == "NCHW" and kernel_layout == "OIHW" N, CI, H, W = get_const_tuple(data.shape) CO, _, KH, KW = get_const_tuple(kernel.shape) - + assert CO % 4 == 0, "Number of output channels should be multiple of 4" new_layout = "NCHW4c" new_attrs["channels"] = CO new_attrs["data_layout"] = new_layout @@ -324,7 +331,7 @@ def _pad_conv2d_NHWC(db, di, do, data, kernel, out_channel, new_attrs, output_te return out -@conv2d_legalize.register("cuda") +@conv2d_legalize.register(["cuda", "gpu"]) def _conv2d_legalize(attrs, inputs, arg_types): """Legalizes Conv2D op. @@ -343,6 +350,10 @@ def _conv2d_legalize(attrs, inputs, arg_types): The legalized expr """ + target = tvm.target.Target.current(allow_none=False) + doit = "vulkan" in target.keys or "cuda" in target.keys + if not doit: + return None # Dilation not supported yet. Return None if dilation is not (1, 1) dilation = attrs.get_int_tuple("dilation") if not (dilation[0] == 1 and dilation[1] == 1): diff --git a/python/tvm/topi/cuda/conv2d_int8.py b/python/tvm/topi/cuda/conv2d_int8.py index 02470bab5228..15120f6a2532 100644 --- a/python/tvm/topi/cuda/conv2d_int8.py +++ b/python/tvm/topi/cuda/conv2d_int8.py @@ -153,13 +153,15 @@ def conv2d_NCHWc_int8(cfg, data, kernel, stride, padding, dilation, layout, out_ kh = te.reduce_axis((0, kernel_h), name="kh") kw = te.reduce_axis((0, kernel_w), name="kw") + packed_kernel_dtype = packed_kernel.dtype + packed_dtype = "int32" if packed_kernel_dtype == "int8" else "uint32" conv = te.compute( oshape, lambda n, oc_chunk, oh, ow, oc_block: te.sum( pad_data[ n, icc, oh * stride_h + kh * dilation_h, ow * stride_w + kw * dilation_w, icb - ].astype("int32") - * packed_kernel[oc_chunk, icc, kh, kw, oc_block, icb].astype("int32"), + ].astype(packed_dtype) + * packed_kernel[oc_chunk, icc, kh, kw, oc_block, icb].astype(packed_dtype), axis=[icc, kh, kw, icb], ), ) @@ -188,9 +190,6 @@ def conv2d_NCHWc_int8(cfg, data, kernel, stride, padding, dilation, layout, out_ return output -_dp4a = dp4a("shared", "shared", "local") - - @autotvm.register_topi_schedule("conv2d_NCHWc_int8.cuda") def schedule_conv2d_NCHWc_int8(cfg, outs): """Schedule conv2d int8 NCHWc template""" @@ -311,7 +310,13 @@ def _schedule_conv2d_NCHWc_int8(cfg, s, output): cfg["reorder_inner"].apply(s, conv, [rci, ryi, rxi]) _, rc_block = s[conv].split(rc_block, factor=4) - s[conv].tensorize(rc_block, _dp4a) + target = tvm.target.Target.current(allow_none=False) + do_tensorize = True + if "vulkan" in target.keys: + do_tensorize = "+dotprod" in target.mattr or target.supports_integer_dot_product + if do_tensorize: + dtypes = (pad_data.dtype, packed_kernel.dtype) + s[conv].tensorize(rc_block, dp4a("shared", "shared", "local", dtypes)) cache_loc = [rco, ryo, rxo][cfg["reorder_inner"].perm[-1]] s[AA].compute_at(s[conv], cache_loc) diff --git a/python/tvm/topi/cuda/tensor_intrin.py b/python/tvm/topi/cuda/tensor_intrin.py index 499f4b3376f2..c0596fc43262 100644 --- a/python/tvm/topi/cuda/tensor_intrin.py +++ b/python/tvm/topi/cuda/tensor_intrin.py @@ -20,7 +20,7 @@ from tvm import te -def dp4a(x_scope="local", y_scope="local", z_scope="local"): +def dp4a(x_scope="local", y_scope="local", z_scope="local", dtypes=("int8", "int8")): """ Int8 dot product reduced by every 4 elements using __dp4a @@ -32,6 +32,8 @@ def dp4a(x_scope="local", y_scope="local", z_scope="local"): The storage scope of buffer for rhs z_scope : str, optional The storage scope of buffer for result + dtypes: tuple of strs, optional + The dtype of x and y Returns ------- @@ -40,28 +42,36 @@ def dp4a(x_scope="local", y_scope="local", z_scope="local"): """ n = 4 # dp4a requires operands packed by 4 - x = te.placeholder((n,), name="x", dtype="int8") - y = te.placeholder((n,), name="y", dtype="int8") + result_dtype = "int32" if dtypes[1] == "int8" else "uint32" + + x = te.placeholder((n,), name="x", dtype=dtypes[0]) + y = te.placeholder((n,), name="y", dtype=dtypes[1]) k = te.reduce_axis((0, n), name="rc") - z = te.compute((1,), lambda i: te.sum(x[k].astype("int32") * y[k].astype("int32"), axis=[k])) + z = te.compute( + (1,), lambda i: te.sum(x[k].astype(result_dtype) * y[k].astype(result_dtype), axis=[k]) + ) def _intrin_func(ins, outs): def _instr(index): xx, yy = ins zz = outs[0] + zz_dtype = zz.dtype if index == 1: - return zz.vstore(0, 0) + return zz.vstore(0, tvm.tir.const(0, zz_dtype)) ib = tvm.tir.ir_builder.create() - vec_x = xx.vload(0, dtype="int8x4") - vec_y = yy.vload(0, dtype="int8x4") + vec_x_dtype = "int8x4" if xx.dtype == "int8" else "uint8x4" + vec_y_dtype = "int8x4" if yy.dtype == "int8" else "uint8x4" + + vec_x = xx.vload(0, dtype=vec_x_dtype) + vec_y = yy.vload(0, dtype=vec_y_dtype) prev_z = 0 if index == 0 else zz.vload(0) - new_z = tvm.tir.call_pure_extern("int32", "__dp4a", vec_x, vec_y, prev_z) + new_z = tvm.tir.call_pure_extern(zz_dtype, "__dp4a", vec_x, vec_y, prev_z) ib.emit(zz.vstore(0, new_z)) return ib.get() From c30a2bc5fe7d5fed730c41f72908ef72c4ccf1d7 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Sun, 6 Mar 2022 19:04:42 +0900 Subject: [PATCH 04/10] add back dense and group conv2d change --- python/tvm/topi/cuda/dense.py | 12 ++++++++---- python/tvm/topi/cuda/group_conv2d_nchw.py | 11 +++++++---- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/python/tvm/topi/cuda/dense.py b/python/tvm/topi/cuda/dense.py index c3c5b6e7cf85..862e7b5bc59d 100644 --- a/python/tvm/topi/cuda/dense.py +++ b/python/tvm/topi/cuda/dense.py @@ -17,6 +17,7 @@ # pylint: disable=invalid-name, unused-argument """Schedule for dense operator""" import logging +import tvm from tvm import te import tvm.autotvm as autotvm from tvm.contrib import cublas @@ -133,9 +134,6 @@ def _callback(op): return s -_dp4a = dp4a("shared", "shared", "local") - - def _schedule_dense_int8(cfg, s, output): data, weight = s[output].op.input_tensors if len(weight.op.input_tensors) == 1 and weight.op.input_tensors[0] == data: @@ -173,7 +171,13 @@ def _schedule_dense_int8(cfg, s, output): ko = CC.op.reduce_axis[0] ko, ki = s[CC].split(ko, factor=4) ko, kt = cfg["tile_k"].apply(s, CC, ko) - s[CC].tensorize(ki, _dp4a) + target = tvm.target.Target.current(allow_none=False) + do_tensorize = True + if "vulkan" in target.keys: + do_tensorize = "+dotprod" in target.mattr or target.supports_integer_dot_product + if do_tensorize: + dtypes = (data.dtype, weight.dtype) + s[CC].tensorize(ki, dp4a("shared", "shared", "local", dtypes)) by, vy, ty, yi = cfg["tile_y"].apply(s, output, n) bx, vx, tx, xi = cfg["tile_x"].apply(s, output, x) diff --git a/python/tvm/topi/cuda/group_conv2d_nchw.py b/python/tvm/topi/cuda/group_conv2d_nchw.py index d75cfffc1af8..f786b0d8d647 100644 --- a/python/tvm/topi/cuda/group_conv2d_nchw.py +++ b/python/tvm/topi/cuda/group_conv2d_nchw.py @@ -396,9 +396,6 @@ def _callback(op): return s -_dp4a = dp4a("shared", "shared", "local") - - def _schedule_group_conv2d_NCHWc_int8(cfg, s, output): """Schedule group conv2d int8 NCHWc template""" workload = output.op.attrs["workload"] @@ -509,7 +506,13 @@ def _schedule_group_conv2d_NCHWc_int8(cfg, s, output): s[conv].reorder(rco, ryo, rxo, rci, ryi, rxi, n, f, y, x, c, rc_block) _, rc_block = s[conv].split(rc_block, factor=4) - s[conv].tensorize(rc_block, _dp4a) + target = tvm.target.Target.current(allow_none=False) + do_tensorize = True + if "vulkan" in target.keys: + do_tensorize = "+dotprod" in target.mattr or target.supports_integer_dot_product + if do_tensorize: + dtypes = (pad_data.dtype, packed_kernel.dtype) + s[conv].tensorize(rc_block, dp4a("shared", "shared", "local", dtypes)) s[AA].compute_at(s[conv], rxo) s[WW].compute_at(s[conv], rxo) From 0d3b16d94e482234280e8538f8e2664af7ac3065 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Mon, 7 Mar 2022 13:37:37 +0900 Subject: [PATCH 05/10] add back test_topi_conv2d_int8.py change (but don't test on vk) --- .../topi/python/test_topi_conv2d_int8.py | 220 +++++++++--------- 1 file changed, 116 insertions(+), 104 deletions(-) diff --git a/tests/python/topi/python/test_topi_conv2d_int8.py b/tests/python/topi/python/test_topi_conv2d_int8.py index 7b36c8a5c1b8..8c310506fa95 100644 --- a/tests/python/topi/python/test_topi_conv2d_int8.py +++ b/tests/python/topi/python/test_topi_conv2d_int8.py @@ -14,6 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +# """Example code to do convolution.""" import numpy as np @@ -32,6 +33,7 @@ from common import Int8Fallback import tvm.testing +import pytest def compile_conv2d_NHWC_gemm_int8_arm( @@ -226,6 +228,7 @@ def check_target(target): def verify_conv2d_NCHWc_int8( + in_dtype, batch, in_channel, in_size, @@ -246,12 +249,15 @@ def verify_conv2d_NCHWc_int8( in_height = in_width = in_size - A = te.placeholder((batch, in_channel, in_height, in_width), name="A", dtype="int8") - W = te.placeholder((num_filter, in_channel, kernel, kernel), name="W", dtype="int8") + A = te.placeholder((batch, in_channel, in_height, in_width), name="A", dtype=in_dtype) + W = te.placeholder((num_filter, in_channel, kernel, kernel), name="W", dtype=in_dtype) a_shape = get_const_tuple(A.shape) w_shape = get_const_tuple(W.shape) dtype = A.dtype + out_dtype = "int32" if in_dtype == "int8" else "uint32" + lo = -128 if in_dtype == "int8" else 0 + hi = 127 if in_dtype == "int8" else 255 def check_target(target, compute, schedule, oc_block_factor): dev = tvm.device(target, 0) @@ -263,17 +269,19 @@ def check_target(target, compute, schedule, oc_block_factor): return bias = te.placeholder( - (num_filter // oc_block_factor, 1, 1, oc_block_factor), name="bias", dtype="int32" + (num_filter // oc_block_factor, 1, 1, oc_block_factor), name="bias", dtype=out_dtype ) bias_shape = get_const_tuple(bias.shape) @memoize("topi.tests.test_topi_conv2d_int8.verify_conv2d_nchw") def get_ref_data(): - a_np = np.random.randint(low=-128, high=127, size=a_shape).astype("int32") - w_np = np.random.randint(low=-128, high=128, size=w_shape).astype("int32") - b_np = np.random.uniform(size=bias_shape).astype("int32") + a_np = np.random.randint(low=lo, high=hi, size=a_shape).astype(out_dtype) + w_np = np.random.randint(low=lo, high=hi, size=w_shape).astype(out_dtype) + b_np = np.random.uniform(size=bias_shape).astype(out_dtype) dw_np = tvm.topi.testing.dilate_python(w_np, (1, 1, dilation, dilation)) - c_np = tvm.topi.testing.conv2d_nchw_python(a_np, dw_np, stride, padding).astype("int32") + c_np = tvm.topi.testing.conv2d_nchw_python(a_np, dw_np, stride, padding).astype( + out_dtype + ) # convert to NCHWc _, _, out_height, out_width = c_np.shape @@ -282,7 +290,7 @@ def get_ref_data(): ).transpose(0, 1, 3, 4, 2) if add_bias: - b_np = np.random.uniform(size=bias_shape).astype("int32") + b_np = np.random.uniform(size=bias_shape).astype(out_dtype) c_np += b_np if add_relu: c_np = np.maximum(c_np, 0) @@ -301,7 +309,7 @@ def get_ref_data(): (dilation, dilation), "NCHW", "NCHW", - "int32", + out_dtype, ) print(C.shape) print(bias.shape) @@ -313,7 +321,7 @@ def get_ref_data(): a = tvm.nd.array(a_np.astype(dtype), dev) w = tvm.nd.array(w_np.astype(dtype), dev) - b = tvm.nd.array(b_np.astype("int32"), dev) + b = tvm.nd.array(b_np.astype(out_dtype), dev) c = tvm.nd.array(np.zeros(get_const_tuple(C.shape), dtype=C.dtype), dev) if add_bias: tvm.build( @@ -375,6 +383,7 @@ def get_ref_data(): def verify_conv2d_nchw_int8( + in_dtype, batch, in_channel, in_size, @@ -395,9 +404,9 @@ def verify_conv2d_nchw_int8( in_height = in_width = in_size - A = te.placeholder((batch, in_channel, in_height, in_width), name="A", dtype="int8") - W = te.placeholder((num_filter, in_channel, kernel, kernel), name="W", dtype="int8") - bias = te.placeholder((num_filter, 1, 1), name="bias", dtype="int8") + A = te.placeholder((batch, in_channel, in_height, in_width), name="A", dtype=in_dtype) + W = te.placeholder((num_filter, in_channel, kernel, kernel), name="W", dtype=in_dtype) + bias = te.placeholder((num_filter, 1, 1), name="bias", dtype=in_dtype) a_shape = get_const_tuple(A.shape) w_shape = get_const_tuple(W.shape) @@ -495,114 +504,118 @@ def check_target(target): check_target(target) -@tvm.testing.requires_cuda -def test_conv2d_nchw(): +@pytest.mark.parametrize("in_dtype", ["int8", "uint8"]) +def test_conv2d_nchw(in_dtype): with Int8Fallback(): # ResNet18 workloads where channels in / out are multiple of oc_block_factor - verify_conv2d_NCHWc_int8(1, 64, 56, 64, 3, 1, 1) - verify_conv2d_NCHWc_int8(1, 64, 56, 64, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 64, 56, 128, 3, 2, 1) - verify_conv2d_NCHWc_int8(1, 64, 56, 128, 1, 2, 0) - verify_conv2d_NCHWc_int8(1, 128, 28, 128, 3, 1, 1) - verify_conv2d_NCHWc_int8(1, 128, 28, 256, 3, 2, 1) - verify_conv2d_NCHWc_int8(1, 128, 28, 256, 1, 2, 0) - verify_conv2d_NCHWc_int8(1, 256, 14, 256, 3, 1, 1) - verify_conv2d_NCHWc_int8(1, 256, 14, 512, 3, 2, 1) - verify_conv2d_NCHWc_int8(1, 256, 14, 512, 1, 2, 0) - verify_conv2d_NCHWc_int8(1, 512, 7, 512, 3, 1, 1) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 56, 64, 3, 1, 1) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 56, 64, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 56, 128, 3, 2, 1) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 56, 128, 1, 2, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 128, 28, 128, 3, 1, 1) + verify_conv2d_NCHWc_int8(in_dtype, 1, 128, 28, 256, 3, 2, 1) + verify_conv2d_NCHWc_int8(in_dtype, 1, 128, 28, 256, 1, 2, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 256, 14, 256, 3, 1, 1) + verify_conv2d_NCHWc_int8(in_dtype, 1, 256, 14, 512, 3, 2, 1) + verify_conv2d_NCHWc_int8(in_dtype, 1, 256, 14, 512, 1, 2, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 512, 7, 512, 3, 1, 1) # bias, relu - verify_conv2d_NCHWc_int8(1, 64, 56, 64, 3, 1, 1, add_relu=True) - verify_conv2d_NCHWc_int8(1, 64, 56, 64, 3, 1, 1, add_bias=True) - verify_conv2d_NCHWc_int8(1, 64, 56, 64, 3, 1, 1, add_bias=True, add_relu=True) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 56, 64, 3, 1, 1, add_relu=True) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 56, 64, 3, 1, 1, add_bias=True) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 56, 64, 3, 1, 1, add_bias=True, add_relu=True) # dilation = 2 - verify_conv2d_NCHWc_int8(1, 64, 56, 64, 3, 1, 1, dilation=2) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 56, 64, 3, 1, 1, dilation=2) # batch size - verify_conv2d_NCHWc_int8(4, 64, 56, 64, 3, 1, 1) - verify_conv2d_NCHWc_int8(9, 64, 56, 64, 3, 1, 1) + verify_conv2d_NCHWc_int8(in_dtype, 4, 64, 56, 64, 3, 1, 1) + verify_conv2d_NCHWc_int8(in_dtype, 9, 64, 56, 64, 3, 1, 1) # weird workloads - verify_conv2d_NCHWc_int8(4, 4, 4, 8, 4, 4, 4) + verify_conv2d_NCHWc_int8(in_dtype, 4, 4, 4, 8, 4, 4, 4) # inception v3 workloads where channels in / out are multiple of oc_block_factor - verify_conv2d_NCHWc_int8(1, 32, 149, 32, 3, 1, 0) - verify_conv2d_NCHWc_int8(1, 32, 147, 64, 3, 1, 1) - verify_conv2d_NCHWc_int8(1, 64, 73, 80, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 80, 73, 192, 3, 1, 0) - verify_conv2d_NCHWc_int8(1, 192, 35, 64, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 192, 35, 48, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 48, 35, 64, 5, 1, 2) - verify_conv2d_NCHWc_int8(1, 64, 35, 96, 3, 1, 1) - verify_conv2d_NCHWc_int8(1, 96, 35, 96, 3, 1, 1) - verify_conv2d_NCHWc_int8(1, 192, 35, 32, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 256, 35, 64, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 256, 35, 48, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 288, 35, 64, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 288, 35, 48, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 288, 35, 384, 3, 2, 0) - verify_conv2d_NCHWc_int8(1, 96, 35, 96, 3, 2, 0) - verify_conv2d_NCHWc_int8(1, 768, 17, 192, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 768, 17, 128, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 128, 17, 128, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 128, 17, 192, 7, 1, 3) - verify_conv2d_NCHWc_int8(1, 128, 17, 128, 7, 1, 3) - verify_conv2d_NCHWc_int8(1, 128, 17, 192, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 768, 17, 160, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 160, 17, 160, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 160, 17, 192, 7, 1, 3) - verify_conv2d_NCHWc_int8(1, 160, 17, 160, 7, 1, 3) - verify_conv2d_NCHWc_int8(1, 160, 17, 192, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 192, 17, 192, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 192, 17, 192, 7, 1, 3) - verify_conv2d_NCHWc_int8(1, 192, 17, 320, 3, 2, 0) - verify_conv2d_NCHWc_int8(1, 192, 17, 192, 3, 2, 0) - verify_conv2d_NCHWc_int8(1, 1280, 8, 320, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 1280, 8, 384, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 384, 8, 384, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 384, 8, 384, 3, 1, 1) - verify_conv2d_NCHWc_int8(1, 1280, 8, 448, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 448, 8, 384, 3, 1, 1) - verify_conv2d_NCHWc_int8(1, 1280, 8, 192, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 2048, 8, 320, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 2048, 8, 384, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 2048, 8, 448, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 2048, 8, 192, 1, 1, 0) - verify_conv2d_NCHWc_int8(1, 1024, 19, 88, 3, 1, 1) + verify_conv2d_NCHWc_int8(in_dtype, 1, 32, 149, 32, 3, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 32, 147, 64, 3, 1, 1) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 73, 80, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 80, 73, 192, 3, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 192, 35, 64, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 192, 35, 48, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 48, 35, 64, 5, 1, 2) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 35, 96, 3, 1, 1) + verify_conv2d_NCHWc_int8(in_dtype, 1, 96, 35, 96, 3, 1, 1) + verify_conv2d_NCHWc_int8(in_dtype, 1, 192, 35, 32, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 256, 35, 64, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 256, 35, 48, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 288, 35, 64, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 288, 35, 48, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 288, 35, 384, 3, 2, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 96, 35, 96, 3, 2, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 768, 17, 192, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 768, 17, 128, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 128, 17, 128, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 128, 17, 192, 7, 1, 3) + verify_conv2d_NCHWc_int8(in_dtype, 1, 128, 17, 128, 7, 1, 3) + verify_conv2d_NCHWc_int8(in_dtype, 1, 128, 17, 192, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 768, 17, 160, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 160, 17, 160, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 160, 17, 192, 7, 1, 3) + verify_conv2d_NCHWc_int8(in_dtype, 1, 160, 17, 160, 7, 1, 3) + verify_conv2d_NCHWc_int8(in_dtype, 1, 160, 17, 192, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 192, 17, 192, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 192, 17, 192, 7, 1, 3) + verify_conv2d_NCHWc_int8(in_dtype, 1, 192, 17, 320, 3, 2, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 192, 17, 192, 3, 2, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 1280, 8, 320, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 1280, 8, 384, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 384, 8, 384, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 384, 8, 384, 3, 1, 1) + verify_conv2d_NCHWc_int8(in_dtype, 1, 1280, 8, 448, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 448, 8, 384, 3, 1, 1) + verify_conv2d_NCHWc_int8(in_dtype, 1, 1280, 8, 192, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 2048, 8, 320, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 2048, 8, 384, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 2048, 8, 448, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 2048, 8, 192, 1, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 1, 1024, 19, 88, 3, 1, 1) # batch > 1 - verify_conv2d_NCHWc_int8(7, 32, 149, 32, 3, 1, 0) - verify_conv2d_NCHWc_int8(8, 32, 149, 32, 3, 1, 0) - verify_conv2d_NCHWc_int8(32, 32, 149, 32, 3, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 7, 32, 149, 32, 3, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 8, 32, 149, 32, 3, 1, 0) + verify_conv2d_NCHWc_int8(in_dtype, 32, 32, 149, 32, 3, 1, 0) # Asymmetric padding - verify_conv2d_NCHWc_int8(1, 32, 35, 64, 7, 2, (0, 0, 1, 1)) - verify_conv2d_NCHWc_int8(1, 64, 8, 128, 3, 1, (3, 3, 2, 2)) - verify_conv2d_NCHWc_int8(1, 64, 8, 64, 1, 1, (1, 2, 2, 1)) - verify_conv2d_NCHWc_int8(1, 64, 17, 192, 1, 1, (1, 2)) - verify_conv2d_NCHWc_int8(1, 64, 8, 64, 3, 1, (3, 1)) - verify_conv2d_NCHWc_int8(1, 128, 8, 384, 3, 1, (0, 2)) - verify_conv2d_NCHWc_int8(1, 64, 8, 64, 1, 1, "VALID") - verify_conv2d_NCHWc_int8(1, 392, 8, 64, 3, 1, "VALID") - verify_conv2d_NCHWc_int8(1, 512, 19, 64, 1, 1, "SAME") - verify_conv2d_NCHWc_int8(1, 64, 16, 32, 2, 1, "SAME") - verify_conv2d_NCHWc_int8(1, 64, 8, 64, 3, 1, (1, 2, 2, 1), add_relu=True) - verify_conv2d_NCHWc_int8(1, 64, 8, 64, 5, 2, (1, 3), add_bias=True) - verify_conv2d_NCHWc_int8(1, 64, 56, 64, 3, 1, "VALID", add_bias=True, add_relu=True) - verify_conv2d_NCHWc_int8(1, 64, 56, 64, 24, 1, "SAME", add_bias=True, add_relu=True) + verify_conv2d_NCHWc_int8(in_dtype, 1, 32, 35, 64, 7, 2, (0, 0, 1, 1)) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 8, 128, 3, 1, (3, 3, 2, 2)) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 8, 64, 1, 1, (1, 2, 2, 1)) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 17, 192, 1, 1, (1, 2)) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 8, 64, 3, 1, (3, 1)) + verify_conv2d_NCHWc_int8(in_dtype, 1, 128, 8, 384, 3, 1, (0, 2)) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 8, 64, 1, 1, "VALID") + verify_conv2d_NCHWc_int8(in_dtype, 1, 392, 8, 64, 3, 1, "VALID") + verify_conv2d_NCHWc_int8(in_dtype, 1, 512, 19, 64, 1, 1, "SAME") + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 16, 32, 2, 1, "SAME") + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 8, 64, 3, 1, (1, 2, 2, 1), add_relu=True) + verify_conv2d_NCHWc_int8(in_dtype, 1, 64, 8, 64, 5, 2, (1, 3), add_bias=True) + verify_conv2d_NCHWc_int8( + in_dtype, 1, 64, 56, 64, 3, 1, "VALID", add_bias=True, add_relu=True + ) + verify_conv2d_NCHWc_int8( + in_dtype, 1, 64, 56, 64, 24, 1, "SAME", add_bias=True, add_relu=True + ) # Conv2d NCHW int8 schedule testing. Internally, it uses NCHWc schedule. So, just # performing basic testing - one test for all different scenarios - batch, dilation etc.. - verify_conv2d_nchw_int8(1, 64, 56, 64, 3, 1, 1) - verify_conv2d_nchw_int8(1, 64, 56, 64, 3, 1, 1, add_relu=True) - verify_conv2d_nchw_int8(1, 64, 56, 64, 3, 1, 1, dilation=2) - verify_conv2d_nchw_int8(9, 64, 56, 64, 3, 1, 1) - verify_conv2d_nchw_int8(4, 4, 4, 4, 4, 4, 4) - verify_conv2d_nchw_int8(1, 32, 149, 32, 3, 1, 0) - verify_conv2d_nchw_int8(7, 32, 149, 32, 3, 1, 0) - verify_conv2d_nchw_int8(1, 32, 35, 64, 7, 2, (0, 0, 1, 1)) - verify_conv2d_nchw_int8(1, 32, 35, 64, 7, 2, (0, 0, 2, 2)) + verify_conv2d_nchw_int8(in_dtype, 1, 64, 56, 64, 3, 1, 1) + verify_conv2d_nchw_int8(in_dtype, 1, 64, 56, 64, 3, 1, 1, add_relu=True) + verify_conv2d_nchw_int8(in_dtype, 1, 64, 56, 64, 3, 1, 1, dilation=2) + verify_conv2d_nchw_int8(in_dtype, 9, 64, 56, 64, 3, 1, 1) + verify_conv2d_nchw_int8(in_dtype, 4, 4, 4, 4, 4, 4, 4) + verify_conv2d_nchw_int8(in_dtype, 1, 32, 149, 32, 3, 1, 0) + verify_conv2d_nchw_int8(in_dtype, 7, 32, 149, 32, 3, 1, 0) + verify_conv2d_nchw_int8(in_dtype, 1, 32, 35, 64, 7, 2, (0, 0, 1, 1)) + verify_conv2d_nchw_int8(in_dtype, 1, 32, 35, 64, 7, 2, (0, 0, 2, 2)) def test_conv2d_nhwc(): @@ -646,5 +659,4 @@ def test_conv2d_nhwc(): if __name__ == "__main__": - test_conv2d_nchw() - test_conv2d_nhwc() + sys.exit(pytest.main(sys.argv)) From aa7328f09a8f0f30d4fbb1cfb8d0455748bb9547 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Mon, 7 Mar 2022 14:18:10 +0900 Subject: [PATCH 06/10] check dot prod availablity in batch matmul schedule --- python/tvm/topi/cuda/batch_matmul.py | 11 ++++++++++- tests/python/topi/python/test_topi_batch_matmul.py | 2 +- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/python/tvm/topi/cuda/batch_matmul.py b/python/tvm/topi/cuda/batch_matmul.py index 85dcd8b1add5..ede1187a3e35 100644 --- a/python/tvm/topi/cuda/batch_matmul.py +++ b/python/tvm/topi/cuda/batch_matmul.py @@ -368,7 +368,16 @@ def _schedule_batch_matmul_int8(cfg, s, output): ko, ki = s[batch_matmul_cache].split(ko, factor=4) ko, kt = cfg["tile_k"].apply(s, batch_matmul_cache, ko) # dp4a tensorize - s[batch_matmul_cache].tensorize(ki, _dp4a) + + target = tvm.target.Target.current(allow_none=False) + do_tensorize = True + + if "vulkan" in target.keys: + do_tensorize = "+dotprod" in target.mattr or target.supports_integer_dot_product + + if do_tensorize: + dtypes = (input_x.dtype, input_y.dtype) + s[batch_matmul_cache].tensorize(ki, dp4a("shared", "shared", "local", dtypes)) # tile axis f, m, n = batch_matmul_op.axis diff --git a/tests/python/topi/python/test_topi_batch_matmul.py b/tests/python/topi/python/test_topi_batch_matmul.py index 9bd9dd286b1a..4c5df5fb5e67 100644 --- a/tests/python/topi/python/test_topi_batch_matmul.py +++ b/tests/python/topi/python/test_topi_batch_matmul.py @@ -128,7 +128,7 @@ def check_device(device): f(a, b, c) tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-5) - for device in ["cuda"]: + for device in ["cuda", "vulkan -from_device=0"]: check_device(device) From 1bb9afd77238d2306206c669c8e87c965c533f91 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Mon, 7 Mar 2022 17:55:54 +0900 Subject: [PATCH 07/10] do not run uint8 tensorization on arm --- .../topi/python/test_topi_conv2d_int8.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/tests/python/topi/python/test_topi_conv2d_int8.py b/tests/python/topi/python/test_topi_conv2d_int8.py index 8c310506fa95..c2af1d74c86f 100644 --- a/tests/python/topi/python/test_topi_conv2d_int8.py +++ b/tests/python/topi/python/test_topi_conv2d_int8.py @@ -370,14 +370,19 @@ def get_ref_data(): lambda a, w, s, p, d, l, ol, o: topi.cuda.conv2d_NCHWc_int8(a, w, s, p, d, l, o), topi.cuda.schedule_conv2d_NCHWc_int8, 4, - ), - ( - "llvm -device arm_cpu -mtriple aarch64-linux-gnu -mattr=+neon", - topi.arm_cpu.conv2d_NCHWc_int8, - topi.arm_cpu.schedule_conv2d_NCHWc_int8, - 8, - ), + ) ] + + if in_dtype == "int8": + targets.append( + ( + "llvm -device arm_cpu -mtriple aarch64-linux-gnu -mattr=+neon", + topi.arm_cpu.conv2d_NCHWc_int8, + topi.arm_cpu.schedule_conv2d_NCHWc_int8, + 8, + ) + ) + for target, compute, schedule, oc_block_factor in targets: check_target(target, compute, schedule, oc_block_factor) From 427009a4876ff1fca5409f341f62fa5bacea2706 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Mon, 7 Mar 2022 18:11:57 +0900 Subject: [PATCH 08/10] add vulkan target to conv2d int8 test but comment out on CI --- tests/python/topi/python/test_topi_conv2d_int8.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/python/topi/python/test_topi_conv2d_int8.py b/tests/python/topi/python/test_topi_conv2d_int8.py index c2af1d74c86f..46257ed387a1 100644 --- a/tests/python/topi/python/test_topi_conv2d_int8.py +++ b/tests/python/topi/python/test_topi_conv2d_int8.py @@ -370,7 +370,14 @@ def get_ref_data(): lambda a, w, s, p, d, l, ol, o: topi.cuda.conv2d_NCHWc_int8(a, w, s, p, d, l, o), topi.cuda.schedule_conv2d_NCHWc_int8, 4, - ) + ), + # Disable on CI since it does not support spirv int8 dot product + # ( + # "vulkan -from_device=0", + # lambda a, w, s, p, d, l, ol, o: topi.cuda.conv2d_NCHWc_int8(a, w, s, p, d, l, o), + # topi.cuda.schedule_conv2d_NCHWc_int8, + # 4, + # ), ] if in_dtype == "int8": From 839d2338a3a750278d77ece3de57c94559015420 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Mon, 7 Mar 2022 18:30:39 +0900 Subject: [PATCH 09/10] do not run vk batch matmul test --- tests/python/topi/python/test_topi_batch_matmul.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/python/topi/python/test_topi_batch_matmul.py b/tests/python/topi/python/test_topi_batch_matmul.py index 4c5df5fb5e67..9bd9dd286b1a 100644 --- a/tests/python/topi/python/test_topi_batch_matmul.py +++ b/tests/python/topi/python/test_topi_batch_matmul.py @@ -128,7 +128,7 @@ def check_device(device): f(a, b, c) tvm.testing.assert_allclose(c.numpy(), c_np, rtol=1e-5) - for device in ["cuda", "vulkan -from_device=0"]: + for device in ["cuda"]: check_device(device) From b5ebe6994fd91b3454c3beae4873e0d4a7e3724e Mon Sep 17 00:00:00 2001 From: Mei Ye Date: Tue, 8 Mar 2022 00:38:38 +0000 Subject: [PATCH 10/10] Fix performance regression due to missing of constant folding in the index expression. --- src/target/spirv/codegen_spirv.cc | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/target/spirv/codegen_spirv.cc b/src/target/spirv/codegen_spirv.cc index 1bb457a8b913..4f875e955576 100644 --- a/src/target/spirv/codegen_spirv.cc +++ b/src/target/spirv/codegen_spirv.cc @@ -452,7 +452,8 @@ spirv::Value CodeGenSPIRV::VisitExpr_(const BufferLoadNode* op) { if (desired_read_type == info.element_type) { // Requested a single value from an array. This may be a scalar load // or a vectorized load, based on the array element type. - spirv::Value index = MakeValue(prim_index); + PrimExpr vec_index = analyzer_->Simplify(prim_index); + spirv::Value index = MakeValue(vec_index); spirv::Value ptr = builder_->StructArrayAccess(ptr_type, buffer, index); spirv::Value loaded = builder_->MakeValue(spv::OpLoad, content_type, ptr, mask); // OpTypeBool have no physical address/storage. Here, cast from