Adding selection condition to hal.executable.variant. (iree-org#15284)

This allows for variants to declare host logic that determines whether
the variant should be selected for loading. When multiple variants are
available their declared conditions will be evaluated in op order along
with the existing executable format match.

Unfortunately the MLIR SymbolTable trait disallows multiple regions on
any op holding it so a new `hal.executable.condition` region op was
added that may be optionally present on any `hal.executable.variant`.
Ideally we clean this up and make it an optional region but that'll need
relaxing of upstream assertions like
https://sourcegraph.com/github.com/llvm/llvm-project/-/blob/mlir/lib/IR/SymbolTable.cpp?L122-123
(ideally either treating region 0 as the symbol table on ops or having
an interface override for selecting the region ala `getCallableRegion`
such as `getSymbolTableRegion`).

This removes the `hal.device.switch` op in favor of `scf.index_switch`.
When we start running const expr hoisting during the HAL pipeline this
should allow variant selection to be completely hoisted to
initialization time (or at least memoized per device). There's decent
low-hanging future work on optimizing the ranking/selection and
improving `scf.index_switch` hoisting/canonicalization to make things
better.

compiler/plugins/target/CUDA/CUDATarget.cpp

@@ -431,7 +431,7 @@ class CUDATargetBackend final : public TargetBackend {
     // Collect all the entry point parameters.
     SmallVector<std::array<int32_t, 3>> workgroupSizes;
     SmallVector<uint32_t> workgroupLocalMemories;
-    for (auto exportOp : variantOp.getOps<IREE::HAL::ExecutableExportOp>()) {
+    for (auto exportOp : variantOp.getExportOps()) {
       std::array<int32_t, 3> workgroupSize;
       if (std::optional<ArrayAttr> workgroupSizeAttr =
               exportOp.getWorkgroupSize()) {
@@ -472,7 +472,7 @@ class CUDATargetBackend final : public TargetBackend {
       // these to match the names in their kernels. We don't support any kind of
       // mangling and if the user was silly enough to rely on nvcc C++ mangling
       // they'll have to figure that out.
-      for (auto exportOp : variantOp.getOps<IREE::HAL::ExecutableExportOp>()) {
+      for (auto exportOp : variantOp.getExportOps()) {
         entryPointNames.emplace_back(exportOp.getSymName());
       }
 
@@ -503,8 +503,7 @@ class CUDATargetBackend final : public TargetBackend {
       }
 
       for (auto [exportOp, workgroupSize] :
-           llvm::zip_equal(variantOp.getOps<IREE::HAL::ExecutableExportOp>(),
-                           workgroupSizes)) {
+           llvm::zip_equal(variantOp.getExportOps(), workgroupSizes)) {
         auto *llvmFunc = llvmModule->getFunction(exportOp.getName());
         if (llvmFunc->isDeclaration())
           continue;

compiler/src/iree/compiler/Codegen/LLVMCPU/BUILD.bazel

@@ -96,7 +96,6 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Dialect/Flow/IR",
         "//compiler/src/iree/compiler/Dialect/HAL/IR",
         "//compiler/src/iree/compiler/Dialect/HAL/IR:HALDialect",
-        "//compiler/src/iree/compiler/Dialect/HAL/Utils",
         "//compiler/src/iree/compiler/Dialect/Util/IR",
         "//compiler/src/iree/compiler/Dialect/Util/Transforms",
         "//compiler/src/iree/compiler/Utils",

compiler/src/iree/compiler/Codegen/LLVMCPU/CMakeLists.txt

@@ -146,7 +146,6 @@ iree_cc_library(
     iree::compiler::Dialect::Flow::IR
     iree::compiler::Dialect::HAL::IR
     iree::compiler::Dialect::HAL::IR::HALDialect
-    iree::compiler::Dialect::HAL::Utils
     iree::compiler::Dialect::Util::IR
     iree::compiler::Dialect::Util::Transforms
     iree::compiler::Utils

compiler/src/iree/compiler/Codegen/LLVMGPU/ConvertToLLVM.cpp

@@ -85,7 +85,7 @@ void ConvertToDynamicSharedMemory(ModuleOp moduleOp) {
   // Add the amount of shared memory required as an attribute.
   auto variantOp = moduleOp->getParentOfType<IREE::HAL::ExecutableVariantOp>();
   if (variantOp != nullptr) {
-    for (auto exportOp : variantOp.getOps<IREE::HAL::ExecutableExportOp>()) {
+    for (auto exportOp : variantOp.getExportOps()) {
       exportOp->setAttr(exportOp.getWorkgroupLocalMemoryAttrName(),
                         builder.getIndexAttr(numberOfBytes));
     }

compiler/src/iree/compiler/Codegen/SPIRV/ConvertToSPIRVPass.cpp

@@ -179,8 +179,7 @@ struct HALInterfaceLoadConstantConverter final
     // TODO(#1519): this conversion should look up the entry point information
     // to get the total push constant count.
     auto variantOp = loadOp->getParentOfType<IREE::HAL::ExecutableVariantOp>();
-    auto exportOps =
-        llvm::to_vector<1>(variantOp.getOps<IREE::HAL::ExecutableExportOp>());
+    auto exportOps = llvm::to_vector<1>(variantOp.getExportOps());
     assert(exportOps.size() == 1);
     auto layoutAttr = exportOps.front().getLayout();
 

compiler/src/iree/compiler/Codegen/Utils/LinkingUtils.cpp

@@ -236,14 +236,14 @@ LogicalResult linkExecutablesInto(
       // Move any constant blocks that need to be preserved for future host
       // translation. There may be duplicates provided but they'll be cleaned
       // up in future passes.
-      for (auto constantBlockOp : llvm::make_early_inc_range(
-               variantOp.getOps<IREE::HAL::ExecutableConstantBlockOp>())) {
+      for (auto constantBlockOp :
+           llvm::make_early_inc_range(variantOp.getConstantBlockOps())) {
         constantBlockOp->moveBefore(&*linkedTargetBuilder.getInsertionPoint());
       }
 
       // Clone export ops and queue remapping ordinals and updating
       // symbol refs.
-      for (auto exportOp : variantOp.getOps<IREE::HAL::ExecutableExportOp>()) {
+      for (auto exportOp : variantOp.getExportOps()) {
         auto newExportOp =
             linkedTargetBuilder.create<IREE::HAL::ExecutableExportOp>(
                 exportOp.getLoc(), exportOp.getSymNameAttr(),
@@ -290,7 +290,7 @@ LogicalResult linkExecutablesInto(
   replaceEntryPointUses(moduleOp, symbolReplacements);
 
   // Remove if we didn't add anything.
-  if (linkedTargetOp.getOps<IREE::HAL::ExecutableExportOp>().empty()) {
+  if (linkedTargetOp.getExportOps().empty()) {
     linkedTargetOp.erase();
     linkedExecutableOp.erase();
   }

compiler/src/iree/compiler/Codegen/Utils/Utils.cpp

@@ -39,7 +39,7 @@ FailureOr<IREE::HAL::ExecutableExportOp> getEntryPoint(func::FuncOp funcOp) {
   if (!variantOp)
     return failure();
 
-  for (auto op : variantOp.getOps<IREE::HAL::ExecutableExportOp>()) {
+  for (auto op : variantOp.getExportOps()) {
     if (op.getSymName() == funcOp.getName()) {
       return op;
     }
@@ -66,7 +66,7 @@ llvm::StringMap<IREE::HAL::ExecutableExportOp>
 getAllEntryPoints(ModuleOp module) {
   auto variantOp = module->getParentOfType<IREE::HAL::ExecutableVariantOp>();
   llvm::StringMap<IREE::HAL::ExecutableExportOp> exportOps;
-  for (auto op : variantOp.getOps<IREE::HAL::ExecutableExportOp>()) {
+  for (auto op : variantOp.getExportOps()) {
     exportOps[op.getSymName()] = op;
   }
   return exportOps;

compiler/src/iree/compiler/Dialect/Flow/Transforms/test/outline_dispatch_regions.mlir

@@ -202,6 +202,15 @@ func.func @dispatchExtern(%arg0: tensor<4xi32>, %arg1: tensor<8xi32>, %arg2: i32
   // CHECK-SAME: hal.interface.bindings = [#hal.interface.binding<0, 0>, #hal.interface.binding<0, 1>]
   // CHECK-SAME: } : (tensor<4xi32>, tensor<8xi32>, i32) -> %arg1
   %result = hal.dispatch.extern "main"[%x, %y](%arg0, %arg1, %arg2) : (tensor<4xi32>, tensor<8xi32>, i32) -> %arg1
+    // Translates the workload (%x and %y captured above) into an XYZ workgroup
+    // count, optionally using device information.
+    count(%device: !hal.device, %x_capture: index, %y_capture: index) -> (index, index, index) {
+      // Shows how device queries can be used when computing the workgroup count.
+      // The device is the one used at runtime.
+      %ok, %z_i32 = hal.device.query<%device : !hal.device> key("some" :: "value") : i1, i32
+      %z = arith.index_cast %z_i32 : i32 to index
+      hal.return %x_capture, %y_capture, %z : index, index, index
+    }
     // Must match the external definition.
     layout(#hal.pipeline.layout<push_constants = 1, sets = [
       <0, bindings = [
@@ -219,15 +228,6 @@ func.func @dispatchExtern(%arg0: tensor<4xi32>, %arg1: tensor<8xi32>, %arg2: i32
       #hal.executable.target<"llvm-cpu", "a"> = [#hal.executable.object<{path = "a.o"}>],
       #hal.executable.target<"llvm-cpu", "b"> = [#hal.executable.object<{path = "b.o"}>]
     }>)
-    // Translates the workload (%x and %y captured above) into an XYZ workgroup
-    // count, optionally using device information.
-    count(%device: !hal.device, %x_capture: index, %y_capture: index) -> (index, index, index) {
-      // Shows how device queries can be used when computing the workgroup count.
-      // The device is the one used at runtime.
-      %ok, %z_i32 = hal.device.query<%device : !hal.device> key("some" :: "value") : i1, i32
-      %z = arith.index_cast %z_i32 : i32 to index
-      hal.return %x_capture, %y_capture, %z : index, index, index
-    }
   // CHECK: return %[[RESULT]]
   return %result : tensor<8xi32>
 }

compiler/src/iree/compiler/Dialect/HAL/Conversion/BUILD.bazel

@@ -25,7 +25,6 @@ iree_compiler_cc_library(
     ],
     deps = [
         "//compiler/src/iree/compiler/Dialect/HAL/IR",
-        "//compiler/src/iree/compiler/Dialect/HAL/Utils",
         "//compiler/src/iree/compiler/Dialect/Util/IR",
         "@llvm-project//mlir:FuncDialect",
         "@llvm-project//mlir:IR",

compiler/src/iree/compiler/Dialect/HAL/Conversion/CMakeLists.txt

@@ -26,7 +26,6 @@ iree_cc_library(
     MLIRMemRefDialect
     MLIRTransforms
     iree::compiler::Dialect::HAL::IR
-    iree::compiler::Dialect::HAL::Utils
     iree::compiler::Dialect::Util::IR
   PUBLIC
 )

compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/BUILD.bazel

@@ -32,7 +32,6 @@ iree_compiler_cc_library(
     deps = [
         "//compiler/src/iree/compiler/Dialect/HAL:hal_imports",
         "//compiler/src/iree/compiler/Dialect/HAL/IR",
-        "//compiler/src/iree/compiler/Dialect/HAL/Utils",
         "//compiler/src/iree/compiler/Dialect/Util/IR",
         "//compiler/src/iree/compiler/Dialect/VM/Conversion",
         "//compiler/src/iree/compiler/Dialect/VM/Conversion/StandardToVM",

compiler/src/iree/compiler/Dialect/HAL/Conversion/HALToVM/CMakeLists.txt

@@ -34,7 +34,6 @@ iree_cc_library(
     MLIRPass
     MLIRTransforms
     iree::compiler::Dialect::HAL::IR
-    iree::compiler::Dialect::HAL::Utils
     iree::compiler::Dialect::HAL::hal_imports
     iree::compiler::Dialect::Util::IR
     iree::compiler::Dialect::VM::Conversion

compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/BUILD.bazel

@@ -25,7 +25,6 @@ iree_compiler_cc_library(
         "//compiler/src/iree/compiler/Dialect/HAL/IR",
         "//compiler/src/iree/compiler/Dialect/HAL/IR:HALDialect",
         "//compiler/src/iree/compiler/Dialect/HAL/Target",
-        "//compiler/src/iree/compiler/Dialect/HAL/Utils",
         "//compiler/src/iree/compiler/Dialect/Stream/IR",
         "//compiler/src/iree/compiler/Dialect/Util/IR",
         "@llvm-project//llvm:Support",

compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/CMakeLists.txt

@@ -29,7 +29,6 @@ iree_cc_library(
     iree::compiler::Dialect::HAL::IR
     iree::compiler::Dialect::HAL::IR::HALDialect
     iree::compiler::Dialect::HAL::Target
-    iree::compiler::Dialect::HAL::Utils
     iree::compiler::Dialect::Stream::IR
     iree::compiler::Dialect::Util::IR
   PUBLIC

compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.cpp

@@ -9,7 +9,6 @@
 #include "iree/compiler/Dialect/HAL/IR/HALDialect.h"
 #include "iree/compiler/Dialect/HAL/IR/HALOps.h"
 #include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
-#include "iree/compiler/Dialect/HAL/Utils/DeviceSwitchBuilder.h"
 #include "iree/compiler/Dialect/Stream/IR/StreamDialect.h"
 #include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
 #include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
@@ -920,16 +919,6 @@ struct CmdCollectiveOpPattern
   }
 };
 
-// Returns a hal.device.switch match expression that selects the given export.
-static Attribute
-getExportConditionAttr(IREE::HAL::ExecutableExportOp exportOp) {
-  // TODO(benvanik): customizable selection logic. Today this just checks
-  // whether the variant target is supported but we can also allow
-  // specialization of entry points based on dispatch site parameters.
-  auto variantOp = exportOp->getParentOfType<IREE::HAL::ExecutableVariantOp>();
-  return variantOp.getTarget().getMatchExpression();
-}
-
 struct CmdDispatchOpPattern
     : public StreamConversionPattern<IREE::Stream::CmdDispatchOp> {
   using StreamConversionPattern::StreamConversionPattern;
@@ -942,42 +931,65 @@ struct CmdDispatchOpPattern
     // Get the device handle we're executing against in this execution region.
     // Note that this is a dynamic value: we have to treat the device as unknown
     // here.
-    auto device = rewriter.create<IREE::HAL::CommandBufferDeviceOp>(
+    auto deviceValue = rewriter.create<IREE::HAL::CommandBufferDeviceOp>(
         loc, rewriter.getType<IREE::HAL::DeviceType>(), commandBuffer);
 
-    // Ask each target backend to record their dispatch logic.
-    IREE::HAL::DeviceSwitchRewriter switchRewriter(loc,
-                                                   /*resultTypes=*/TypeRange{},
-                                                   device, rewriter);
+    // Prepare for variant switch table by gathering the conditions selecting
+    // each variant.
+    SmallVector<int64_t> caseIndices;
+    SmallVector<std::pair<SymbolRefAttr, IREE::HAL::ExecutableExportOp>>
+        caseExportOps;
     dispatchOp.forEachEntryPointAttr([&](SymbolRefAttr entryPointAttr) {
       // NOTE: slow lookup!
       auto exportOp =
           SymbolTable::lookupNearestSymbolFrom<IREE::HAL::ExecutableExportOp>(
               dispatchOp, entryPointAttr);
       assert(exportOp && "dispatch target export not found");
+      caseIndices.push_back(caseIndices.size());
+      caseExportOps.push_back(std::make_pair(entryPointAttr, exportOp));
+    });
 
-      // Setup the case condition for the entry point.
-      auto *caseRegion =
-          switchRewriter.addConditionRegion(getExportConditionAttr(exportOp));
-      auto &entryBlock = caseRegion->front();
-      auto caseBuilder = OpBuilder::atBlockBegin(&entryBlock);
+    // Select the variant index.
+    Value selectedIndex = buildIfElseTree(
+        loc, caseExportOps.size(),
+        [&](Location loc, size_t i, OpBuilder &builder) {
+          auto exportOp = caseExportOps[i].second;
+          auto variantOp =
+              exportOp->getParentOfType<IREE::HAL::ExecutableVariantOp>();
+          return variantOp.buildCondition(deviceValue, rewriter);
+        },
+        rewriter);
+
+    // Allow each variant to define how it is dispatched.
+    auto switchOp = rewriter.replaceOpWithNewOp<scf::IndexSwitchOp>(
+        dispatchOp, TypeRange{}, selectedIndex, caseIndices,
+        caseIndices.size());
+    for (size_t i = 0; i < caseExportOps.size(); ++i) {
+      auto entryPointAttr = caseExportOps[i].first;
+      auto exportOp = caseExportOps[i].second;
+      auto &caseBlock = switchOp.getCaseRegions()[i].emplaceBlock();
+      auto caseBuilder = OpBuilder::atBlockBegin(&caseBlock);
 
       // Record push constants and buffer bindings.
-      recordParameters(loc, device, commandBuffer, dispatchOp, adaptor,
+      recordParameters(loc, deviceValue, commandBuffer, dispatchOp, adaptor,
                        exportOp.getLayout(), caseBuilder);
 
       // Dispatch with a target-specific workgroup count.
       auto caseWorkgroupCount = exportOp.calculateWorkgroupCount(
-          loc, device, adaptor.getWorkload(), caseBuilder);
+          loc, deviceValue, adaptor.getWorkload(), caseBuilder);
       caseBuilder.create<IREE::HAL::CommandBufferDispatchSymbolOp>(
           loc, commandBuffer, entryPointAttr, caseWorkgroupCount[0],
           caseWorkgroupCount[1], caseWorkgroupCount[2]);
 
-      caseBuilder.create<IREE::HAL::ReturnOp>(loc);
-    });
-    switchRewriter.build();
+      caseBuilder.create<scf::YieldOp>(loc);
+    }
+
+    // Fallback for no available variant. Today we just no-op as executable
+    // loading should have already failed.
+    auto &defaultBlock = switchOp.getDefaultRegion().emplaceBlock();
+    auto defaultBuilder = OpBuilder::atBlockBegin(&defaultBlock);
+    defaultBuilder.create<scf::YieldOp>(loc);
 
-    rewriter.eraseOp(dispatchOp);
     return success();
   }
 

compiler/src/iree/compiler/Dialect/HAL/Conversion/StreamToHAL/test/cmd_ops.mlir

@@ -164,9 +164,10 @@ func.func @cmdExecute(%arg0: !stream.resource<transient>, %arg1: index, %arg2: !
 
 // -----
 
-#executable_target_embedded_elf_x86_64 = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64">
+#executable_target_aarch64 = #hal.executable.target<"llvm-cpu", "embedded-elf-aarch64">
+#executable_target_x86_64 = #hal.executable.target<"llvm-cpu", "embedded-elf-x86_64">
 #device_target_cpu = #hal.device.target<"llvm-cpu", {
-  executable_targets = [#executable_target_embedded_elf_x86_64]
+  executable_targets = [#executable_target_aarch64, #executable_target_x86_64]
 }>
 #pipeline_layout = #hal.pipeline.layout<push_constants = 0, sets = [
   #hal.descriptor_set.layout<0, bindings = [
@@ -177,7 +178,24 @@ func.func @cmdExecute(%arg0: !stream.resource<transient>, %arg1: index, %arg2: !
   ]>
 ]>
 hal.executable private @ex {
-  hal.executable.variant public @embedded_elf_x86_64 target(#executable_target_embedded_elf_x86_64) {
+  hal.executable.variant public @aarch64 target(#executable_target_aarch64) {
+    hal.executable.condition(%device: !hal.device) -> i1 {
+      %ok, %selected = hal.device.query<%device : !hal.device> key("some" :: "feature") : i1, i1
+      hal.return %selected : i1
+    }
+    hal.executable.export public @dispatch ordinal(0) layout(#pipeline_layout) attributes {
+      translation_info = #iree_codegen.translation_info<CPUDefault>
+    } {
+    ^bb0(%device: !hal.device, %arg0: index, %arg1: index, %arg2: index):  // no predecessors
+      %c1 = arith.constant 1 : index
+      %0 = affine.apply affine_map<()[s0] -> (s0 ceildiv 4)>()[%arg0]
+      hal.return %0, %c1, %c1 : index, index, index
+    }
+    builtin.module {
+      // Opaque at this point (in some target-specific dialects).
+    }
+  }
+  hal.executable.variant public @x86_64 target(#executable_target_x86_64) {
     hal.executable.export public @dispatch ordinal(0) layout(#pipeline_layout) attributes {
       translation_info = #iree_codegen.translation_info<CPUDefault>
     } {
@@ -203,9 +221,19 @@ func.func @cmdDispatch(%arg0: !stream.resource<transient>, %arg1: index, %arg2:
   %c128 = arith.constant 128 : index
   // CHECK: %[[CMD:.+]] = hal.command_buffer.create
   %0 = stream.cmd.execute with(%arg0 as %arg4: !stream.resource<transient>{%arg1}, %arg2 as %arg5: !stream.resource<external>{%arg3}) {
-    // Switch for each executable variant:
-    // CHECK: hal.device.switch
-    // CHECK-NEXT: #hal.device.match.executable.format<"embedded-elf-x86_64">
+    // Switch for each executable variant by checking conditions and ranking:
+    // CHECK: %[[DEVICE:.+]] = hal.command_buffer.device<%[[CMD]] : !hal.command_buffer>
+    //  CHECK-DAG: %{{.+}}, %[[AARCH64_FORMAT:.+]] = hal.device.query<%[[DEVICE]] : !hal.device> key("hal.executable.format" :: "embedded-elf-aarch64")
+    //  CHECK-DAG: %[[AARCH64_FEATURE:.+]] = scf.execute_region -> i1 {
+    // CHECK-NEXT:   %{{.+}}, %[[FEATURE:.+]] = hal.device.query<%[[DEVICE]] : !hal.device> key("some" :: "feature")
+    // CHECK-NEXT:   scf.yield %[[FEATURE]]
+    // CHECK-NEXT: }
+    //  CHECK-DAG: %[[AARCH64_SELECTED:.+]] = arith.andi %[[AARCH64_FORMAT]], %[[AARCH64_FEATURE]]
+    //  CHECK-DAG: %{{.+}}, %[[X86_64_SELECTED:.+]] = hal.device.query<%[[DEVICE]] : !hal.device> key("hal.executable.format" :: "embedded-elf-x86_64")
+    // CHECK: %[[VARIANT1:.+]] = arith.select %[[X86_64_SELECTED]], %c1
+    // CHECK: %[[VARIANT0:.+]] = arith.select %[[AARCH64_SELECTED]], %c0{{.+}}, %[[VARIANT1]]
+    // CHECK: scf.index_switch %[[VARIANT0]]
+    // CHECK-NEXT: case 0 {
 
     // Cache queries:
     //  CHECK-DAG:   %[[LAYOUT:.+]] = hal.pipeline_layout.lookup {{.+}} layout(#pipeline_layout)
@@ -230,9 +258,14 @@ func.func @cmdDispatch(%arg0: !stream.resource<transient>, %arg1: index, %arg2:
 
     // Dispatch:
     // CHECK: hal.command_buffer.dispatch.symbol<%[[CMD]]
-    // CHECK-SAME: target(@ex::@embedded_elf_x86_64::@dispatch)
+    // CHECK-SAME: target(@ex::@aarch64::@dispatch)
     // CHECK-SAME: workgroups([%[[X]], %[[YZ]], %[[YZ]]])
-    stream.cmd.dispatch @ex::@embedded_elf_x86_64::@dispatch[%c1, %c2, %c3](%c4_i32, %c5_i32 : i32, i32) {
+
+    // Other variant, when selected:
+    // CHECK: case 1 {
+    // CHECK: hal.command_buffer.dispatch.symbol<%[[CMD]]
+    // CHECK-SAME: target(@ex::@x86_64::@dispatch)
+    stream.cmd.dispatch {@ex::@aarch64::@dispatch, @ex::@x86_64::@dispatch}[%c1, %c2, %c3](%c4_i32, %c5_i32 : i32, i32) {
       ro %arg4[%c0 for %c128] : !stream.resource<transient>{%arg1},
       wo %arg5[%c0 for %c128] : !stream.resource<external>{%arg3}
     } attributes {

compiler/src/iree/compiler/Dialect/HAL/IR/BUILD.bazel

@@ -81,6 +81,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:MemRefDialect",
         "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:SideEffectInterfaces",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TransformUtils",
@@ -105,6 +106,7 @@ iree_compiler_cc_library(
         "@llvm-project//mlir:IR",
         "@llvm-project//mlir:MemRefDialect",
         "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:SCFDialect",
         "@llvm-project//mlir:TransformUtils",
     ],
 )