Adding --iree-rocm-container-type= flag. (#20902)

This should almost always be 'auto' but when performing executable
compilation (`--compile-mode=hal-executable`) it can be explicitly set
to `hsaco` to have the output of the compile command be the raw HSACO
ELF image.

Example:
```
iree-compile \
  --compile-mode=hal-executable \
  --iree-hal-target-device=hip --iree-hip-target=gfx942 \
  --iree-rocm-container-type=hsaco \
  tools/test/iree-benchmark-executable.mlir \
  -o=hsaco.elf
```

Author: benvanik

compiler/plugins/target/ROCM/ROCMTarget.cpp

@@ -61,10 +61,22 @@ namespace mlir::iree_compiler::IREE::HAL {
 
 namespace {
 
+enum class ContainerType {
+  // Automatically detect the container type from the target ABI attribute.
+  Auto,
+  // HIP ExecutableDef flatbuffer.
+  HIP,
+  // AMDGPU ExecutableDef flatbuffer.
+  AMDGPU,
+  // Raw HSACO image (ELF).
+  HSACO,
+};
+
 // TODO(#18792): rename flags back to iree-rocm- as they are not HIP-specific.
 struct ROCMOptions {
   std::string target = "";
   std::string targetFeatures = "";
+  ContainerType containerType = ContainerType::Auto;
   std::string bitcodeDirectory = getDefaultBitcodeDirectory();
   int wavesPerEu = 0;
   std::string enableROCMUkernels = "none";
@@ -80,6 +92,7 @@ struct ROCMOptions {
   void bindOptions(OptionsBinder &binder) {
     using namespace llvm;
     static cl::OptionCategory category("HIP HAL Target");
+
     binder.opt<std::string>(
         "iree-hip-target", target, cl::cat(category),
         cl::desc(
@@ -93,16 +106,34 @@ struct ROCMOptions {
             "for more details."
             // clang-format on
             ));
+
     binder.opt<std::string>(
         "iree-hip-target-features", targetFeatures, cl::cat(category),
         cl::desc("HIP target features as expected by LLVM AMDGPU backend; "
                  "e.g., '+sramecc,+xnack'."));
+
+    binder.opt<ContainerType>(
+        "iree-rocm-container-type", containerType,
+        llvm::cl::desc("Serialized executable container type."),
+        llvm::cl::cat(category),
+        llvm::cl::values(clEnumValN(ContainerType::Auto, "auto",
+                                    "Automatically detect the container type "
+                                    "from the target ABI attribute."),
+                         clEnumValN(ContainerType::HIP, "hip",
+                                    "HIP ExecutableDef flatbuffer."),
+                         clEnumValN(ContainerType::AMDGPU, "amdgpu",
+                                    "AMDGPU ExecutableDef flatbuffer."),
+                         clEnumValN(ContainerType::HSACO, "hsaco",
+                                    "Raw HSACO image (ELF).")));
+
     binder.opt<std::string>("iree-hip-bc-dir", bitcodeDirectory,
                             cl::cat(category),
                             cl::desc("Directory of HIP Bitcode."));
+
     binder.opt<int>("iree-hip-waves-per-eu", wavesPerEu, cl::cat(category),
                     cl::desc("Optimization hint specifying minimum "
                              "number of waves per execution unit."));
+
     binder.opt<std::string>(
         "iree-hip-enable-ukernels", enableROCMUkernels, cl::cat(category),
         cl::desc("Enables microkernels in the HIP compiler backend. May be "
@@ -124,6 +155,7 @@ struct ROCMOptions {
                  "to be passed to the target backend compiler during HIP "
                  "executable serialization"),
         cl::ZeroOrMore, cl::cat(category));
+
     binder.opt<bool>("iree-hip-llvm-slp-vec", slpVectorization,
                      cl::cat(category),
                      cl::desc("Enable slp vectorization in llvm opt."));
@@ -673,14 +705,44 @@ class ROCMTargetBackend final : public TargetBackend {
                      ".hsaco", targetHSACO);
     }
 
-    // Wrap the HSACO ELF binary in a Flatbuffers container.
+    // Determine container type from the target ABI attribute.
+    ContainerType containerType = options.containerType;
+    if (containerType == ContainerType::Auto) {
+      if (getABI(targetAttr) == "amdgpu") {
+        containerType = ContainerType::AMDGPU;
+      } else {
+        containerType = ContainerType::HIP;
+      }
+    }
+
+    // Wrap the HSACO ELF binary in the requested container type (if any).
     FailureOr<DenseIntElementsAttr> binaryContainer;
-    if (getABI(targetAttr) == "amdgpu") {
+    switch (containerType) {
+    case ContainerType::Auto: {
+      // Resolved above; unreachable. Fall-through to error case.
+      assert(false && "auto container type must have resolved earlier");
+      break;
+    }
+    case ContainerType::AMDGPU: {
       binaryContainer = serializeAMDGPUBinaryContainer(
           serializationOptions, variantOp, exportOps, targetHSACO);
-    } else {
+      break;
+    }
+    case ContainerType::HIP: {
       binaryContainer = serializeHIPBinaryContainer(
           serializationOptions, variantOp, exportOps, targetHSACO);
+      break;
+    }
+    case ContainerType::HSACO: {
+      SmallVector<uint8_t> image;
+      image.resize(targetHSACO.size());
+      std::memcpy(image.data(), targetHSACO.data(), image.size());
+      binaryContainer = DenseIntElementsAttr::get(
+          VectorType::get({static_cast<int64_t>(targetHSACO.size())},
+                          executableBuilder.getI8Type()),
+          image);
+      break;
+    }
     }
     if (failed(binaryContainer) || !binaryContainer.value()) {
       return failure();

compiler/plugins/target/ROCM/test/BUILD.bazel

@@ -4,7 +4,6 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
-load("//build_tools/bazel:build_defs.oss.bzl", "iree_cmake_extra_content")
 load("//build_tools/bazel:iree_lit_test.bzl", "iree_lit_test_suite")
 
 package(

compiler/plugins/target/ROCM/test/CMakeLists.txt

@@ -42,6 +42,7 @@ iree_lit_test_suite(
   SRCS
     "external_function_validation.mlir"
     "smoketest.mlir"
+    "smoketest_hsaco.mlir"
     "target_device_features.mlir"
   TOOLS
     FileCheck

compiler/plugins/target/ROCM/test/smoketest_hsaco.mlir

@@ -0,0 +1,31 @@
+// RUN: iree-opt --split-input-file --iree-hal-serialize-all-executables --iree-rocm-container-type=hsaco %s | FileCheck %s
+
+// This smoketest verifies that serializing with the --iree-rocm-container-type=
+// flag set to raw HSACO ELF files produces an embedded ELF file. This cannot be
+// used with the IREE runtime but can be useful when using
+// --compile-mode=hal-executable and wanting an HSACO to pass to other tooling
+// without needing to unwrap the Flatbuffer. To avoid test churn we just check
+// that the `.ELF` magic bytes are present at the start and ignore the contents.
+
+//      CHECK: hal.executable public @executable
+//      CHECK: hal.executable.binary public @rocm_hsaco_fb attributes {
+// CHECK-SAME:   data = dense<"0x7F454C46
+// CHECK-SAME:   format = "rocm-hsaco-fb"
+
+#pipeline_layout = #hal.pipeline.layout<bindings = [
+  #hal.pipeline.binding<storage_buffer>
+]>
+#executable_target = #hal.executable.target<"rocm", "rocm-hsaco-fb", {
+  abi = "hip",
+  iree.gpu.target = #iree_gpu.target<arch = "gfx942", features = "", wgp = <compute =  fp64|fp32|fp16|int64|int32|int16|int8, storage =  b64|b32|b16|b8, subgroup =  shuffle|arithmetic, dot =  dp4xi8toi32, mma = [<MFMA_F32_16x16x16_BF16>, <MFMA_F32_32x32x8_BF16>, <MFMA_F32_16x16x32_F8E5M2FNUZ>, <MFMA_F32_16x16x32_F8E5M2FNUZ_F8E4M3FNUZ>, <MFMA_F32_16x16x32_F8E4M3FNUZ>, <MFMA_F32_16x16x32_F8E4M3FNUZ_F8E5M2FNUZ>, <MFMA_F32_32x32x16_F8E5M2FNUZ>, <MFMA_F32_32x32x16_F8E5M2FNUZ_F8E4M3FNUZ>, <MFMA_F32_32x32x16_F8E4M3FNUZ>, <MFMA_F32_32x32x16_F8E4M3FNUZ_F8E5M2FNUZ>, <MFMA_I32_16x16x32_I8>, <MFMA_I32_32x32x16_I8>, <MFMA_F64_16x16x4_F64>, <MFMA_F32_16x16x4_F32>, <MFMA_F32_16x16x16_F16>, <MFMA_F32_32x32x8_F16>], subgroup_size_choices = [64], max_workgroup_sizes = [1024, 1024, 1024], max_thread_count_per_workgroup = 1024, max_workgroup_memory_bytes = 65536, max_workgroup_counts = [2147483647, 2147483647, 2147483647], max_load_instruction_bits = 128, simds_per_wgp = 4, vgpr_space_bits = 16384>>
+}>
+hal.executable public @executable {
+  hal.executable.variant public @rocm_hsaco_fb target(#executable_target) {
+    hal.executable.export public @export ordinal(0) layout(#pipeline_layout)
+    builtin.module {
+      llvm.func @export() {
+        llvm.return
+      }
+    }
+  }
+}