e2e microkernel pipeline + argmax ukernel on ROCM backend. (#15943)
This commit presents is an end-to-end ukernel pipeline on ROCm with argmax as the first ukernel. It is
end-to-end in a because we automatically generate and link the bitcode, which means end users do not need to manually link bitcodes or hand-modify the kernel to have ukernel.generics.
This work is based on Ukernel lowerings on CPU written by Benoit Jacob
as well as a ukernel on CUDA sample by Mahesh Ravishankar.
diff --git a/compiler/plugins/target/ROCM/BUILD.bazel b/compiler/plugins/target/ROCM/BUILD.bazel
index 3186984..ce37323 100644
--- a/compiler/plugins/target/ROCM/BUILD.bazel
+++ b/compiler/plugins/target/ROCM/BUILD.bazel
@@ -29,12 +29,15 @@
deps = [
"//compiler/src/iree/compiler/Codegen/Dialect/Codegen/IR:IREECodegenDialect",
"//compiler/src/iree/compiler/Codegen/LLVMGPU",
+ "//compiler/src/iree/compiler/Codegen/Utils",
"//compiler/src/iree/compiler/Dialect/HAL/Target",
+ "//compiler/src/iree/compiler/Dialect/HAL/Target:LLVMLinkerUtils",
"//compiler/src/iree/compiler/PluginAPI",
"//compiler/src/iree/compiler/Utils",
"//runtime/src/iree/schemas:rocm_executable_def_c_fbs",
"@llvm-project//llvm:AMDGPUCodeGen",
"@llvm-project//llvm:Analysis",
+ "@llvm-project//llvm:BitWriter",
"@llvm-project//llvm:Core",
"@llvm-project//llvm:IPO",
"@llvm-project//llvm:IRReader",
diff --git a/compiler/plugins/target/ROCM/CMakeLists.txt b/compiler/plugins/target/ROCM/CMakeLists.txt
index c5cc9f4..400a07f 100644
--- a/compiler/plugins/target/ROCM/CMakeLists.txt
+++ b/compiler/plugins/target/ROCM/CMakeLists.txt
@@ -28,6 +28,7 @@
DEPS
LLVMAMDGPUCodeGen
LLVMAnalysis
+ LLVMBitWriter
LLVMCore
LLVMIRReader
LLVMLinker
@@ -52,7 +53,9 @@
MLIRTargetLLVMIRExport
iree::compiler::Codegen::Dialect::Codegen::IR::IREECodegenDialect
iree::compiler::Codegen::LLVMGPU
+ iree::compiler::Codegen::Utils
iree::compiler::Dialect::HAL::Target
+ iree::compiler::Dialect::HAL::Target::LLVMLinkerUtils
iree::compiler::PluginAPI
iree::compiler::Utils
iree::schemas::rocm_executable_def_c_fbs
diff --git a/compiler/plugins/target/ROCM/ROCMTarget.cpp b/compiler/plugins/target/ROCM/ROCMTarget.cpp
index 1b57579..19700a5 100644
--- a/compiler/plugins/target/ROCM/ROCMTarget.cpp
+++ b/compiler/plugins/target/ROCM/ROCMTarget.cpp
@@ -11,12 +11,14 @@
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
#include "iree/compiler/Codegen/LLVMGPU/Passes.h"
+#include "iree/compiler/Dialect/HAL/Target/LLVMLinkerUtils.h"
#include "iree/compiler/Dialect/HAL/Target/TargetRegistry.h"
#include "iree/compiler/PluginAPI/Client.h"
#include "iree/compiler/Utils/FlatbufferUtils.h"
#include "iree/compiler/Utils/ToolUtils.h"
#include "iree/schemas/rocm_executable_def_builder.h"
#include "llvm/Analysis/TargetTransformInfo.h"
+#include "llvm/Bitcode/BitcodeWriter.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Module.h"
@@ -47,6 +49,7 @@
bool linkBitcode = false;
std::string bitcodeDirectory;
int wavesPerEu = 0;
+ std::string enableROCMUkernels = "none";
void bindOptions(OptionsBinder &binder) {
static llvm::cl::OptionCategory category("ROCM HAL Target");
@@ -62,10 +65,27 @@
llvm::cl::cat(category),
llvm::cl::desc("Optimization hint specifying minimum "
"number of waves per execution unit"));
+ binder.opt<std::string>(
+ "iree-rocm-enable-ukernels", enableROCMUkernels,
+ llvm::cl::cat(category),
+ llvm::cl::desc(
+ "Enables microkernels in the llvmcpu backend. May be "
+ "`default`, `none`, `all`, or a comma-separated list of "
+ "specific unprefixed microkernels to enable, e.g. `mmt4d`."));
}
};
} // namespace
+static void dumpModuleToPath(StringRef path, StringRef baseName,
+ StringRef suffix, StringRef extension,
+ llvm::Module &module) {
+ llvm::SmallVector<char, 0> data;
+ llvm::raw_svector_ostream ostream(data);
+ module.print(ostream, nullptr);
+ dumpDataToPath(path, baseName, suffix, extension,
+ StringRef(data.data(), data.size()));
+}
+
static std::string translateModuleToObj(llvm::Module &module,
llvm::TargetMachine &targetMachine) {
std::string targetObj;
@@ -309,6 +329,23 @@
iree_compiler::FlatbufferBuilder builder;
iree_hal_rocm_ExecutableDef_start_as_root(builder);
+ // Link user modules and libdevice (if required).
+ // Note that linking order matters:
+ llvm::Linker linker(*llvmModule);
+ if (failed(linkCmdlineBitcodeFiles(
+ variantOp.getLoc(), linker, llvm::Linker::OverrideFromSrc,
+ *targetMachine, llvmModule->getContext()))) {
+ return failure();
+ }
+
+ if (!options.enableROCMUkernels.empty() ||
+ options.enableROCMUkernels != "none") {
+ auto enabledUkernelsStr = StringRef(options.enableROCMUkernels);
+ linkUkernelBCFiles(llvmModule.get(), variantOp.getLoc(),
+ enabledUkernelsStr, options.targetChip,
+ options.bitcodeDirectory,
+ llvm::Linker::OverrideFromSrc, *targetMachine);
+ }
// Link module to Device Library
if (options.linkBitcode) {
if (options.bitcodeDirectory.empty()) {
@@ -320,9 +357,19 @@
linkROCDLIfNecessary(llvmModule.get(), options.targetChip,
options.bitcodeDirectory);
}
+ if (!serOptions.dumpIntermediatesPath.empty()) {
+ dumpModuleToPath(serOptions.dumpIntermediatesPath,
+ serOptions.dumpBaseName, variantOp.getName(),
+ ".linked.ll", *llvmModule);
+ }
// Add Optimize module
optimizeModule(*llvmModule, *targetMachine);
-
+ // Store optimized ll.
+ if (!serOptions.dumpIntermediatesPath.empty()) {
+ dumpModuleToPath(serOptions.dumpIntermediatesPath,
+ serOptions.dumpBaseName, variantOp.getName(),
+ ".optimized.ll", *llvmModule);
+ }
// Serialize hsaco kernel into the binary that we will embed in the
// final FlatBuffer.
std::unique_ptr<llvm::Module> moduleCopy;
@@ -406,6 +453,8 @@
// Set target arch
addConfig("target_arch", StringAttr::get(context, options.targetChip));
+ addConfig("ukernels", StringAttr::get(context, options.enableROCMUkernels));
+
auto configAttr = b.getDictionaryAttr(configItems);
return IREE::HAL::ExecutableTargetAttr::get(
context, b.getStringAttr("rocm"), b.getStringAttr("rocm-hsaco-fb"),
diff --git a/compiler/plugins/target/ROCM/ROCMTargetUtils.cpp b/compiler/plugins/target/ROCM/ROCMTargetUtils.cpp
index 51b1e2f..7798693 100644
--- a/compiler/plugins/target/ROCM/ROCMTargetUtils.cpp
+++ b/compiler/plugins/target/ROCM/ROCMTargetUtils.cpp
@@ -6,6 +6,8 @@
#include "./ROCMTargetUtils.h"
+#include "iree/compiler/Codegen/Utils/GPUUtils.h"
+#include "iree/compiler/Dialect/HAL/Target/LLVMLinkerUtils.h"
#include "iree/compiler/Utils/ToolUtils.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Module.h"
@@ -92,6 +94,48 @@
return success();
}
+LogicalResult linkPathBitcodeFiles(Location loc, llvm::Linker &linker,
+ unsigned linkerFlags, StringRef path,
+ llvm::TargetMachine &targetMachine,
+ llvm::LLVMContext &context) {
+ auto bitcodeBufferRef = llvm::MemoryBuffer::getFile(path);
+ if (auto ec = bitcodeBufferRef.getError()) {
+ return mlir::emitError(loc) << "failed reading user bitcode file `" << path
+ << "`: " << ec.message();
+ }
+ auto setAlwaysInline = [&](llvm::Module &module) {
+ if (targetMachine.getTargetCPU().contains("gfx10") ||
+ targetMachine.getTargetCPU().contains("gfx11")) {
+ // some ROCM/HIP functions for gfx10 or gfx11 has accuracy issue if
+ // inlined.
+ return;
+ }
+ for (auto &func : module.getFunctionList()) {
+ // Some ROCM/HIP builtin functions have Optnone and NoInline for default.
+ if (targetMachine.getTargetTriple().isAMDGCN()) {
+ if (func.hasFnAttribute(llvm::Attribute::OptimizeNone)) {
+ func.removeFnAttr(llvm::Attribute::OptimizeNone);
+ }
+ if (targetMachine.getTargetTriple().isAMDGCN() &&
+ func.hasFnAttribute(llvm::Attribute::NoInline)) {
+ func.removeFnAttr(llvm::Attribute::NoInline);
+ }
+ }
+ func.addFnAttr(llvm::Attribute::AlwaysInline);
+ }
+ };
+ if (failed(linkBitcodeModule(
+ loc, linker, linkerFlags, targetMachine, path,
+ llvm::parseBitcodeFile(*bitcodeBufferRef->get(), context),
+ setAlwaysInline))) {
+ return mlir::emitError(loc) << "failed linking in user bitcode file `"
+ << path << "` for target triple '"
+ << targetMachine.getTargetTriple().str() << "'";
+ }
+
+ return success();
+}
+
static std::vector<std::string> getROCDLPaths(std::string targetChip,
std::string bitCodeDir) {
// AMDGPU bitcodes.
@@ -107,6 +151,36 @@
return result;
}
+static std::vector<std::string> getUkernelPaths(StringRef enabledUkernelsStr,
+ StringRef targetChip,
+ StringRef bitCodeDir) {
+ // AMD bitcodes.
+ std::vector<std::string> selectedUkernelNames;
+ if (enabledUkernelsStr == "all") {
+ const char *allUkernelNames[] = {"argmax"};
+ size_t numUkernels = sizeof(allUkernelNames) / sizeof(allUkernelNames[0]);
+ for (int i = 0; i < numUkernels; i++) {
+ selectedUkernelNames.push_back(allUkernelNames[i]);
+ }
+ } else {
+ while (!enabledUkernelsStr.empty()) {
+ auto split = enabledUkernelsStr.split(',');
+ selectedUkernelNames.push_back(split.first.str());
+ enabledUkernelsStr = split.second;
+ }
+ }
+
+ // Construct full path to ROCDL bitcode libraries.
+ std::vector<std::string> result;
+ std::string app = "/";
+ for (auto &kernelName : selectedUkernelNames) {
+ std::string filename =
+ "rocm_" + kernelName + "_ukernel_" + targetChip.str();
+ result.push_back(bitCodeDir.str() + app + filename + ".bc");
+ }
+ return result;
+}
+
static void overridePlatformGlobal(llvm::Module *module, StringRef globalName,
uint32_t newValue, llvm::Type *globalTy) {
// NOTE: the global will not be defined if it is not used in the module.
@@ -180,6 +254,26 @@
};
}
+// Links optimized Ukernel bitcodes into the given module if the module needs
+// it.
+void linkUkernelBCFiles(llvm::Module *module, Location loc,
+ StringRef enabledUkernelsStr, StringRef targetChip,
+ StringRef bitCodeDir, unsigned linkerFlags,
+ llvm::TargetMachine &targetMachine) {
+ // Early exit if Ukernel not supported on target chip.
+ if (!iree_compiler::hasUkernelSupportedRocmArch(targetChip))
+ return;
+ std::vector<std::string> ukernelPaths =
+ getUkernelPaths(enabledUkernelsStr, targetChip, bitCodeDir);
+ llvm::Linker linker(*module);
+ for (auto &path : ukernelPaths) {
+ if (failed(linkPathBitcodeFiles(loc, linker, linkerFlags, StringRef(path),
+ targetMachine, module->getContext()))) {
+ llvm::WithColor::error(llvm::errs()) << "Fail to Link Ukernel.\n";
+ }
+ }
+}
+
//===========Link LLVM Module to ROCDL End===================/
//=====================Create HSACO Begin=============//
diff --git a/compiler/plugins/target/ROCM/ROCMTargetUtils.h b/compiler/plugins/target/ROCM/ROCMTargetUtils.h
index 39bb89f..171fded 100644
--- a/compiler/plugins/target/ROCM/ROCMTargetUtils.h
+++ b/compiler/plugins/target/ROCM/ROCMTargetUtils.h
@@ -9,6 +9,7 @@
#include "iree/compiler/Dialect/HAL/Target/TargetBackend.h"
#include "llvm/IR/Module.h"
+#include "llvm/Target/TargetMachine.h"
namespace mlir::iree_compiler::IREE::HAL {
@@ -16,9 +17,17 @@
void linkROCDLIfNecessary(llvm::Module *module, std::string targetChip,
std::string bitCodeDir);
+// Links optimized Ukernel module.
+void linkUkernelBCFiles(llvm::Module *module, Location loc,
+ StringRef enabledUkernelsStr, StringRef targetChip,
+ StringRef bitCodeDir, unsigned linkerFlags,
+ llvm::TargetMachine &targetMachine);
// Compiles ISAToHsaco Code
std::string createHsaco(Location loc, const std::string isa, StringRef name);
+// Returns true if the rocm archtecture target is supported for ukernels.
+bool hasUkernelSupportedRocmArch(IREE::HAL::ExecutableTargetAttr targetAttr);
+
} // namespace mlir::iree_compiler::IREE::HAL
#endif // IREE_COMPILER_PLUGINS_TARGET_ROCM_ROCMTARGETUTILS_H_
diff --git a/compiler/plugins/target/ROCM/builtins/CMakeLists.txt b/compiler/plugins/target/ROCM/builtins/CMakeLists.txt
new file mode 100644
index 0000000..02d4284
--- /dev/null
+++ b/compiler/plugins/target/ROCM/builtins/CMakeLists.txt
@@ -0,0 +1,16 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+iree_cc_library(
+ NAME
+ defs
+ INCLUDES
+ ${CMAKE_CURRENT_SOURCE_DIR}/../..
+ ${CMAKE_CURRENT_BINARY_DIR}/../..
+ PUBLIC
+)
+
+iree_add_all_subdirs()
diff --git a/compiler/plugins/target/ROCM/builtins/ukernel/CMakeLists.txt b/compiler/plugins/target/ROCM/builtins/ukernel/CMakeLists.txt
new file mode 100644
index 0000000..da3a8d9
--- /dev/null
+++ b/compiler/plugins/target/ROCM/builtins/ukernel/CMakeLists.txt
@@ -0,0 +1,174 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+if (NOT IREE_TARGET_BACKEND_ROCM)
+ return()
+endif()
+
+# Check if HIP is installed on system.
+# HIP is required to compile ukernels.
+if (NOT IREE_ROCM_PATH)
+ set(IREE_ROCM_PATH "/opt/rocm")
+endif()
+set (IREE_ROCM_VERSION "${IREE_ROCM_PATH}/.info/version")
+if (NOT EXISTS ${IREE_ROCM_VERSION})
+ message(STATUS
+ "hip runtime cannot be found in ${IREE_ROCM_PATH}.
+ Please try setting IREE_ROCM_PATH to rocm directory.
+ Ukernel will not be compiled.")
+ return()
+endif()
+
+
+iree_add_all_subdirs()
+
+set(_platform_lib_reldir "iree_platform_libs/rocm")
+set(_device_bc_path "${IREE_COMPILER_DYLIB_DIR}/iree_platform_libs/rocm")
+set (_amd_ukernel_libs)
+set (_amd_ukernel_targets)
+function(iree_rocm_bitcode_library)
+ cmake_parse_arguments(
+ _RULE
+ ""
+ "NAME;OUT;ROCM_ARCH"
+ "SRCS;COPTS"
+ ${ARGN}
+ )
+
+ if(DEFINED _RULE_OUT)
+ set(_OUT "${_RULE_OUT}")
+ else()
+ set(_OUT "${_RULE_NAME}_${_RULE_ROCM_ARCH}.bc")
+ endif()
+
+ set(_ROCM_ARCH "${_RULE_ROCM_ARCH}")
+ set (OPT_FLAG "-O0")
+ if (_ROCM_ARCH MATCHES "GFX9")
+ set (OPT_FLAG "-O3")
+ endif()
+ set(_COPTS
+ "-x" "hip"
+ # Target architecture.
+ "--offload-arch=${_ROCM_ARCH}"
+
+ # Suppress warnings about missing path to rocm lib,
+ # and benign warning about ROCM version.
+ "-nogpulib"
+ "-D_ALLOW_COMPILER_AND_STL_VERSION_MISMATCH"
+ # Avoid error caused by clang bundling.
+ "--offload-device-only"
+ # Only enable necessary optimizations S.T we can use -O3.
+ "-Xclang" "-disable-llvm-optzns"
+ "${OPT_FLAG}"
+
+ # Object file only in bitcode format:
+ "-c"
+ "-emit-llvm"
+ )
+
+ set(_BITCODE_FILES)
+ foreach(_SRC ${_RULE_SRCS})
+ get_filename_component(_BITCODE_SRC_PATH "${_SRC}" REALPATH)
+ set(_BITCODE_FILE "${_RULE_NAME}_${_SRC}_${_ROCM_ARCH}.bc")
+ list(APPEND _BITCODE_FILES ${_BITCODE_FILE})
+ add_custom_command(
+ OUTPUT
+ "${_BITCODE_FILE}"
+ COMMAND
+ "${IREE_CLANG_BINARY}"
+ ${_COPTS}
+ "${_BITCODE_SRC_PATH}"
+ "-o"
+ "${_BITCODE_FILE}"
+ DEPENDS
+ "${IREE_CLANG_BINARY}"
+ "${_SRC}"
+ COMMENT
+ "Compiling ${_SRC} to ${_BITCODE_FILE}"
+ VERBATIM
+ )
+ endforeach()
+
+ add_custom_command(
+ OUTPUT
+ "${_OUT}"
+ COMMAND
+ ${IREE_LLVM_LINK_BINARY}
+ ${_BITCODE_FILES}
+ "-o"
+ "${_OUT}"
+ DEPENDS
+ ${IREE_LLVM_LINK_BINARY}
+ ${_BITCODE_FILES}
+ COMMENT
+ "Linking bitcode to ${_OUT}"
+ VERBATIM
+ )
+ # Only add iree_${NAME} as custom target doesn't support aliasing to
+ # iree::${NAME}.
+ iree_package_name(_PACKAGE_NAME)
+ add_custom_target("${_PACKAGE_NAME}_${_RULE_NAME}_${_ROCM_ARCH}"
+ DEPENDS "${_OUT}"
+ )
+ set(_amd_ukernel_libs ${_amd_ukernel_libs} ${_OUT} PARENT_SCOPE)
+ set(_amd_ukernel_targets ${_amd_ukernel_targets} "${_PACKAGE_NAME}_${_RULE_NAME}_${_ROCM_ARCH}" PARENT_SCOPE)
+endfunction()
+
+# TODO: Decide what to build by default. No real constaints here
+# except compile-time cost, so just picked out the popular ones.
+set(_ukernel_supported_chips "gfx90a" "gfx940" "gfx1030" "gfx1100")
+foreach(_amd_chip ${_ukernel_supported_chips})
+ iree_rocm_bitcode_library(
+ NAME
+ rocm_argmax_ukernel
+ ROCM_ARCH
+ ${_amd_chip}
+ SRCS
+ "argmax_ukernel.c"
+ )
+endforeach()
+
+# Copy UKernel into platform dir.
+set(_all_ukernel_bc_copy_commands)
+set(_all_ukernel_bc_files)
+set(_ukernel_lib_srcdir ${CMAKE_CURRENT_BINARY_DIR})
+foreach(_amd_ukernel_name ${_amd_ukernel_libs})
+ # Copy to lib/ tree.
+ set(_ukernel_bc_srcpath "${_ukernel_lib_srcdir}/${_amd_ukernel_name}")
+ set(_ukernel_bc_relpath "${_platform_lib_reldir}/${_amd_ukernel_name}")
+ list(APPEND _all_ukernel_bc_files "${IREE_COMPILER_DYLIB_DIR}/${_ukernel_bc_relpath}")
+ list(APPEND _all_ukernel_bc_deps "${_ukernel_bc_path}")
+ list(APPEND _all_ukernel_bc_copy_commands
+ COMMAND ${CMAKE_COMMAND} -E copy
+ "${_ukernel_bc_srcpath}"
+ "${IREE_COMPILER_DYLIB_DIR}/${_ukernel_bc_relpath}"
+ )
+ # Note this bc file as being part of the bundle that must be included with
+ # the compiler dylib.
+ set_property(GLOBAL APPEND PROPERTY IREE_COMPILER_DYLIB_RELPATHS "${_ukernel_bc_relpath}")
+endforeach()
+
+# Generate a custom target with all file level dependencies and commands to
+# copy to our build tree locations.
+# Our GenDeviceLibs target depends on all of the defined device lib targets.
+add_custom_command(
+ OUTPUT ${_all_ukernel_bc_files}
+ DEPENDS ${_amd_ukernel_targets}
+ POST_BUILD
+ ${_all_ukernel_bc_copy_commands}
+)
+
+add_custom_target(iree_builtin_ROCM_UkernelDeviceLibs
+ DEPENDS
+ ${_all_ukernel_bc_files}
+)
+
+# Ensure that the device libs are built when the compiler dylib is built.
+set_property(GLOBAL APPEND PROPERTY IREE_COMPILER_DYLIB_DEPENDS
+ iree_builtin_ROCM_UkernelDeviceLibs)
+
+# Install.
+install(FILES ${_all_ukernel_bc_files}
+ DESTINATION "${IREE_COMPILER_DYLIB_INSTALL_PREFIX}/${_platform_lib_reldir}")
diff --git a/compiler/plugins/target/ROCM/builtins/ukernel/argmax_ukernel.c b/compiler/plugins/target/ROCM/builtins/ukernel/argmax_ukernel.c
new file mode 100644
index 0000000..9c2caf3
--- /dev/null
+++ b/compiler/plugins/target/ROCM/builtins/ukernel/argmax_ukernel.c
@@ -0,0 +1,198 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <float.h>
+#include <hip/hip_fp16.h>
+#include <hip/hip_runtime.h>
+
+extern "C" __device__ __attribute__((const)) half __ockl_wfred_max_f16(half);
+extern "C" __device__ __attribute__((const))
+int64_t __ockl_wfred_min_i64(int64_t);
+extern "C" __device__ __attribute__((const))
+int32_t __ockl_wfred_min_i32(int32_t);
+
+/*
+Constraint/Tiling note:
+For simplicity, we distribute all parallel dim across different workgroup, and
+only use single subgroup/warp per workgroup. This constraint is also set during
+tiling phase in KernelConfig.
+*/
+
+extern "C" __device__ void __iree_uk_rocm_argmax_F32I32(float *inputBuffer,
+ size_t input_offset,
+ int32_t *outputBuffer,
+ size_t output_offset,
+ size_t reductionSize) {
+ uint laneID = __builtin_amdgcn_workitem_id_x();
+ // Set identity value to handle problem non divisible by subgroupSize.
+ float laneMax =
+ laneID >= reductionSize ? -FLT_MAX : inputBuffer[input_offset + laneID];
+ int32_t laneResult = laneID;
+
+ // NOTE: On F32 kernels with clang, reductionSize/blockDim.x has numerical
+ // inaccuracy.
+ uint numBatches = (reductionSize + warpSize - 1) / warpSize;
+ for (int i = 1; i < numBatches; ++i) {
+ uint idx = warpSize * i + laneID;
+ float newIn =
+ idx >= reductionSize ? -FLT_MAX : inputBuffer[input_offset + idx];
+ if (newIn == laneMax)
+ continue;
+ laneMax = __ocml_fmax_f32(newIn, laneMax);
+ laneResult = newIn == laneMax ? idx : laneResult;
+ }
+
+ // Final reduction with one subgroup
+ // NOTE: __ockl_wfred_max_f32 has correctness issue on gfx1100 documented on
+ // https://github.com/openxla/iree/issues/16112.
+ float wgMax = laneMax;
+ for (int i = 1; i < warpSize; i *= 2) {
+ wgMax = __ocml_fmax_f32(__shfl_xor(wgMax, i), wgMax);
+ }
+ // Check if there are multiple max value holders.
+ uint64_t laneHasMaxValmask = __ballot(wgMax == laneMax);
+ // if there is only one max value holder, write and exit.
+ if (__popcll(laneHasMaxValmask) == 1) {
+ if (wgMax == laneMax)
+ outputBuffer[output_offset] = laneResult;
+ return;
+ }
+ // if there are multiple max value holder, find smallest index (argmax
+ // semantics).
+ int32_t indexVal = wgMax == laneMax ? laneResult : __INT32_MAX__;
+ laneResult = __ockl_wfred_min_i32(indexVal);
+ if (laneID == 0)
+ outputBuffer[output_offset] = laneResult;
+}
+
+extern "C" __device__ void __iree_uk_rocm_argmax_F32I64(float *inputBuffer,
+ size_t input_offset,
+ int64_t *outputBuffer,
+ size_t output_offset,
+ size_t reductionSize) {
+ uint laneID = __builtin_amdgcn_workitem_id_x();
+ // Set identity value to handle problem non divisible by subgroupSize.
+ float laneMax =
+ laneID >= reductionSize ? -FLT_MAX : inputBuffer[input_offset + laneID];
+ int64_t laneResult = laneID;
+
+ // NOTE: On F32 kernels with clang, reductionSize/blockDim.x has numerical
+ // inaccuracy.
+ uint numBatches = (reductionSize + warpSize - 1) / warpSize;
+ for (int i = 1; i < numBatches; ++i) {
+ uint idx = warpSize * i + laneID;
+ float newIn =
+ idx >= reductionSize ? -FLT_MAX : inputBuffer[input_offset + idx];
+ if (newIn == laneMax)
+ continue;
+ laneMax = __ocml_fmax_f32(newIn, laneMax);
+ laneResult = newIn == laneMax ? idx : laneResult;
+ }
+
+ // Final reduction with one subgroup
+ // NOTE: __ockl_wfred_max_f32 has correctness issue on gfx1100 documented on
+ // https://github.com/openxla/iree/issues/16112.
+ float wgMax = laneMax;
+ for (int i = 1; i < warpSize; i *= 2) {
+ wgMax = __ocml_fmax_f32(__shfl_xor(wgMax, i), wgMax);
+ }
+ // Check if there are multiple max value holders.
+ uint64_t laneHasMaxValmask = __ballot(wgMax == laneMax);
+ // if there is only one max value holder, write and exit.
+ if (__popcll(laneHasMaxValmask) == 1) {
+ if (wgMax == laneMax)
+ outputBuffer[output_offset] = laneResult;
+ return;
+ }
+ // if there are multiple max value holder, find smallest index (argmax
+ // semantics).
+ int32_t indexVal = wgMax == laneMax ? laneResult : __INT64_MAX__;
+ laneResult = __ockl_wfred_min_i64(indexVal);
+ if (laneID == 0)
+ outputBuffer[output_offset] = laneResult;
+}
+
+extern "C" __device__ void __iree_uk_rocm_argmax_F16I32(half *inputBuffer,
+ size_t input_offset,
+ int32_t *outputBuffer,
+ size_t output_offset,
+ size_t reductionSize) {
+ half NEG_F16_MAX = __float2half(-65504.0f);
+ uint laneID = __builtin_amdgcn_workitem_id_x();
+ // Set identity value to handle problem non divisible by subgroupSize.
+ half laneMax = laneID >= reductionSize ? NEG_F16_MAX
+ : inputBuffer[input_offset + laneID];
+ int32_t laneResult = laneID;
+
+ uint numBatches = (reductionSize + warpSize - 1) / warpSize;
+ for (int i = 1; i < numBatches; ++i) {
+ uint idx = warpSize * i + laneID;
+ half newIn =
+ idx >= reductionSize ? NEG_F16_MAX : inputBuffer[input_offset + idx];
+ if (newIn == laneMax)
+ continue;
+ laneMax = __ocml_fmax_f16(newIn, laneMax);
+ laneResult = newIn == laneMax ? idx : laneResult;
+ }
+
+ // Final reduction with one subgroup
+ half wgMax = __ockl_wfred_max_f16(laneMax);
+ // Check if there are multiple max value holders.
+ uint64_t laneHasMaxValmask = __ballot(wgMax == laneMax);
+ // if there is only one max value holder, write and exit.
+ if (__popcll(laneHasMaxValmask) == 1) {
+ if (wgMax == laneMax)
+ outputBuffer[output_offset] = laneResult;
+ return;
+ }
+ // if there are multiple max value holder, find smallest index (argmax
+ // semantics).
+ int32_t indexVal = wgMax == laneMax ? laneResult : __INT32_MAX__;
+ laneResult = __ockl_wfred_min_i32(indexVal);
+ if (laneID == 0)
+ outputBuffer[output_offset] = laneResult;
+}
+
+extern "C" __device__ void __iree_uk_rocm_argmax_F16I64(half *inputBuffer,
+ size_t input_offset,
+ int64_t *outputBuffer,
+ size_t output_offset,
+ size_t reductionSize) {
+ half NEG_F16_MAX = __float2half(-65504.0f);
+ uint laneID = __builtin_amdgcn_workitem_id_x();
+ // Set identity value to handle problem non divisible by subgroupSize.
+ half laneMax = laneID >= reductionSize ? NEG_F16_MAX
+ : inputBuffer[input_offset + laneID];
+ int64_t laneResult = laneID;
+
+ uint numBatches = (reductionSize + warpSize - 1) / warpSize;
+ for (int i = 1; i < numBatches; ++i) {
+ uint idx = warpSize * i + laneID;
+ half newIn =
+ idx >= reductionSize ? NEG_F16_MAX : inputBuffer[input_offset + idx];
+ if (newIn == laneMax)
+ continue;
+ laneMax = __ocml_fmax_f16(newIn, laneMax);
+ laneResult = newIn == laneMax ? idx : laneResult;
+ }
+
+ // Final reduction with one subgroup
+ half wgMax = __ockl_wfred_max_f16(laneMax);
+ // Check if there are multiple max value holders.
+ uint64_t laneHasMaxValmask = __ballot(wgMax == laneMax);
+ // if there is only one max value holder, write and exit.
+ if (__popcll(laneHasMaxValmask) == 1) {
+ if (wgMax == laneMax)
+ outputBuffer[output_offset] = laneResult;
+ return;
+ }
+ // if there are multiple max value holder, find smallest index (argmax
+ // semantics).
+ int32_t indexVal = wgMax == laneMax ? laneResult : __INT64_MAX__;
+ laneResult = __ockl_wfred_min_i64(indexVal);
+ if (laneID == 0)
+ outputBuffer[output_offset] = laneResult;
+}
diff --git a/compiler/plugins/target/ROCM/builtins/ukernel/test/CMakeLists.txt b/compiler/plugins/target/ROCM/builtins/ukernel/test/CMakeLists.txt
new file mode 100644
index 0000000..e18c17e
--- /dev/null
+++ b/compiler/plugins/target/ROCM/builtins/ukernel/test/CMakeLists.txt
@@ -0,0 +1,21 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+iree_add_all_subdirs()
+
+iree_lit_test_suite(
+ NAME
+ lit
+ SRCS
+ "argmax_linking.mlir"
+ TOOLS
+ ${IREE_LLD_TARGET}
+ FileCheck
+ iree-compile
+ LABELS
+ "driver=rocm"
+ "hostonly"
+)
diff --git a/compiler/plugins/target/ROCM/builtins/ukernel/test/argmax_linking.mlir b/compiler/plugins/target/ROCM/builtins/ukernel/test/argmax_linking.mlir
new file mode 100644
index 0000000..7907af5
--- /dev/null
+++ b/compiler/plugins/target/ROCM/builtins/ukernel/test/argmax_linking.mlir
@@ -0,0 +1,201 @@
+// RUN: [[ $IREE_ROCM_DISABLE == 1 ]] || iree-compile --split-input-file --iree-hal-target-backends=rocm --iree-rocm-enable-ukernels=all --iree-rocm-target-chip=gfx1100 --compile-to=executable-targets %s | FileCheck %s
+
+// We want to check that uKernel is indeed generated from e2e workflow.
+
+// CHECK: llvm.func @__iree_uk_rocm_argmax_F32I64
+// CHECK: llvm.call @__iree_uk_rocm_argmax_F32I64
+func.func @argmax_1d_f32i64(%arg0: tensor<1x?xf32>) -> tensor<1x1xi64> {
+ %c0 = arith.constant 0 : index
+ %c0_i64 = arith.constant 0 : i64
+ %cst = arith.constant 0xFF800000 : f32
+ %cst_0 = arith.constant 0.000000e+00 : f16
+ %c-1_i64 = arith.constant -1 : i64
+ %15 = tensor.empty() : tensor<1xi64>
+ %16 = linalg.fill ins(%c0_i64 : i64) outs(%15 : tensor<1xi64>) -> tensor<1xi64>
+ %17 = tensor.empty() : tensor<1xf32>
+ %18 = linalg.fill ins(%cst : f32) outs(%17 : tensor<1xf32>) -> tensor<1xf32>
+ %19:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<1x?xf32>) outs(%18, %16 : tensor<1xf32>, tensor<1xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_2: i64):
+ %20 = linalg.index 1 : index
+ %21 = arith.index_cast %20 : index to i64
+ %22 = arith.maximumf %in, %out : f32
+ %23 = arith.cmpf ogt, %in, %out : f32
+ %24 = arith.select %23, %21, %out_2 : i64
+ linalg.yield %22, %24 : f32, i64
+ } -> (tensor<1xf32>, tensor<1xi64>)
+ %expanded_1 = tensor.expand_shape %19#1 [[0, 1]] : tensor<1xi64> into tensor<1x1xi64>
+ return %expanded_1 : tensor<1x1xi64>
+}
+
+// -----
+
+// CHECK: llvm.func @__iree_uk_rocm_argmax_F16I64
+// CHECK: llvm.call @__iree_uk_rocm_argmax_F16I64
+func.func @argmax_1d_f16i64(%arg0: tensor<1x?xf16>) -> tensor<1x1xi64> {
+ %c0 = arith.constant 0 : index
+ %c0_i64 = arith.constant 0 : i64
+ %cst = arith.constant 0xFC00 : f16
+ %cst_0 = arith.constant 0.000000e+00 : f16
+ %c-1_i64 = arith.constant -1 : i64
+ %15 = tensor.empty() : tensor<1xi64>
+ %16 = linalg.fill ins(%c0_i64 : i64) outs(%15 : tensor<1xi64>) -> tensor<1xi64>
+ %17 = tensor.empty() : tensor<1xf16>
+ %18 = linalg.fill ins(%cst : f16) outs(%17 : tensor<1xf16>) -> tensor<1xf16>
+ %19:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<1x?xf16>) outs(%18, %16 : tensor<1xf16>, tensor<1xi64>) {
+ ^bb0(%in: f16, %out: f16, %out_2: i64):
+ %20 = linalg.index 1 : index
+ %21 = arith.index_cast %20 : index to i64
+ %22 = arith.maximumf %in, %out : f16
+ %23 = arith.cmpf ogt, %in, %out : f16
+ %24 = arith.select %23, %21, %out_2 : i64
+ linalg.yield %22, %24 : f16, i64
+ } -> (tensor<1xf16>, tensor<1xi64>)
+ %expanded_1 = tensor.expand_shape %19#1 [[0, 1]] : tensor<1xi64> into tensor<1x1xi64>
+ return %expanded_1 : tensor<1x1xi64>
+}
+
+
+// -----
+
+// CHECK: llvm.func @__iree_uk_rocm_argmax_F32I64
+// CHECK: llvm.call @__iree_uk_rocm_argmax_F32I64
+func.func @argmax_2d_f32i64(%arg0: tensor<16x?xf32>) -> tensor<16x1xi64> {
+ %c0 = arith.constant 0 : index
+ %c0_i64 = arith.constant 0 : i64
+ %cst = arith.constant 0xFF800000 : f32
+ %cst_0 = arith.constant 0.000000e+00 : f16
+ %c-1_i64 = arith.constant -1 : i64
+ %15 = tensor.empty() : tensor<16xi64>
+ %16 = linalg.fill ins(%c0_i64 : i64) outs(%15 : tensor<16xi64>) -> tensor<16xi64>
+ %17 = tensor.empty() : tensor<16xf32>
+ %18 = linalg.fill ins(%cst : f32) outs(%17 : tensor<16xf32>) -> tensor<16xf32>
+ %19:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<16x?xf32>) outs(%18, %16 : tensor<16xf32>, tensor<16xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_2: i64):
+ %20 = linalg.index 1 : index
+ %21 = arith.index_cast %20 : index to i64
+ %22 = arith.maximumf %in, %out : f32
+ %23 = arith.cmpf ogt, %in, %out : f32
+ %24 = arith.select %23, %21, %out_2 : i64
+ linalg.yield %22, %24 : f32, i64
+ } -> (tensor<16xf32>, tensor<16xi64>)
+ %expanded_1 = tensor.expand_shape %19#1 [[0, 1]] : tensor<16xi64> into tensor<16x1xi64>
+ return %expanded_1 : tensor<16x1xi64>
+}
+
+// -----
+
+// CHECK: llvm.func @__iree_uk_rocm_argmax_F32I32
+// CHECK: llvm.call @__iree_uk_rocm_argmax_F32I32
+#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
+#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>
+func.func @argmax_3d_dyn_f32i32(%arg0: tensor<?x?x?xf32>) -> tensor<?x?xf32> {
+ %c0 = arith.constant 0 : index
+ %c0_i32 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0xFF800000 : f32
+ %cst_0 = arith.constant 0.000000e+00 : f16
+ %dim = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
+ %dim_1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
+ %0 = tensor.empty(%dim, %dim_1) : tensor<?x?xi32>
+ %1 = linalg.fill ins(%c0_i32 : i32) outs(%0 : tensor<?x?xi32>) -> tensor<?x?xi32>
+ %2 = tensor.empty(%dim, %dim_1) : tensor<?x?xf32>
+ %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0 : tensor<?x?x?xf32>) outs(%3, %1 : tensor<?x?xf32>, tensor<?x?xi32>) {
+ ^bb0(%in: f32, %out: f32, %out_2: i32):
+ %6 = linalg.index 2 : index
+ %7 = arith.index_cast %6 : index to i32
+ %8 = arith.maximumf %in, %out : f32
+ %9 = arith.cmpf ogt, %in, %out : f32
+ %10 = arith.select %9, %7, %out_2 : i32
+ linalg.yield %8, %10 : f32, i32
+ } -> (tensor<?x?xf32>, tensor<?x?xi32>)
+ %5 = arith.sitofp %4#1 : tensor<?x?xi32> to tensor<?x?xf32>
+ return %5 : tensor<?x?xf32>
+}
+
+// -----
+
+// CHECK: llvm.func @__iree_uk_rocm_argmax_F32I64
+// CHECK: llvm.call @__iree_uk_rocm_argmax_F32I64
+func.func @argmax_3d_dyn_f32i64(%arg0: tensor<?x?x?xf32>) -> tensor<?x?xf32> {
+ %c0 = arith.constant 0 : index
+ %c0_i64 = arith.constant 0 : i64
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0xFF800000 : f32
+ %cst_0 = arith.constant 0.000000e+00 : f16
+ %dim = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
+ %dim_1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
+ %0 = tensor.empty(%dim, %dim_1) : tensor<?x?xi64>
+ %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<?x?xi64>) -> tensor<?x?xi64>
+ %2 = tensor.empty(%dim, %dim_1) : tensor<?x?xf32>
+ %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?x?xf32>) -> tensor<?x?xf32>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0 : tensor<?x?x?xf32>) outs(%3, %1 : tensor<?x?xf32>, tensor<?x?xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_2: i64):
+ %6 = linalg.index 2 : index
+ %7 = arith.index_cast %6 : index to i64
+ %8 = arith.maximumf %in, %out : f32
+ %9 = arith.cmpf ogt, %in, %out : f32
+ %10 = arith.select %9, %7, %out_2 : i64
+ linalg.yield %8, %10 : f32, i64
+ } -> (tensor<?x?xf32>, tensor<?x?xi64>)
+ %5 = arith.sitofp %4#1 : tensor<?x?xi64> to tensor<?x?xf32>
+ return %5 : tensor<?x?xf32>
+}
+
+// -----
+
+// CHECK: llvm.func @__iree_uk_rocm_argmax_F16I32
+// CHECK: llvm.call @__iree_uk_rocm_argmax_F16I32
+func.func @argmax_3d_dyn_f16i32(%arg0: tensor<?x?x?xf16>) -> tensor<?x?xf32> {
+ %c0 = arith.constant 0 : index
+ %c0_i32 = arith.constant 0 : i32
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0xFC00 : f16
+ %cst_0 = arith.constant 0.000000e+00 : f16
+ %dim = tensor.dim %arg0, %c0 : tensor<?x?x?xf16>
+ %dim_1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf16>
+ %0 = tensor.empty(%dim, %dim_1) : tensor<?x?xi32>
+ %1 = linalg.fill ins(%c0_i32 : i32) outs(%0 : tensor<?x?xi32>) -> tensor<?x?xi32>
+ %2 = tensor.empty(%dim, %dim_1) : tensor<?x?xf16>
+ %3 = linalg.fill ins(%cst : f16) outs(%2 : tensor<?x?xf16>) -> tensor<?x?xf16>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0 : tensor<?x?x?xf16>) outs(%3, %1 : tensor<?x?xf16>, tensor<?x?xi32>) {
+ ^bb0(%in: f16, %out: f16, %out_2: i32):
+ %6 = linalg.index 2 : index
+ %7 = arith.index_cast %6 : index to i32
+ %8 = arith.maximumf %in, %out : f16
+ %9 = arith.cmpf ogt, %in, %out : f16
+ %10 = arith.select %9, %7, %out_2 : i32
+ linalg.yield %8, %10 : f16, i32
+ } -> (tensor<?x?xf16>, tensor<?x?xi32>)
+ %5 = arith.sitofp %4#1 : tensor<?x?xi32> to tensor<?x?xf32>
+ return %5 : tensor<?x?xf32>
+}
+
+// -----
+
+// CHECK: llvm.func @__iree_uk_rocm_argmax_F16I64
+// CHECK: llvm.call @__iree_uk_rocm_argmax_F16I64
+func.func @argmax_3d_dyn_f16i64(%arg0: tensor<?x?x?xf16>) -> tensor<?x?xf32> {
+ %c0 = arith.constant 0 : index
+ %c0_i64 = arith.constant 0 : i64
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0xFC00 : f16
+ %cst_0 = arith.constant 0.000000e+00 : f16
+ %dim = tensor.dim %arg0, %c0 : tensor<?x?x?xf16>
+ %dim_1 = tensor.dim %arg0, %c1 : tensor<?x?x?xf16>
+ %0 = tensor.empty(%dim, %dim_1) : tensor<?x?xi64>
+ %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<?x?xi64>) -> tensor<?x?xi64>
+ %2 = tensor.empty(%dim, %dim_1) : tensor<?x?xf16>
+ %3 = linalg.fill ins(%cst : f16) outs(%2 : tensor<?x?xf16>) -> tensor<?x?xf16>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>, affine_map<(d0, d1, d2) -> (d0, d1)>], iterator_types = ["parallel", "parallel", "reduction"]} ins(%arg0 : tensor<?x?x?xf16>) outs(%3, %1 : tensor<?x?xf16>, tensor<?x?xi64>) {
+ ^bb0(%in: f16, %out: f16, %out_2: i64):
+ %6 = linalg.index 2 : index
+ %7 = arith.index_cast %6 : index to i64
+ %8 = arith.maximumf %in, %out : f16
+ %9 = arith.cmpf ogt, %in, %out : f16
+ %10 = arith.select %9, %7, %out_2 : i64
+ linalg.yield %8, %10 : f16, i64
+ } -> (tensor<?x?xf16>, tensor<?x?xi64>)
+ %5 = arith.sitofp %4#1 : tensor<?x?xi64> to tensor<?x?xf32>
+ return %5 : tensor<?x?xf32>
+}
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel
index 5c28076..fde7e6d 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/BUILD.bazel
@@ -53,6 +53,7 @@
"GPUDistributeSharedMemoryCopy.cpp",
"GPUDistributionPatterns.cpp",
"GPUGeneralizeNamedOps.cpp",
+ "GPULowerToUKernels.cpp",
"GPUMultiBuffering.cpp",
"GPUPatterns.cpp",
"GPUPipelining.cpp",
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt
index 235a09c..185ba86 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/CMakeLists.txt
@@ -52,6 +52,7 @@
"GPUDistributeSharedMemoryCopy.cpp"
"GPUDistributionPatterns.cpp"
"GPUGeneralizeNamedOps.cpp"
+ "GPULowerToUKernels.cpp"
"GPUMultiBuffering.cpp"
"GPUPatterns.cpp"
"GPUPipelining.cpp"
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPULowerToUKernels.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPULowerToUKernels.cpp
new file mode 100644
index 0000000..a21c26b
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPULowerToUKernels.cpp
@@ -0,0 +1,206 @@
+// Copyright 2023 The IREE Authors
+//
+// Licensed under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "iree-dialects/Dialect/LinalgExt/IR/LinalgExtOps.h"
+#include "iree/compiler/Codegen/Common/GPU/PassDetail.h"
+#include "iree/compiler/Codegen/Common/GPU/Passes.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenDialect.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
+#include "iree/compiler/Codegen/Dialect/Codegen/IR/UKernelOps.h"
+#include "iree/compiler/Codegen/Utils/GPUUtils.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Attributes.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/IR/TypeRange.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace mlir {
+namespace iree_compiler {
+
+namespace {
+class GPULowerToUKernelsPass
+ : public GPULowerToUKernelsBase<GPULowerToUKernelsPass> {
+public:
+ void getDependentDialects(DialectRegistry ®istry) const override {
+ registry.insert<IREE::Codegen::IREECodegenDialect>();
+ }
+
+ void runOnOperation() override;
+};
+} // namespace
+
+/// Holds a function name and attributes.
+struct FnNameAndDefAttrs {
+ std::string name;
+ SmallVector<NamedAttribute> defAttrs;
+};
+
+/// Returns the function name and attributes to use for a ukernel with given
+/// `ukernelName` on the target described by `targetAttr`.
+static FnNameAndDefAttrs
+getFnNameAndDefAttrs(const char *ukernelName, std::string &typeSuffixID,
+ RewriterBase &rewriter,
+ IREE::HAL::ExecutableTargetAttr targetAttr) {
+ FnNameAndDefAttrs result;
+ if (isROCMBackend(targetAttr)) {
+ result.name =
+ std::string("__iree_uk_rocm_") + ukernelName + "_" + typeSuffixID;
+ result.defAttrs.emplace_back(rewriter.getStringAttr("vm.import.module"),
+ rewriter.getStringAttr("rocm"));
+ }
+ return result;
+}
+
+/// Matches generic that represent argmax and check if
+/// we have the ukernel that matches it shape constraint, and types.
+/// If we do, then we convert into iree_codegen.ukernel.argmax operation,
+/// that is later lowered into a call to the microkernel.
+static FailureOr<IREE::Codegen::UKernelOpInterface>
+matchArgmaxDAGForUKernel(RewriterBase &rewriter, linalg::GenericOp op) {
+ auto targetAttr = IREE::HAL::ExecutableTargetAttr::lookup(op);
+ const char ukernelName[] = "argmax";
+ if (!hasUkernel(targetAttr, ukernelName) ||
+ !hasUkernelSupportedGpuArch(targetAttr)) {
+ return failure();
+ }
+
+ // Currently only support argmax where parallel dims are 1.
+ // Tiling pipeline is also set to tile all parallel dims to 1, and
+ // reduction dim to be size of whole reduction problem. Which allow
+ // this constraint to be true for a lot of argmax variances.
+ // TODO: Support multi-row or grid-strided argmax ukernel.
+ SmallVector<int64_t, 4> bounds = op.getStaticLoopRanges();
+ SmallVector<unsigned> parallelDims;
+ op.getParallelDims(parallelDims);
+ int64_t parallelSize = 1;
+ for (int64_t dim : parallelDims) {
+ if (ShapedType::isDynamic(bounds[dim])) {
+ return failure();
+ }
+ parallelSize *= bounds[dim];
+ }
+ if (parallelSize != 1)
+ return failure();
+
+ // Get value/input type.
+ Value input = op.getDpsInputOperand(0)->get();
+ auto inputType = llvm::cast<ShapedType>(input.getType());
+ Type inputElemType = inputType.getElementType();
+ // Only support f16 and f32 values.
+ if (!inputElemType.isF16() && !inputElemType.isF32()) {
+ return failure();
+ }
+
+ // Get index type.
+ Value index = op.getDpsInitOperand(1)->get();
+ auto indexType = llvm::cast<ShapedType>(index.getType());
+ Type indexElemType = indexType.getElementType();
+ // Only support i32 and i64 index.
+ if (!indexElemType.isInteger(32) && !indexElemType.isInteger(64)) {
+ return failure();
+ }
+
+ std::string typeSuffixID = "";
+ if (inputElemType.isF16() && indexElemType.isInteger(32)) {
+ typeSuffixID = "F16I32";
+ } else if (inputElemType.isF16() && indexElemType.isInteger(64)) {
+ typeSuffixID = "F16I64";
+ } else if (inputElemType.isF32() && indexElemType.isInteger(32)) {
+ typeSuffixID = "F32I32";
+ } else if (inputElemType.isF32() && indexElemType.isInteger(64)) {
+ typeSuffixID = "F32I64";
+ } else {
+ return rewriter.notifyMatchFailure(
+ op, "unsupported combination of element types");
+ }
+
+ Location loc = op.getLoc();
+ // Currently only support 1D reduction, where reduc is on fastest dim.
+ // Tiling argmax ukernel is also set to enforce this structure.
+ const int kReductionDim = op.getNumLoops() - 1;
+ Value reductionDimSize =
+ rewriter.create<tensor::DimOp>(loc, input, kReductionDim);
+ auto fn =
+ getFnNameAndDefAttrs(ukernelName, typeSuffixID, rewriter, targetAttr);
+ auto genericMicroKernelOp = rewriter.create<IREE::Codegen::UKernelGenericOp>(
+ loc, indexType, fn.name, ValueRange{input}, index,
+ ValueRange{reductionDimSize},
+ /*fn_def_attrs=*/rewriter.getDictionaryAttr(fn.defAttrs),
+ /*strided_outer_dims=*/rewriter.getIndexAttr(0));
+ return cast<IREE::Codegen::UKernelOpInterface>(
+ genericMicroKernelOp.getOperation());
+}
+
+namespace {
+
+using TargetPredicate = std::function<bool(IREE::HAL::ExecutableTargetAttr)>;
+
+struct LowerArgmaxToUKernelPattern : OpRewritePattern<linalg::GenericOp> {
+ LowerArgmaxToUKernelPattern(MLIRContext *context,
+ TargetPredicate targetPredicate)
+ : OpRewritePattern<linalg::GenericOp>(context),
+ targetPredicate(targetPredicate) {}
+
+ LogicalResult matchAndRewrite(linalg::GenericOp op,
+ PatternRewriter &rewriter) const override {
+ if (targetPredicate &&
+ !targetPredicate(IREE::HAL::ExecutableTargetAttr::lookup(op))) {
+ return failure();
+ }
+ if (failed(isArgmaxOp(op))) {
+ return failure();
+ }
+ FailureOr<IREE::Codegen::UKernelOpInterface> ukernelOp =
+ matchArgmaxDAGForUKernel(rewriter, op);
+ if (failed(ukernelOp)) {
+ return rewriter.notifyMatchFailure(
+ op, "failed to find microkernel op to replace with");
+ }
+ rewriter.replaceAllUsesWith(op.getResults()[1],
+ ukernelOp.value()->getResults());
+ return success();
+ }
+
+ TargetPredicate targetPredicate;
+};
+
+} // namespace
+
+void GPULowerToUKernelsPass::runOnOperation() {
+ MLIRContext *context = &getContext();
+ RewritePatternSet patterns(context);
+ // Enabling a lowering of an op to a microkernel is a trade-off between the
+ // potential performance advantage of a microkernel over pure code generation
+ // for that op, and the potential benefits of fusions. Indeed, once an op
+ // lowered into a microkernel, it will never be fused at any MLIR level.
+ // Since microkernels are linked as bitcode, they will still undergo LTO-like
+ // optimization in their calling contexts, but we shouldn't expect this to
+ // achieve similar results as fusing structured ops.
+
+ // These patterns are unconditionally enabled, because we have strong evidence
+ // that it is difficult for codegen to consistently approach microkernels
+ // performance, and that consideration overrides the benefit of fusions for
+ // these ops.
+ patterns.insert<LowerArgmaxToUKernelPattern>(context, isROCMBackend);
+ if (failed(
+ applyPatternsAndFoldGreedily(getOperation(), std::move(patterns)))) {
+ return signalPassFailure();
+ }
+}
+
+std::unique_ptr<OperationPass<>> createGPULowerToUKernelsPass() {
+ return std::make_unique<GPULowerToUKernelsPass>();
+}
+
+} // namespace iree_compiler
+} // namespace mlir
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.h b/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.h
index 07ca997..baafb2b 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.h
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.h
@@ -131,6 +131,10 @@
// This pass generalizes named Linalg ops that are better off as generics.
std::unique_ptr<OperationPass<func::FuncOp>> createGPUGeneralizeNamedOpsPass();
+/// Pass to lower a sequence of operations to a iree_codegen.ukernel.*
+/// operation.
+std::unique_ptr<OperationPass<>> createGPULowerToUKernelsPass();
+
/// Register Common GPU passes.
void registerCodegenCommonGPUPasses();
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td b/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
index 5b88f4e..879a8b2 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/Passes.td
@@ -37,6 +37,14 @@
let constructor = "mlir::iree_compiler::createGPUGeneralizeNamedOpsPass()";
}
+def GPULowerToUKernels :
+ Pass<"iree-codegen-gpu-lower-to-ukernels", ""> {
+ let summary =
+ "Separate out parts of the IR that lower to a micro-kernel";
+ let constructor =
+ "mlir::iree_compiler::createGPULowerToUKernelsPass()";
+}
+
def GPUMultiBuffering :
Pass<"iree-codegen-gpu-multi-buffering", "func::FuncOp"> {
let summary = "Pass to do multi buffering.";
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel
index 73ddd5c..78db631 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/BUILD.bazel
@@ -22,6 +22,7 @@
"gpu_distribute.mlir",
"gpu_distribute_shared_memory.mlir",
"gpu_generalize_named_ops.mlir",
+ "gpu_lower_to_ukernels.mlir",
"gpu_pipeline.mlir",
"gpu_tensor_alloc.mlir",
"gpu_tensor_tile.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt
index b748cca..011b72c 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/CMakeLists.txt
@@ -18,6 +18,7 @@
"gpu_distribute.mlir"
"gpu_distribute_shared_memory.mlir"
"gpu_generalize_named_ops.mlir"
+ "gpu_lower_to_ukernels.mlir"
"gpu_pipeline.mlir"
"gpu_tensor_alloc.mlir"
"gpu_tensor_tile.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_lower_to_ukernels.mlir b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_lower_to_ukernels.mlir
new file mode 100644
index 0000000..2118681
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_lower_to_ukernels.mlir
@@ -0,0 +1,282 @@
+// RUN: iree-opt --pass-pipeline="builtin.module(func.func(iree-codegen-gpu-lower-to-ukernels,cse,canonicalize))" %s | FileCheck %s
+
+func.func @argmax_2d_f32i64(%arg0 : tensor<1x?xf32>) -> tensor<1xi64> attributes {
+ hal.executable.target = #hal.executable.target<"rocm", "rocm-hsaco-fb", {target_arch = "gfx1100", ukernels = "all"}>
+} {
+ %c0_i64 = arith.constant 0 : i64
+ %cst = arith.constant 0xFF800000 : f32
+ %0 = tensor.empty() : tensor<1xi64>
+ %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<1xi64>) -> tensor<1xi64>
+ %2 = tensor.empty() : tensor<1xf32>
+ %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<1xf32>) -> tensor<1xf32>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<1x?xf32>) outs(%3, %1 : tensor<1xf32>, tensor<1xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_0: i64):
+ %5 = linalg.index 1 : index
+ %6 = arith.index_cast %5 : index to i64
+ %7 = arith.maximumf %in, %out : f32
+ %8 = arith.cmpf ogt, %in, %out : f32
+ %9 = arith.select %8, %6, %out_0 : i64
+ linalg.yield %7, %9 : f32, i64
+ } -> (tensor<1xf32>, tensor<1xi64>)
+ return %4#1 : tensor<1xi64>
+}
+
+// CHECK: func @argmax_2d_f32i64(
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<1x?xf32>
+// CHECK-DAG: %[[C1_index:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C0_i64:.+]] = arith.constant 0
+// CHECK-DAG: %[[FILL:.+]] = linalg.fill ins(%[[C0_i64]]
+// CHECK: %[[MICRO_KERNEL:.+]] = iree_codegen.ukernel.generic "__iree_uk_rocm_argmax_F32I64"
+// CHECK-SAME: ins(%[[ARG0]] :
+// CHECK-SAME: outs(%[[FILL]] :
+// CHECK: return %[[MICRO_KERNEL]]
+
+// -----
+
+func.func @argmax_4d_unit_parallel_f32i64(%arg0 : tensor<1x1x1x?xf32>) -> tensor<1x1x1xi64> attributes {
+ hal.executable.target = #hal.executable.target<"rocm", "rocm-hsaco-fb", {target_arch = "gfx1100", ukernels = "all"}>
+} {
+ %c0_i64 = arith.constant 0 : i64
+ %cst = arith.constant 0xFF800000 : f32
+ %0 = tensor.empty() : tensor<1x1x1xi64>
+ %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<1x1x1xi64>) -> tensor<1x1x1xi64>
+ %2 = tensor.empty() : tensor<1x1x1xf32>
+ %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<1x1x1xf32>) -> tensor<1x1x1xf32>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>, affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel", "reduction"]} ins(%arg0 : tensor<1x1x1x?xf32>) outs(%3, %1 : tensor<1x1x1xf32>, tensor<1x1x1xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_0: i64):
+ %5 = linalg.index 1 : index
+ %6 = arith.index_cast %5 : index to i64
+ %7 = arith.maximumf %in, %out : f32
+ %8 = arith.cmpf ogt, %in, %out : f32
+ %9 = arith.select %8, %6, %out_0 : i64
+ linalg.yield %7, %9 : f32, i64
+ } -> (tensor<1x1x1xf32>, tensor<1x1x1xi64>)
+ return %4#1 : tensor<1x1x1xi64>
+}
+
+// CHECK: func @argmax_4d_unit_parallel_f32i64(
+// CHECK: iree_codegen.ukernel.generic
+// CHECK-NOT: linalg.generic
+
+// -----
+
+func.func @argmax_2d_non_unit_parallel_f32i64(%arg0 : tensor<4x?xf32>) -> tensor<4xi64> attributes {
+ hal.executable.target = #hal.executable.target<"rocm", "rocm-hsaco-fb", {target_arch = "gfx1100", ukernels = "all"}>
+} {
+ %c0_i64 = arith.constant 0 : i64
+ %cst = arith.constant 0xFF800000 : f32
+ %0 = tensor.empty() : tensor<4xi64>
+ %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<4xi64>) -> tensor<4xi64>
+ %2 = tensor.empty() : tensor<4xf32>
+ %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<4xf32>) -> tensor<4xf32>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<4x?xf32>) outs(%3, %1 : tensor<4xf32>, tensor<4xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_0: i64):
+ %5 = linalg.index 1 : index
+ %6 = arith.index_cast %5 : index to i64
+ %7 = arith.maximumf %in, %out : f32
+ %8 = arith.cmpf ogt, %in, %out : f32
+ %9 = arith.select %8, %6, %out_0 : i64
+ linalg.yield %7, %9 : f32, i64
+ } -> (tensor<4xf32>, tensor<4xi64>)
+ return %4#1 : tensor<4xi64>
+}
+
+// CHECK: func @argmax_2d_non_unit_parallel_f32i64(
+// CHECK-NOT: iree_codegen.ukernel.generic
+// CHECK: linalg.generic
+
+// -----
+
+func.func @argmax_2d_dyn_parallel_f32i64(%arg0 : tensor<?x?xf32>) -> tensor<?xi64> attributes {
+ hal.executable.target = #hal.executable.target<"rocm", "rocm-hsaco-fb", {target_arch = "gfx1100", ukernels = "all"}>
+} {
+ %c0 = arith.constant 0 : index
+ %c0_i64 = arith.constant 0 : i64
+ %cst = arith.constant 0xFF800000 : f32
+ %dim = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+ %0 = tensor.empty(%dim) : tensor<?xi64>
+ %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<?xi64>) -> tensor<?xi64>
+ %2 = tensor.empty(%dim) : tensor<?xf32>
+ %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<?xf32>) -> tensor<?xf32>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<?x?xf32>) outs(%3, %1 : tensor<?xf32>, tensor<?xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_0: i64):
+ %5 = linalg.index 1 : index
+ %6 = arith.index_cast %5 : index to i64
+ %7 = arith.maximumf %in, %out : f32
+ %8 = arith.cmpf ogt, %in, %out : f32
+ %9 = arith.select %8, %6, %out_0 : i64
+ linalg.yield %7, %9 : f32, i64
+ } -> (tensor<?xf32>, tensor<?xi64>)
+ return %4#1 : tensor<?xi64>
+}
+
+// CHECK: func @argmax_2d_dyn_parallel_f32i64(
+// CHECK-NOT: iree_codegen.ukernel.generic
+// CHECK: linalg.generic
+
+// -----
+
+func.func @argmax_none_ukernel_enabled(%arg0 : tensor<1x?xf32>) -> tensor<1xi64> attributes {
+ hal.executable.target = #hal.executable.target<"rocm", "rocm-hsaco-fb", {target_arch = "gfx1100", ukernels = "none"}>
+} {
+ %c0_i64 = arith.constant 0 : i64
+ %cst = arith.constant 0xFF800000 : f32
+ %0 = tensor.empty() : tensor<1xi64>
+ %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<1xi64>) -> tensor<1xi64>
+ %2 = tensor.empty() : tensor<1xf32>
+ %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<1xf32>) -> tensor<1xf32>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<1x?xf32>) outs(%3, %1 : tensor<1xf32>, tensor<1xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_0: i64):
+ %5 = linalg.index 1 : index
+ %6 = arith.index_cast %5 : index to i64
+ %7 = arith.maximumf %in, %out : f32
+ %8 = arith.cmpf ogt, %in, %out : f32
+ %9 = arith.select %8, %6, %out_0 : i64
+ linalg.yield %7, %9 : f32, i64
+ } -> (tensor<1xf32>, tensor<1xi64>)
+ return %4#1 : tensor<1xi64>
+}
+
+// CHECK: func @argmax_none_ukernel_enabled(
+// CHECK-NOT: iree_codegen.ukernel.generic
+// CHECK: linalg.generic
+
+// -----
+
+func.func @argmax_only_argmax_ukernel_enabled(%arg0 : tensor<1x?xf32>) -> tensor<1xi64> attributes {
+ hal.executable.target = #hal.executable.target<"rocm", "rocm-hsaco-fb", {target_arch = "gfx90a", ukernels = "argmax"}>
+} {
+ %c0_i64 = arith.constant 0 : i64
+ %cst = arith.constant 0xFF800000 : f32
+ %0 = tensor.empty() : tensor<1xi64>
+ %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<1xi64>) -> tensor<1xi64>
+ %2 = tensor.empty() : tensor<1xf32>
+ %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<1xf32>) -> tensor<1xf32>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<1x?xf32>) outs(%3, %1 : tensor<1xf32>, tensor<1xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_0: i64):
+ %5 = linalg.index 1 : index
+ %6 = arith.index_cast %5 : index to i64
+ %7 = arith.maximumf %in, %out : f32
+ %8 = arith.cmpf ogt, %in, %out : f32
+ %9 = arith.select %8, %6, %out_0 : i64
+ linalg.yield %7, %9 : f32, i64
+ } -> (tensor<1xf32>, tensor<1xi64>)
+ return %4#1 : tensor<1xi64>
+}
+
+// CHECK: func @argmax_only_argmax_ukernel_enabled(
+// CHECK: iree_codegen.ukernel.generic
+// CHECK-NOT: linalg.generic
+
+// -----
+
+func.func @argmax_only_foo_argmax_bar_ukernel_enabled(%arg0 : tensor<1x?xf32>) -> tensor<1xi64> attributes {
+ hal.executable.target = #hal.executable.target<"rocm", "rocm-hsaco-fb", {target_arch = "gfx1100", ukernels = "foo,argmax,bar"}>
+} {
+ %c0_i64 = arith.constant 0 : i64
+ %cst = arith.constant 0xFF800000 : f32
+ %0 = tensor.empty() : tensor<1xi64>
+ %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<1xi64>) -> tensor<1xi64>
+ %2 = tensor.empty() : tensor<1xf32>
+ %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<1xf32>) -> tensor<1xf32>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<1x?xf32>) outs(%3, %1 : tensor<1xf32>, tensor<1xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_0: i64):
+ %5 = linalg.index 1 : index
+ %6 = arith.index_cast %5 : index to i64
+ %7 = arith.maximumf %in, %out : f32
+ %8 = arith.cmpf ogt, %in, %out : f32
+ %9 = arith.select %8, %6, %out_0 : i64
+ linalg.yield %7, %9 : f32, i64
+ } -> (tensor<1xf32>, tensor<1xi64>)
+ return %4#1 : tensor<1xi64>
+}
+
+// CHECK: func @argmax_only_foo_argmax_bar_ukernel_enabled(
+// CHECK: iree_codegen.ukernel.generic
+// CHECK-NOT: linalg.generic
+
+// -----
+
+func.func @argmax_only_foo_ukernel_enabled(%arg0 : tensor<1x?xf32>) -> tensor<1xi64> attributes {
+ hal.executable.target = #hal.executable.target<"rocm", "rocm-hsaco-fb", {target_arch = "gfx1100", ukernels = "foo"}>
+} {
+ %c0_i64 = arith.constant 0 : i64
+ %cst = arith.constant 0xFF800000 : f32
+ %0 = tensor.empty() : tensor<1xi64>
+ %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<1xi64>) -> tensor<1xi64>
+ %2 = tensor.empty() : tensor<1xf32>
+ %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<1xf32>) -> tensor<1xf32>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<1x?xf32>) outs(%3, %1 : tensor<1xf32>, tensor<1xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_0: i64):
+ %5 = linalg.index 1 : index
+ %6 = arith.index_cast %5 : index to i64
+ %7 = arith.maximumf %in, %out : f32
+ %8 = arith.cmpf ogt, %in, %out : f32
+ %9 = arith.select %8, %6, %out_0 : i64
+ linalg.yield %7, %9 : f32, i64
+ } -> (tensor<1xf32>, tensor<1xi64>)
+ return %4#1 : tensor<1xi64>
+}
+
+// CHECK: func @argmax_only_foo_ukernel_enabled(
+// CHECK-NOT: iree_codegen.ukernel.generic
+// CHECK: linalg.generic
+
+// -----
+
+// Currently we do only handle -Inf case as initial values.
+func.func @argmax_2d_f32i64_not_neg_inf_init(%arg0 : tensor<1x?xf32>) -> tensor<1xi64> attributes {
+ hal.executable.target = #hal.executable.target<"rocm", "rocm-hsaco-fb", {target_arch = "gfx1100", ukernels = "all"}>
+} {
+ %c0_i64 = arith.constant 0 : i64
+ %cst = arith.constant 0.0 : f32
+ %0 = tensor.empty() : tensor<1xi64>
+ %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<1xi64>) -> tensor<1xi64>
+ %2 = tensor.empty() : tensor<1xf32>
+ %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<1xf32>) -> tensor<1xf32>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<1x?xf32>) outs(%3, %1 : tensor<1xf32>, tensor<1xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_0: i64):
+ %5 = linalg.index 1 : index
+ %6 = arith.index_cast %5 : index to i64
+ %7 = arith.maximumf %in, %out : f32
+ %8 = arith.cmpf ogt, %in, %out : f32
+ %9 = arith.select %8, %6, %out_0 : i64
+ linalg.yield %7, %9 : f32, i64
+ } -> (tensor<1xf32>, tensor<1xi64>)
+ return %4#1 : tensor<1xi64>
+}
+
+// CHECK: func @argmax_2d_f32i64_not_neg_inf_init(
+// CHECK-NOT: iree_codegen.ukernel.generic
+// CHECK: linalg.generic
+
+// -----
+
+// TODO: No technical reason this architecture is not supported.
+// Currently just picking out popular chips to support,
+// to minimize compile time and space.
+
+func.func @argmax_ukernel_unsupported_arch(%arg0 : tensor<1x?xf32>) -> tensor<1xi64> attributes {
+ hal.executable.target = #hal.executable.target<"rocm", "rocm-hsaco-fb", {target_arch = "gfx800", ukernels = "all"}>
+} {
+ %c0_i64 = arith.constant 0 : i64
+ %cst = arith.constant 0xFF800000 : f32
+ %0 = tensor.empty() : tensor<1xi64>
+ %1 = linalg.fill ins(%c0_i64 : i64) outs(%0 : tensor<1xi64>) -> tensor<1xi64>
+ %2 = tensor.empty() : tensor<1xf32>
+ %3 = linalg.fill ins(%cst : f32) outs(%2 : tensor<1xf32>) -> tensor<1xf32>
+ %4:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%arg0 : tensor<1x?xf32>) outs(%3, %1 : tensor<1xf32>, tensor<1xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_0: i64):
+ %5 = linalg.index 1 : index
+ %6 = arith.index_cast %5 : index to i64
+ %7 = arith.maximumf %in, %out : f32
+ %8 = arith.cmpf ogt, %in, %out : f32
+ %9 = arith.select %8, %6, %out_0 : i64
+ linalg.yield %7, %9 : f32, i64
+ } -> (tensor<1xf32>, tensor<1xi64>)
+ return %4#1 : tensor<1xi64>
+}
+
+// CHECK: func @argmax_ukernel_unsupported_arch(
+// CHECK-NOT: iree_codegen.ukernel.generic
+// CHECK: linalg.generic
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
index 520f12e..d36220f 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/KernelConfig.cpp
@@ -1062,6 +1062,83 @@
workgroupSize);
}
+/// Set the configuration for argmax that can be mapped to argmax uKernel.
+/// Distribute all parallel dim across different workgroups, and only use single
+/// subgroup per workgroup.
+static LogicalResult setArgmaxUkernelConfig(func::FuncOp entryPoint,
+ linalg::GenericOp op,
+ const TargetInfo &targetInfo) {
+
+ // Checks if UKernels are enabled.
+ if (auto variantOp =
+ entryPoint->getParentOfType<IREE::HAL::ExecutableVariantOp>()) {
+ auto target = variantOp.getTarget();
+ const char ukernelName[] = "argmax";
+ if (!hasUkernel(target, ukernelName) ||
+ !hasUkernelSupportedGpuArch(target)) {
+ return failure();
+ }
+ }
+
+ if (!targetInfo.hasWarpShuffle)
+ return failure();
+
+ if (failed(isArgmaxOp(op)))
+ return failure();
+ SmallVector<unsigned> parallelDims;
+ SmallVector<unsigned> reductionDims;
+ op.getParallelDims(parallelDims);
+ op.getReductionDims(reductionDims);
+
+ // Currently Argmax UKernel only support 1 reduction dim.
+ if (reductionDims.size() != 1)
+ return failure();
+
+ // Make sure reduction dimensions are static and innermost ones.
+ SmallVector<int64_t, 4> bounds = op.getStaticLoopRanges();
+ int64_t numParallelDims = op.getNumParallelLoops();
+ int64_t numDynamicReductionDims = 0;
+ for (unsigned dim : reductionDims) {
+ if (ShapedType::isDynamic(bounds[dim])) {
+ numDynamicReductionDims++;
+ }
+ if (dim < numParallelDims) {
+ return failure();
+ }
+ }
+
+ // Distribution of multi-dim masked writes currently aren't fully supported.
+ if (numDynamicReductionDims > 1) {
+ return failure();
+ }
+
+ // Tile all the parallel dimension to 1.
+ SmallVector<unsigned> partitionedLoops =
+ cast<PartitionableLoopsInterface>(op.getOperation())
+ .getPartitionableLoops(kNumMaxParallelDims);
+ size_t numLoops = partitionedLoops.empty() ? 0 : partitionedLoops.back() + 1;
+ SmallVector<int64_t> workgroupTileSizes(numLoops, 1);
+
+ // Currently Argmax Ukernel let's every thread reduce reductionDim/WarpSize
+ // number of elements, and then it does a single step butterfly warp reduce.
+ // Hence it expects workgroupSize to be warpSize(subgroupSize), and
+ // reductionTileSize to be size of the reduction dim.
+ SmallVector<int64_t> reductionTileSizes(op.getNumLoops(), 0);
+ int64_t preferredSubgroupSize = targetInfo.supportedSubgroupSizes.front();
+ reductionTileSizes[reductionDims[0]] = preferredSubgroupSize;
+ TileSizesListType tileSizes;
+ tileSizes.emplace_back(std::move(workgroupTileSizes)); // Workgroup level
+ tileSizes.emplace_back(std::move(reductionTileSizes)); // Reduction level
+ std::array<int64_t, 3> workgroupSize = {preferredSubgroupSize, 1, 1};
+ if (failed(setOpConfigAndEntryPointFnTranslation(
+ entryPoint, op, tileSizes,
+ IREE::Codegen::DispatchLoweringPassPipeline::LLVMGPUDefault,
+ workgroupSize))) {
+ return failure();
+ }
+ return success();
+}
+
/// Make UKernels take the LLVMGPUDefault lowering pipeline.
static LogicalResult
setUKernelConfig(func::FuncOp entryPoint,
@@ -1256,6 +1333,9 @@
auto genericOp = dyn_cast<linalg::GenericOp>(computeOp);
if (genericOp && succeeded(setTransposeConfig(entryPointFn, genericOp))) {
return success();
+ } else if (genericOp && succeeded(setArgmaxUkernelConfig(
+ entryPointFn, genericOp, targetInfo))) {
+ return success();
}
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPULowerExecutableTarget.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPULowerExecutableTarget.cpp
index cef1407..350a8fe 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPULowerExecutableTarget.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/LLVMGPULowerExecutableTarget.cpp
@@ -75,10 +75,11 @@
return signalPassFailure();
}
+ bool enableMicrokernels = hasUkernel(variantOp.getTarget());
OpPassManager pipeline(IREE::HAL::ExecutableVariantOp::getOperationName());
switch (translationInfo.value().getDispatchLoweringPassPipeline()) {
case IREE::Codegen::DispatchLoweringPassPipeline::LLVMGPUDefault:
- addGPUDefaultPassPipeline(pipeline);
+ addGPUDefaultPassPipeline(pipeline, enableMicrokernels);
break;
case IREE::Codegen::DispatchLoweringPassPipeline::LLVMGPUDistribute:
addGPUSimpleDistributePassPipeline(pipeline);
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
index 4727977..69b5858 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.cpp
@@ -484,9 +484,16 @@
createRemoveSingleIterationLoopPass());
}
-void addGPUDefaultPassPipeline(OpPassManager &pm) {
- tileAndBufferize(pm);
+void addGPUDefaultPassPipeline(OpPassManager &pm, bool enableMicrokernels) {
+ tileAndDistributeToWorkgroup(pm, /*useWARForCooperativeMatrixCodegen=*/true);
auto &nestedModulePM = pm.nest<ModuleOp>();
+ if (enableMicrokernels) {
+ nestedModulePM.addPass(createGPULowerToUKernelsPass());
+ }
+ nestedModulePM.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+ nestedModulePM.addNestedPass<func::FuncOp>(createCSEPass());
+
+ addBufferizePasses(nestedModulePM);
nestedModulePM.addNestedPass<func::FuncOp>(
createRemoveSingleIterationLoopPass());
}
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.h b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.h
index 9860197..0a97fb6 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.h
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/Passes.h
@@ -53,7 +53,7 @@
void addGPUWarpReductionPassPipeline(OpPassManager &pm);
/// Default pass pipeline on GPU, currently used only for the ukernel path.
-void addGPUDefaultPassPipeline(OpPassManager &pm);
+void addGPUDefaultPassPipeline(OpPassManager &pm, bool enableMicrokernels);
/// Populates passes needed to preprocess and select the translation strategy.
void buildLLVMGPUCodegenConfigurationPassPipeline(OpPassManager &pm);
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/BUILD.bazel b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/BUILD.bazel
index be733ae..bd80615 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/BUILD.bazel
@@ -59,6 +59,7 @@
"transform_gpu_pipelining.mlir",
"transform_vector_to_mma.mlir",
"transpose_pipeline_test.mlir",
+ "ukernel_pipeline_transform.mlir",
"vector_lowering.mlir",
"vector_to_gpu.mlir",
"workgroup_specialization_pipeline_test.mlir",
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/CMakeLists.txt
index caa6dcb..d20d23d 100644
--- a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/CMakeLists.txt
@@ -55,6 +55,7 @@
"transform_gpu_pipelining.mlir"
"transform_vector_to_mma.mlir"
"transpose_pipeline_test.mlir"
+ "ukernel_pipeline_transform.mlir"
"vector_lowering.mlir"
"vector_to_gpu.mlir"
"workgroup_specialization_pipeline_test.mlir"
diff --git a/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ukernel_pipeline_transform.mlir b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ukernel_pipeline_transform.mlir
new file mode 100644
index 0000000..809022f
--- /dev/null
+++ b/compiler/src/iree/compiler/Codegen/LLVMGPU/test/ukernel_pipeline_transform.mlir
@@ -0,0 +1,215 @@
+// RUN: iree-opt --split-input-file --pass-pipeline="builtin.module(hal.executable(hal.executable.variant(iree-llvmgpu-select-lowering-strategy, iree-llvmgpu-lower-executable-target)))" %s | FileCheck %s
+
+hal.executable @argmax_1d_f16i64 {
+hal.executable.variant public @rocm_hsaco_fb target(<"rocm", "rocm-hsaco-fb", {target_arch = "gfx1100", ukernels = "argmax"}>) {
+ hal.executable.export public @argmax_1d_f16i64 ordinal(0) layout(#hal.pipeline.layout<push_constants = 2, sets = [<0, bindings = [<0, storage_buffer, ReadOnly>, <1, storage_buffer>]>]>) {
+ ^bb0(%arg0: !hal.device, %arg1: index):
+ %x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg1
+ hal.return %x, %y, %z : index, index, index
+ }
+ builtin.module {
+ func.func @argmax_1d_f16i64() {
+ %c32_i64 = arith.constant 32 : i64
+ %cst = arith.constant 0xFC00 : f16
+ %c0_i64 = arith.constant 0 : i64
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.constant.load[0] : i32
+ %1 = hal.interface.constant.load[1] : i32
+ %2 = arith.extui %0 : i32 to i64
+ %3 = arith.extui %1 : i32 to i64
+ %4 = arith.shli %3, %c32_i64 : i64
+ %5 = arith.ori %2, %4 : i64
+ %6 = arith.index_castui %5 : i64 to index
+ %7 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<i64>>
+ %8 = flow.dispatch.workload.ordinal %6, 0 : index
+ %9 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<?xf16>>{%8}
+ %10 = flow.dispatch.tensor.load %9, offsets = [0], sizes = [%8], strides = [1] : !flow.dispatch.tensor<readonly:tensor<?xf16>>{%8} -> tensor<?xf16>
+ %11 = tensor.empty() : tensor<i64>
+ %12 = tensor.empty() : tensor<f16>
+ %13 = linalg.fill ins(%c0_i64 : i64) outs(%11 : tensor<i64>) -> tensor<i64>
+ %14 = linalg.fill ins(%cst : f16) outs(%12 : tensor<f16>) -> tensor<f16>
+ %15:2 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>, affine_map<(d0) -> ()>], iterator_types = ["reduction"]} ins(%10 : tensor<?xf16>) outs(%14, %13 : tensor<f16>, tensor<i64>) {
+ ^bb0(%in: f16, %out: f16, %out_0: i64):
+ %16 = linalg.index 0 : index
+ %17 = arith.index_cast %16 : index to i64
+ %18 = arith.maximumf %in, %out : f16
+ %19 = arith.cmpf ogt, %in, %out : f16
+ %20 = arith.select %19, %17, %out_0 : i64
+ linalg.yield %18, %20 : f16, i64
+ } -> (tensor<f16>, tensor<i64>)
+ flow.dispatch.tensor.store %15#1, %7, offsets = [], sizes = [], strides = [] : tensor<i64> -> !flow.dispatch.tensor<writeonly:tensor<i64>>
+ return
+ }
+ }
+}
+}
+
+// CHECK: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUDefault>
+// CHECK-LABEL: hal.executable.export public @argmax_1d_f16i64
+// CHECK-SAME: translation_info = #[[$TRANSLATION]]
+// CHECK-SAME: workgroup_size = [32 : index, 1 : index, 1 : index]
+//CHECK-LABEL: func.func @argmax_1d_f16i64
+// CHECK: iree_codegen.ukernel.generic "__iree_uk_rocm_argmax_F16I64"
+
+// -----
+
+hal.executable @argmax_2d_f32i64 {
+hal.executable.variant public @rocm_hsaco_fb target(<"rocm", "rocm-hsaco-fb", {target_arch = "gfx1100", ukernels = "argmax"}>) {
+ hal.executable.export public @argmax_2d_f32i64 ordinal(0) layout(#hal.pipeline.layout<push_constants = 2, sets = [<0, bindings = [<0, storage_buffer, ReadOnly>, <1, storage_buffer>]>]>) {
+ ^bb0(%arg0: !hal.device, %arg1: index):
+ %x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg1
+ hal.return %x, %y, %z : index, index, index
+ }
+ builtin.module {
+ func.func @argmax_2d_f32i64() {
+ %c32_i64 = arith.constant 32 : i64
+ %cst = arith.constant 0xFF800000 : f32
+ %c0_i64 = arith.constant 0 : i64
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.constant.load[0] : i32
+ %1 = hal.interface.constant.load[1] : i32
+ %2 = arith.extui %0 : i32 to i64
+ %3 = arith.extui %1 : i32 to i64
+ %4 = arith.shli %3, %c32_i64 : i64
+ %5 = arith.ori %2, %4 : i64
+ %6 = arith.index_castui %5 : i64 to index
+ %7 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<16xi64>>
+ %8 = flow.dispatch.workload.ordinal %6, 0 : index
+ %9 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<16x?xf32>>{%8}
+ %10 = flow.dispatch.tensor.load %9, offsets = [0, 0], sizes = [16, %8], strides = [1, 1] : !flow.dispatch.tensor<readonly:tensor<16x?xf32>>{%8} -> tensor<16x?xf32>
+ %11 = tensor.empty() : tensor<16xi64>
+ %12 = tensor.empty() : tensor<16xf32>
+ %13 = linalg.fill ins(%c0_i64 : i64) outs(%11 : tensor<16xi64>) -> tensor<16xi64>
+ %14 = linalg.fill ins(%cst : f32) outs(%12 : tensor<16xf32>) -> tensor<16xf32>
+ %15:2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d0)>], iterator_types = ["parallel", "reduction"]} ins(%10 : tensor<16x?xf32>) outs(%14, %13 : tensor<16xf32>, tensor<16xi64>) {
+ ^bb0(%in: f32, %out: f32, %out_0: i64):
+ %16 = linalg.index 1 : index
+ %17 = arith.index_cast %16 : index to i64
+ %18 = arith.maximumf %in, %out : f32
+ %19 = arith.cmpf ogt, %in, %out : f32
+ %20 = arith.select %19, %17, %out_0 : i64
+ linalg.yield %18, %20 : f32, i64
+ } -> (tensor<16xf32>, tensor<16xi64>)
+ flow.dispatch.tensor.store %15#1, %7, offsets = [0], sizes = [16], strides = [1] : tensor<16xi64> -> !flow.dispatch.tensor<writeonly:tensor<16xi64>>
+ return
+ }
+ }
+}
+}
+
+// CHECK: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUDefault>
+// CHECK-LABEL: hal.executable.export public @argmax_2d_f32i64
+// CHECK-SAME: translation_info = #[[$TRANSLATION]]
+// CHECK-SAME: workgroup_size = [32 : index, 1 : index, 1 : index]
+//CHECK-LABEL: func.func @argmax_2d_f32i64
+// CHECK-DAG: %[[SUBVIEW:.*]] = memref.subview{{.*}} memref<16x?xf32
+// CHECK-SAME: to memref<1x?xf32
+// CHECK: iree_codegen.ukernel.generic "__iree_uk_rocm_argmax_F32I64" ins(%[[SUBVIEW]]
+
+// -----
+
+// When the ukernel attribute is not set, we do not go through ukernel pipeline.
+hal.executable @no_ukernel_argmax_1d_f16i64 {
+hal.executable.variant public @rocm_hsaco_fb target(<"rocm", "rocm-hsaco-fb", {target_arch = "gfx1100"}>) {
+ hal.executable.export public @no_ukernel_argmax_1d_f16i64 ordinal(0) layout(#hal.pipeline.layout<push_constants = 2, sets = [<0, bindings = [<0, storage_buffer, ReadOnly>, <1, storage_buffer>]>]>) {
+ ^bb0(%arg0: !hal.device, %arg1: index):
+ %x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg1
+ hal.return %x, %y, %z : index, index, index
+ }
+ builtin.module {
+ func.func @no_ukernel_argmax_1d_f16i64() {
+ %c32_i64 = arith.constant 32 : i64
+ %cst = arith.constant 0xFC00 : f16
+ %c0_i64 = arith.constant 0 : i64
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.constant.load[0] : i32
+ %1 = hal.interface.constant.load[1] : i32
+ %2 = arith.extui %0 : i32 to i64
+ %3 = arith.extui %1 : i32 to i64
+ %4 = arith.shli %3, %c32_i64 : i64
+ %5 = arith.ori %2, %4 : i64
+ %6 = arith.index_castui %5 : i64 to index
+ %7 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<i64>>
+ %8 = flow.dispatch.workload.ordinal %6, 0 : index
+ %9 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<?xf16>>{%8}
+ %10 = flow.dispatch.tensor.load %9, offsets = [0], sizes = [%8], strides = [1] : !flow.dispatch.tensor<readonly:tensor<?xf16>>{%8} -> tensor<?xf16>
+ %11 = tensor.empty() : tensor<i64>
+ %12 = tensor.empty() : tensor<f16>
+ %13 = linalg.fill ins(%c0_i64 : i64) outs(%11 : tensor<i64>) -> tensor<i64>
+ %14 = linalg.fill ins(%cst : f16) outs(%12 : tensor<f16>) -> tensor<f16>
+ %15:2 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>, affine_map<(d0) -> ()>], iterator_types = ["reduction"]} ins(%10 : tensor<?xf16>) outs(%14, %13 : tensor<f16>, tensor<i64>) {
+ ^bb0(%in: f16, %out: f16, %out_0: i64):
+ %16 = linalg.index 0 : index
+ %17 = arith.index_cast %16 : index to i64
+ %18 = arith.maximumf %in, %out : f16
+ %19 = arith.cmpf ogt, %in, %out : f16
+ %20 = arith.select %19, %17, %out_0 : i64
+ linalg.yield %18, %20 : f16, i64
+ } -> (tensor<f16>, tensor<i64>)
+ flow.dispatch.tensor.store %15#1, %7, offsets = [], sizes = [], strides = [] : tensor<i64> -> !flow.dispatch.tensor<writeonly:tensor<i64>>
+ return
+ }
+ }
+}
+}
+
+// CHECK: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUDistribute>
+// CHECK-LABEL: hal.executable.export public @no_ukernel_argmax_1d_f16i64
+// CHECK-SAME: translation_info = #[[$TRANSLATION]]
+// CHECK-SAME: workgroup_size = [1 : index, 1 : index, 1 : index]
+//CHECK-LABEL: func.func @no_ukernel_argmax_1d_f16i64
+// CHECK-NOT: iree_codegen.ukernel.generic
+
+// -----
+
+// Currently we do only handle -Inf case as initial values.
+hal.executable @not_neg_inf_init_argmax_1d {
+hal.executable.variant public @rocm_hsaco_fb target(<"rocm", "rocm-hsaco-fb", {target_arch = "gfx1100", ukernels = "argmax"}>) {
+ hal.executable.export public @not_neg_inf_init_argmax_1d ordinal(0) layout(#hal.pipeline.layout<push_constants = 2, sets = [<0, bindings = [<0, storage_buffer, ReadOnly>, <1, storage_buffer>]>]>) {
+ ^bb0(%arg0: !hal.device, %arg1: index):
+ %x, %y, %z = flow.dispatch.workgroup_count_from_slice %arg1
+ hal.return %x, %y, %z : index, index, index
+ }
+ builtin.module {
+ func.func @not_neg_inf_init_argmax_1d() {
+ %c32_i64 = arith.constant 32 : i64
+ %cst = arith.constant 0.0 : f16
+ %c0_i64 = arith.constant 0 : i64
+ %c0 = arith.constant 0 : index
+ %0 = hal.interface.constant.load[0] : i32
+ %1 = hal.interface.constant.load[1] : i32
+ %2 = arith.extui %0 : i32 to i64
+ %3 = arith.extui %1 : i32 to i64
+ %4 = arith.shli %3, %c32_i64 : i64
+ %5 = arith.ori %2, %4 : i64
+ %6 = arith.index_castui %5 : i64 to index
+ %7 = hal.interface.binding.subspan set(0) binding(1) type(storage_buffer) alignment(64) offset(%c0) : !flow.dispatch.tensor<writeonly:tensor<i64>>
+ %8 = flow.dispatch.workload.ordinal %6, 0 : index
+ %9 = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) alignment(64) offset(%c0) flags(ReadOnly) : !flow.dispatch.tensor<readonly:tensor<?xf16>>{%8}
+ %10 = flow.dispatch.tensor.load %9, offsets = [0], sizes = [%8], strides = [1] : !flow.dispatch.tensor<readonly:tensor<?xf16>>{%8} -> tensor<?xf16>
+ %11 = tensor.empty() : tensor<i64>
+ %12 = tensor.empty() : tensor<f16>
+ %13 = linalg.fill ins(%c0_i64 : i64) outs(%11 : tensor<i64>) -> tensor<i64>
+ %14 = linalg.fill ins(%cst : f16) outs(%12 : tensor<f16>) -> tensor<f16>
+ %15:2 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>, affine_map<(d0) -> ()>], iterator_types = ["reduction"]} ins(%10 : tensor<?xf16>) outs(%14, %13 : tensor<f16>, tensor<i64>) {
+ ^bb0(%in: f16, %out: f16, %out_0: i64):
+ %16 = linalg.index 0 : index
+ %17 = arith.index_cast %16 : index to i64
+ %18 = arith.maximumf %in, %out : f16
+ %19 = arith.cmpf ogt, %in, %out : f16
+ %20 = arith.select %19, %17, %out_0 : i64
+ linalg.yield %18, %20 : f16, i64
+ } -> (tensor<f16>, tensor<i64>)
+ flow.dispatch.tensor.store %15#1, %7, offsets = [], sizes = [], strides = [] : tensor<i64> -> !flow.dispatch.tensor<writeonly:tensor<i64>>
+ return
+ }
+ }
+}
+}
+
+// CHECK: #[[$TRANSLATION:.+]] = #iree_codegen.translation_info<LLVMGPUDistribute>
+// CHECK-LABEL: hal.executable.export public @not_neg_inf_init_argmax_1d
+// CHECK-SAME: translation_info = #[[$TRANSLATION]]
+// CHECK-SAME: workgroup_size = [1 : index, 1 : index, 1 : index]
+//CHECK-LABEL: func.func @not_neg_inf_init_argmax_1d
+// CHECK-NOT: iree_codegen.ukernel.generic
diff --git a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp
index 2ddd654..69ab4e4 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp
@@ -881,4 +881,42 @@
return false;
}
+//===----------------------------------------------------------------------===//
+// GPU UKernel Utils
+//===----------------------------------------------------------------------===//
+
+// TODO: Add more popular kernels into this list and the ukernel cmake.
+// No real technical reason to only allow these aside from compile
+// time and diskspace.
+bool hasUkernelSupportedRocmArch(StringRef targetChip) {
+ const char *kSupportedTargetChip[] = {"gfx90a", "gfx940", "gfx1030",
+ "gfx1100"};
+ size_t arraySize =
+ sizeof(kSupportedTargetChip) / sizeof(kSupportedTargetChip[0]);
+ for (int i = 0; i < arraySize; i++) {
+ // return true if targetChip is found inside kSupportedTargetChip.
+ if (targetChip.compare(kSupportedTargetChip[i]) == 0)
+ return true;
+ }
+ return false;
+}
+
+bool hasUkernelSupportedRocmArch(IREE::HAL::ExecutableTargetAttr targetAttr) {
+ auto targetArch = getConfigStringAttr(targetAttr, "target_arch");
+ if (!targetArch) {
+ return false;
+ }
+ StringRef targetArchStr = targetArch->getValue();
+ return hasUkernelSupportedRocmArch(targetArchStr);
+}
+
+/// Checks if target GPU has UKernel support.
+bool hasUkernelSupportedGpuArch(IREE::HAL::ExecutableTargetAttr targetAttr) {
+ if (isROCMBackend(targetAttr) && hasUkernelSupportedRocmArch(targetAttr)) {
+ return true;
+ }
+ // TODO: Once plumbed, add a CUDA backend and supported cuda arch check.
+ return false;
+}
+
} // namespace mlir::iree_compiler
diff --git a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h
index e885226..0abe311 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h
+++ b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h
@@ -110,6 +110,16 @@
/// using shared memory when CodeGen towards the GPU.
bool sharedMemTransposeFilter(AffineMap indexMap);
+//===----------------------------------------------------------------------===//
+// GPU UKernel Utils
+//===----------------------------------------------------------------------===//
+
+/// Checks if target Chip(StringRef) has UKernel support.
+bool hasUkernelSupportedRocmArch(StringRef targetChip);
+
+/// Checks if targetAttr's GPU target has UKernel support.
+bool hasUkernelSupportedGpuArch(IREE::HAL::ExecutableTargetAttr targetAttr);
+
} // namespace mlir::iree_compiler
#endif // IREE_COMPILER_CODEGEN_UTILS_GPUUTILS_H_
diff --git a/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp b/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
index ad7844c..544501f 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Utils/Utils.cpp
@@ -148,6 +148,10 @@
return targetAttr && targetAttr.getBackend().getValue().starts_with("vmvx");
}
+bool isROCMBackend(IREE::HAL::ExecutableTargetAttr targetAttr) {
+ return targetAttr && targetAttr.getBackend().getValue().starts_with("rocm");
+}
+
bool hasUkernel(IREE::HAL::ExecutableTargetAttr targetAttr,
StringRef ukernelName) {
auto enabledUkernels = getConfigStringAttr(targetAttr, "ukernels");
@@ -800,6 +804,104 @@
}
}
+LogicalResult isArgmaxOp(linalg::GenericOp genericOp) {
+ // Check for 2 results(value, index), and 1 input
+ if (genericOp.getNumDpsInits() != 2) {
+ return failure();
+ }
+ if (genericOp.getNumDpsInputs() != 1) {
+ return failure();
+ }
+
+ // If max value is being used, it is not a pure argmax.
+ if (!genericOp.getResults()[0].use_empty()) {
+ return failure();
+ }
+
+ // Check that the rank is at least 3 and all loops are parallel
+ unsigned numLoops = genericOp.getNumLoops();
+ unsigned numParallelLoops = genericOp.getNumParallelLoops();
+
+ // Argmax will require 1D reduction.
+ if (numParallelLoops != (numLoops - 1)) {
+ return failure();
+ }
+ // TODO: Add better affine map checks.
+ auto indexing_maps = genericOp.getIndexingMapsArray();
+ if (!indexing_maps[0].isIdentity())
+ return failure();
+
+ // Check that initial value is negative Infinite.
+ // TODO: Move this check to ukernel once we implement
+ // variant to handle non neg-Inf initial value.
+ Value initVal = genericOp.getDpsInitOperand(0)->get();
+ auto fillOp = initVal.getDefiningOp<linalg::FillOp>();
+ if (!fillOp)
+ return failure();
+ Value fillVal = fillOp.getDpsInputOperand(0)->get();
+ if (!matchPattern(fillVal, m_NegInfFloat()))
+ return failure();
+
+ // Work back from linalg.yield and check body of genericOp.
+ // The genericOp should yield the result of an arith.select,
+ // preceded by an arith.cmpf, arith.maximumf, and arith.extui
+ auto yieldOp = cast<linalg::YieldOp>(genericOp.getBody()->getTerminator());
+ Value producerOutput;
+ Operation *producer;
+
+ // Producer of linalg.yield 1st arg is arith.maximumf
+ {
+ producerOutput = yieldOp->getOperand(0);
+ producer = producerOutput.getDefiningOp();
+ if (!producer || producer->getNumOperands() == 0) {
+ return failure();
+ }
+ if (!matchPattern(producer, m_Op<arith::MaximumFOp>())) {
+ return failure();
+ }
+ }
+
+ // Producer of linalg.yield op 2nd arg is arith.select
+ // TODO: Add check that select is selecting between linalg.index and index of
+ // current max.
+ {
+ producerOutput = yieldOp->getOperand(1);
+ producer = producerOutput.getDefiningOp();
+ if (!producer || producer->getNumOperands() == 0) {
+ return failure();
+ }
+ if (!matchPattern(producer, m_Op<arith::SelectOp>())) {
+ return failure();
+ }
+ }
+
+ // Producer of arith.select op is arith.cmpf
+ {
+ producerOutput = producer->getOperand(0);
+ producer = producerOutput.getDefiningOp();
+ if (!producer || producer->getNumOperands() == 0) {
+ return failure();
+ }
+ auto producerCmpFOp = dyn_cast<arith::CmpFOp>(producer);
+ if (!producerCmpFOp) {
+ return failure();
+ }
+ if (producerCmpFOp.getPredicate() != arith::CmpFPredicate::OGT) {
+ return failure();
+ }
+
+ // Check that in and out of cmpf are loop variables.
+ // Currently first operand is disabled because it may be mixed type
+ // which would lead it to be extf(%arg0).
+ // TODO: Add better mixed type support check.
+ if (producer->getOperand(1) != genericOp.getBody()->getArgument(1)) {
+ return failure();
+ }
+ }
+
+ return success();
+}
+
//===---------------------------------------------------------------------===//
// Replace Memref users (transitively)
//===---------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Codegen/Utils/Utils.h b/compiler/src/iree/compiler/Codegen/Utils/Utils.h
index 5a75bf0..2c1ab9a 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/Utils.h
+++ b/compiler/src/iree/compiler/Codegen/Utils/Utils.h
@@ -73,6 +73,9 @@
/// Methods to get target information.
bool isVMVXBackend(IREE::HAL::ExecutableTargetAttr targetAttr);
+/// Methods to get target information.
+bool isROCMBackend(IREE::HAL::ExecutableTargetAttr targetAttr);
+
// Returns true if the ukernel with given `ukernelName` is enabled.
// If `ukernelName` is empty (the default), returns true if any ukernel
// is enabled at all.
@@ -204,6 +207,9 @@
OpFoldResult byteOffset,
Type elementType);
+/// Check if a linalg.generic is representing an argmax operation.
+LogicalResult isArgmaxOp(linalg::GenericOp genericOp);
+
/// Replace the uses of memref value `origValue` with the given
/// `replacementValue`. Some uses of the memref value might require changes to
/// the operation itself. Create new operations which can carry the change, and
diff --git a/experimental/regression_suite/pyproject.toml b/experimental/regression_suite/pyproject.toml
index 15c207c..7f590ee 100644
--- a/experimental/regression_suite/pyproject.toml
+++ b/experimental/regression_suite/pyproject.toml
@@ -7,6 +7,8 @@
"plat_host_cpu: mark tests as running on the host CPU",
"plat_rdna3_vulkan: mark tests as running on AMD RDNA3 Vulkan device",
"plat_nvidia_a100: mark tests as running on NVIDIA A100 device",
+ "plat_gfx90a_rocm: mark tests as running on AMD GFX90A ROCm device",
+ "plat_gfx940_rocm: mark tests as running on AMD GFX940 ROCm device",
"plat_rdna3_rocm: mark tests as running on AMD RDNA3 ROCm device",
"presubmit: mark test as running on presubmit",
"postsubmit: mark test as running on postsubmit",
diff --git a/experimental/regression_suite/tests/pregenerated/test_ukernel.py b/experimental/regression_suite/tests/pregenerated/test_ukernel.py
new file mode 100644
index 0000000..1fc5015
--- /dev/null
+++ b/experimental/regression_suite/tests/pregenerated/test_ukernel.py
@@ -0,0 +1,185 @@
+# Copyright 2023 The IREE Authors
+#
+# Licensed under the Apache License v2.0 with LLVM Exceptions.
+# See https://llvm.org/LICENSE.txt for license information.
+# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+import pytest
+from ireers import *
+
+###############################################################################
+# Fixtures
+###############################################################################
+
+COMMON_FLAGS = [
+ "--iree-input-type=none",
+ "--iree-stream-resource-index-bits=64",
+ "--iree-vm-target-index-bits=64",
+ "--iree-stream-resource-max-allocation-size=3221225472",
+]
+
+argmax_ukernel_source = fetch_source_fixture(
+ "https://storage.googleapis.com/shark_tank/ukernel_regression/20231217/argmax/argmax_3d_linalg.mlir",
+ group="argmax_ukernel_linalg",
+)
+
+
+@pytest.fixture
+def argmax_ukernel_gfx90a_rocm_vmfb(argmax_ukernel_source):
+ return iree_compile(
+ argmax_ukernel_source,
+ "gfx90a_rocm",
+ flags=COMMON_FLAGS
+ + [
+ "--iree-hal-target-backends=rocm",
+ "--iree-rocm-target-chip=gfx90a",
+ "--iree-rocm-link-bc=true",
+ "--iree-rocm-enable-ukernels=argmax",
+ ],
+ )
+
+
+@pytest.fixture
+def argmax_ukernel_gfx940_rocm_vmfb(argmax_ukernel_source):
+ return iree_compile(
+ argmax_ukernel_source,
+ "gfx940_rocm",
+ flags=COMMON_FLAGS
+ + [
+ "--iree-hal-target-backends=rocm",
+ "--iree-rocm-target-chip=gfx940",
+ "--iree-rocm-link-bc=true",
+ "--iree-rocm-enable-ukernels=argmax",
+ ],
+ )
+
+
+###############################################################################
+# Correctness
+###############################################################################
+
+# Generation script:
+# argmax_input_f16 = np.random.normal(size=[2, 4, 33000]).astype(np.float32)
+# argmax_output_f16 = np.argmax(argmax_input_f16,axis=-1).astype(np.float32)
+# argmax_input_f32 = np.random.normal(size=[2, 4, 33000]).astype(np.float32)
+# argmax_output_f32 = np.argmax(argmax_input_f32,axis=-1).astype(np.float32)
+# TODO: Currently forcing sitofp (i32 -> f32) and (i64 -> f32) because expected_output
+# cannot compare signless i64 from vmfb and by default si64 from npy.
+
+argmax_input_f16 = fetch_source_fixture(
+ "https://storage.googleapis.com/shark_tank/ukernel_regression/20231217/argmax/argmax_3d_input_f16.npy",
+ group="argmax_ukernel_input_f16",
+)
+
+argmax_output_f16 = fetch_source_fixture(
+ "https://storage.googleapis.com/shark_tank/ukernel_regression/20231217/argmax/argmax_3d_output_f16.npy",
+ group="argmax_ukernel_output_f16",
+)
+
+argmax_input_f32 = fetch_source_fixture(
+ "https://storage.googleapis.com/shark_tank/ukernel_regression/20231217/argmax/argmax_3d_input_f32.npy",
+ group="argmax_ukernel_input_f32",
+)
+
+argmax_output_f32 = fetch_source_fixture(
+ "https://storage.googleapis.com/shark_tank/ukernel_regression/20231217/argmax/argmax_3d_output_f32.npy",
+ group="argmax_ukernel_output_f32",
+)
+
+
+@pytest.mark.presubmit
+@pytest.mark.unstable_linalg
+@pytest.mark.plat_gfx90a_rocm
+def test_correctness_gfx90a_rocm(
+ argmax_ukernel_gfx90a_rocm_vmfb,
+ argmax_input_f16,
+ argmax_output_f16,
+ argmax_input_f32,
+ argmax_output_f32,
+):
+ iree_run_module(
+ argmax_ukernel_gfx90a_rocm_vmfb,
+ device="rocm",
+ function="argmax_3d_dyn_f16i32",
+ args=[
+ f"--input=@{argmax_input_f16.path}",
+ f"--expected_output=@{argmax_output_f16.path}",
+ ],
+ )
+ iree_run_module(
+ argmax_ukernel_gfx90a_rocm_vmfb,
+ device="rocm",
+ function="argmax_3d_dyn_f16i64",
+ args=[
+ f"--input=@{argmax_input_f16.path}",
+ f"--expected_output=@{argmax_output_f16.path}",
+ ],
+ )
+
+ iree_run_module(
+ argmax_ukernel_gfx90a_rocm_vmfb,
+ device="rocm",
+ function="argmax_3d_dyn_f32i32",
+ args=[
+ f"--input=@{argmax_input_f32.path}",
+ f"--expected_output=@{argmax_output_f32.path}",
+ ],
+ )
+ iree_run_module(
+ argmax_ukernel_gfx90a_rocm_vmfb,
+ device="rocm",
+ function="argmax_3d_dyn_f32i64",
+ args=[
+ f"--input=@{argmax_input_f32.path}",
+ f"--expected_output=@{argmax_output_f32.path}",
+ ],
+ )
+
+
+@pytest.mark.presubmit
+@pytest.mark.unstable_linalg
+@pytest.mark.plat_gfx940_rocm
+def test_correctness_gfx940_rocm(
+ argmax_ukernel_gfx940_rocm_vmfb,
+ argmax_input_f16,
+ argmax_output_f16,
+ argmax_input_f32,
+ argmax_output_f32,
+):
+ iree_run_module(
+ argmax_ukernel_gfx940_rocm_vmfb,
+ device="rocm",
+ function="argmax_3d_dyn_f16i32",
+ args=[
+ f"--input=@{argmax_input_f16.path}",
+ f"--expected_output=@{argmax_output_f16.path}",
+ ],
+ )
+ iree_run_module(
+ argmax_ukernel_gfx940_rocm_vmfb,
+ device="rocm",
+ function="argmax_3d_dyn_f16i64",
+ args=[
+ f"--input=@{argmax_input_f16.path}",
+ f"--expected_output=@{argmax_output_f16.path}",
+ ],
+ )
+
+ iree_run_module(
+ argmax_ukernel_gfx940_rocm_vmfb,
+ device="rocm",
+ function="argmax_3d_dyn_f32i32",
+ args=[
+ f"--input=@{argmax_input_f32.path}",
+ f"--expected_output=@{argmax_output_f32.path}",
+ ],
+ )
+ iree_run_module(
+ argmax_ukernel_gfx940_rocm_vmfb,
+ device="rocm",
+ function="argmax_3d_dyn_f32i64",
+ args=[
+ f"--input=@{argmax_input_f32.path}",
+ f"--expected_output=@{argmax_output_f32.path}",
+ ],
+ )
