// Copyright 2021 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/compiler/Dialect/HAL/Conversion/StreamToHAL/Patterns.h"

#include "iree/compiler/Dialect/HAL/Analysis/Captures.h"
#include "iree/compiler/Dialect/HAL/Conversion/StreamToHAL/Utils.h"
#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/Stream/IR/StreamDialect.h"
#include "iree/compiler/Dialect/Stream/IR/StreamOps.h"
#include "iree/compiler/Dialect/Stream/IR/StreamTypes.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Dominance.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/RegionUtils.h"

namespace mlir::iree_compiler {

namespace {

static llvm::cl::opt<bool> clIndirectCommandBuffers{
    "iree-hal-indirect-command-buffers",
    llvm::cl::desc("Whether to turn buffer bindings into indirect references "
                   "when recording command buffers."),
    llvm::cl::init(true),
};

// TODO(benvanik): remove when we support capturing dynamic values for reuse.
static llvm::cl::opt<bool> clForceIndirectCommandBuffers{
    "iree-hal-force-indirect-command-buffers",
    llvm::cl::desc("Forces indirect command buffers when they would otherwise "
                   "not be chosen due to the values they capture. They may not "
                   "be reusable but will still be outlined."),
    llvm::cl::init(false),
};

struct ContextResolveOpPattern
    : public StreamConversionPattern<IREE::Stream::ContextResolveOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::ContextResolveOp resolveOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto resultTypes = llvm::to_vector(resolveOp.getResultTypes());
    assert(!resultTypes.empty() && "must have at least one result");

    // Get the affinity from the op or an ancestor. Note that there may be no
    // affinity specified at all.
    auto affinityAttr = IREE::Stream::AffinityAttr::lookupOrDefault(resolveOp);

    // If no affinity was specified then resolve as 'any'.
    if (!affinityAttr) {
      rewriter.replaceOpWithNewOp<IREE::HAL::DeviceResolveOp>(
          resolveOp, resolveOp.getResultTypes(),
          IREE::HAL::DeviceAffinityAttr{});
      return success();
    }

    // We currently only handle HAL device affinities.
    // We could make this an interface to select the device and allow users to
    // provide their own affinities to convert to HAL. In the future users may
    // also want to provide devices as function arguments post-initialization.
    // For now we just have one way to specify device globals.
    if (auto deviceAffinityAttr =
            dyn_cast_if_present<IREE::HAL::DeviceAffinityAttr>(affinityAttr)) {
      rewriter.replaceOpWithNewOp<IREE::HAL::DeviceResolveOp>(
          resolveOp, resolveOp.getResultTypes(), deviceAffinityAttr);
      return success();
    }

    resolveOp.emitOpError() << "failed to resolve affinity: only HAL device "
                               "affinities are supported";
    return rewriter.notifyMatchFailure(
        resolveOp, "only HAL device affinities are supported");
  }
};

struct ResourceAllocOpPattern
    : public StreamConversionPattern<IREE::Stream::ResourceAllocOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::ResourceAllocOp allocOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto [allocator, queueAffinity] =
        lookupAllocatorAndQueueAffinityFor(allocOp, rewriter);
    auto bufferType = rewriter.getType<IREE::HAL::BufferType>();

    auto resourceType =
        cast<IREE::Stream::ResourceType>(allocOp.getResult().getType());

    auto memoryTypes = IREE::HAL::MemoryTypeBitfield::None;
    auto bufferUsage = IREE::HAL::BufferUsageBitfield::None;
    if (failed(deriveAllowedResourceBufferBits(allocOp.getLoc(), resourceType,
                                               memoryTypes, bufferUsage))) {
      return failure();
    }

    rewriter.replaceOpWithNewOp<IREE::HAL::AllocatorAllocateOp>(
        allocOp, bufferType, allocator, queueAffinity, memoryTypes, bufferUsage,
        adaptor.getStorageSize());
    return success();
  }
};

struct ResourceAllocaOpPattern
    : public StreamConversionPattern<IREE::Stream::ResourceAllocaOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::ResourceAllocaOp allocaOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = allocaOp.getLoc();
    auto [device, queueAffinity] =
        lookupDeviceAndQueueAffinityFor(allocaOp, rewriter);
    auto bufferType = rewriter.getType<IREE::HAL::BufferType>();

    auto resourceType =
        cast<IREE::Stream::ResourceType>(allocaOp.getResult().getType());
    auto memoryTypes = IREE::HAL::MemoryTypeBitfield::None;
    auto bufferUsage = IREE::HAL::BufferUsageBitfield::None;
    if (failed(deriveAllowedResourceBufferBits(loc, resourceType, memoryTypes,
                                               bufferUsage))) {
      return failure();
    }

    // Gather wait/signal fence, which are optional.
    Value waitFence =
        getOrCreateWaitFence(loc, adaptor.getAwaitTimepoint(), rewriter);
    Value signalFence = getOrCreateSignalFence(
        loc, device, allocaOp.getResultTimepoint(), rewriter);

    // Queue allocation.
    auto pool = rewriter.create<arith::ConstantIntOp>(loc, 0, 64);
    auto allocateOp = rewriter.create<IREE::HAL::DeviceQueueAllocaOp>(
        loc, bufferType, device, queueAffinity, waitFence, signalFence, pool,
        memoryTypes, bufferUsage, adaptor.getStorageSize());

    rewriter.replaceOp(allocaOp, {allocateOp.getResult(), signalFence});
    return success();
  }
};

struct ResourceDeallocaOpPattern
    : public StreamConversionPattern<IREE::Stream::ResourceDeallocaOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::ResourceDeallocaOp deallocaOp,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = deallocaOp.getLoc();
    auto [device, queueAffinity] =
        lookupDeviceAndQueueAffinityFor(deallocaOp, rewriter);

    // Gather wait/signal fence, which are optional.
    Value waitFence =
        getOrCreateWaitFence(loc, adaptor.getAwaitTimepoint(), rewriter);
    Value signalFence = getOrCreateSignalFence(
        loc, device, deallocaOp.getResultTimepoint(), rewriter);

    // Queue allocation.
    rewriter.create<IREE::HAL::DeviceQueueDeallocaOp>(
        loc, device, queueAffinity, waitFence, signalFence,
        adaptor.getOperand());

    rewriter.replaceOp(deallocaOp, {signalFence});
    return success();
  }
};

struct ResourceSizeOpPattern
    : public StreamConversionPattern<IREE::Stream::ResourceSizeOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::ResourceSizeOp sizeOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    rewriter.replaceOpWithNewOp<IREE::HAL::BufferLengthOp>(
        sizeOp, rewriter.getIndexType(), adaptor.getOperand());
    return success();
  }
};

struct ResourceTryMapOpPattern
    : public StreamConversionPattern<IREE::Stream::ResourceTryMapOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::ResourceTryMapOp tryMapOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto [allocator, queueAffinity] =
        lookupAllocatorAndQueueAffinityFor(tryMapOp, rewriter);
    auto resourceType =
        llvm::cast<IREE::Stream::ResourceType>(tryMapOp.getResult().getType());
    auto bufferType = rewriter.getType<IREE::HAL::BufferType>();

    auto memoryTypes = IREE::HAL::MemoryTypeBitfield::None;
    auto bufferUsage = IREE::HAL::BufferUsageBitfield::None;
    switch (resourceType.getLifetime()) {
    default:
      return tryMapOp.emitOpError()
             << "unsupported resource lifetime: "
             << IREE::Stream::stringifyLifetime(resourceType.getLifetime());
    case IREE::Stream::Lifetime::Constant:
      // Device local; copies required to get into external resources.
      memoryTypes = memoryTypes | IREE::HAL::MemoryTypeBitfield::DeviceLocal;
      bufferUsage =
          bufferUsage | IREE::HAL::BufferUsageBitfield::SharingImmutable;
      // TODO(benvanik): refine usage based on analysis.
      bufferUsage = bufferUsage | IREE::HAL::BufferUsageBitfield::Transfer |
                    IREE::HAL::BufferUsageBitfield::DispatchStorage;
      break;
    case IREE::Stream::Lifetime::Variable:
      // Device local; copies required to get into external resources.
      memoryTypes = memoryTypes | IREE::HAL::MemoryTypeBitfield::DeviceLocal;
      // TODO(benvanik): refine usage based on analysis.
      bufferUsage = bufferUsage | IREE::HAL::BufferUsageBitfield::Transfer |
                    IREE::HAL::BufferUsageBitfield::DispatchStorage;
      break;
    case IREE::Stream::Lifetime::Staging:
      // Host local; copies required to get into device resources.
      // We could vary this based on staging usage (upload/download) by
      // making it device-local|host-visible, but host-local means we have
      // a better chance of mapping it during uploads.
      memoryTypes = memoryTypes | IREE::HAL::MemoryTypeBitfield::HostVisible |
                    IREE::HAL::MemoryTypeBitfield::DeviceVisible;
      bufferUsage =
          bufferUsage | IREE::HAL::BufferUsageBitfield::TransferSource;
      break;
    }

    rewriter.replaceOpWithNewOp<IREE::HAL::AllocatorImportOp>(
        tryMapOp, rewriter.getI1Type(), bufferType, allocator, queueAffinity,
        memoryTypes, bufferUsage, adaptor.getSource(),
        adaptor.getSourceOffset(), adaptor.getResultSize());
    return success();
  }
};

struct ResourceLoadOpPattern
    : public StreamConversionPattern<IREE::Stream::ResourceLoadOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::ResourceLoadOp loadOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loadType =
        getTypeConverter()->convertType(loadOp.getResult().getType());
    rewriter.replaceOpWithNewOp<IREE::HAL::BufferLoadOp>(
        loadOp, loadType, adaptor.getSource(), adaptor.getSourceOffset());
    return success();
  }
};

struct ResourceStoreOpPattern
    : public StreamConversionPattern<IREE::Stream::ResourceStoreOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::ResourceStoreOp storeOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    rewriter.replaceOpWithNewOp<IREE::HAL::BufferStoreOp>(
        storeOp, adaptor.getValue(), adaptor.getTarget(),
        adaptor.getTargetOffset());
    return success();
  }
};

struct ResourceSubviewOpPattern
    : public StreamConversionPattern<IREE::Stream::ResourceSubviewOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::ResourceSubviewOp subviewOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto bufferType = rewriter.getType<IREE::HAL::BufferType>();
    // NOTE: this aliases! We assume at this point all useful alias analysis
    // has been performed and it's fine to lose the tie information here.
    rewriter.replaceOpWithNewOp<IREE::HAL::BufferSubspanOp>(
        subviewOp, bufferType, adaptor.getSource(), adaptor.getSourceOffset(),
        adaptor.getResultSize());
    return success();
  }
};

struct FileConstantOpPattern
    : public StreamConversionPattern<IREE::Stream::FileConstantOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::FileConstantOp constantOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto [device, queueAffinity] =
        lookupDeviceAndQueueAffinityFor(constantOp, rewriter);
    rewriter.replaceOpWithNewOp<IREE::HAL::ExFileFromMemoryOp>(
        constantOp, rewriter.getType<IREE::HAL::FileType>(), device,
        queueAffinity, IREE::HAL::MemoryAccessBitfield::Read,
        constantOp.getSource(), constantOp.getSourceOffset(),
        constantOp.getSourceLength(),
        rewriter.create<arith::ConstantIntOp>(constantOp.getLoc(), 0, 32));
    return success();
  }
};

struct FileReadOpPattern
    : public StreamConversionPattern<IREE::Stream::FileReadOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::FileReadOp readOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = readOp.getLoc();
    auto [device, queueAffinity] =
        lookupDeviceAndQueueAffinityFor(readOp, rewriter);

    // Gather wait/signal fence, which are optional.
    Value waitFence =
        getOrCreateWaitFence(loc, adaptor.getAwaitTimepoint(), rewriter);
    Value signalFence = getOrCreateSignalFence(
        loc, device, readOp.getResultTimepoint(), rewriter);

    // Queue read.
    rewriter.create<IREE::HAL::DeviceQueueReadOp>(
        loc, device, queueAffinity, waitFence, signalFence, adaptor.getSource(),
        adaptor.getSourceOffset(), adaptor.getTarget(),
        adaptor.getTargetOffset(), adaptor.getLength(),
        /*flags=*/0);

    rewriter.replaceOp(readOp, {signalFence});
    return success();
  }
};

struct FileWriteOpPattern
    : public StreamConversionPattern<IREE::Stream::FileWriteOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::FileWriteOp writeOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = writeOp.getLoc();
    auto [device, queueAffinity] =
        lookupDeviceAndQueueAffinityFor(writeOp, rewriter);

    // Gather wait/signal fence, which are optional.
    Value waitFence =
        getOrCreateWaitFence(loc, adaptor.getAwaitTimepoint(), rewriter);
    Value signalFence = getOrCreateSignalFence(
        loc, device, writeOp.getResultTimepoint(), rewriter);

    // Queue write.
    rewriter.create<IREE::HAL::DeviceQueueWriteOp>(
        loc, device, queueAffinity, waitFence, signalFence, adaptor.getSource(),
        adaptor.getSourceOffset(), adaptor.getTarget(),
        adaptor.getTargetOffset(), adaptor.getLength(),
        /*flags=*/0);

    rewriter.replaceOp(writeOp, {signalFence});
    return success();
  }
};

// Inserts IR to assert that the underlying buffer storage is compatible with
// the intended usage in the program. The allocator used to allocate the
// buffer must have compatibility with our target device allocator and the
// buffer must have at least the minimum expected size (additional padding is
// ok).
static LogicalResult
buildStorageAssertions(Location loc, Value buffer, StringAttr message,
                       Value allocator, Value minimumLength,
                       IREE::Stream::ResourceType resourceType,
                       OpBuilder &builder) {
  auto memoryTypes = IREE::HAL::MemoryTypeBitfield::None;
  auto bufferUsage = IREE::HAL::BufferUsageBitfield::None;
  if (failed(deriveRequiredResourceBufferBits(loc, resourceType, memoryTypes,
                                              bufferUsage))) {
    return failure();
  }

  auto requiredTypes =
      IREE::HAL::MemoryTypeBitfieldAttr::get(builder.getContext(), memoryTypes);
  auto requiredUsage = IREE::HAL::BufferUsageBitfieldAttr::get(
      builder.getContext(), bufferUsage);

  builder.create<IREE::HAL::BufferAssertOp>(loc, buffer, message, allocator,
                                            minimumLength, requiredTypes,
                                            requiredUsage);
  return success();
}

struct TensorImportBufferOpPattern
    : public StreamConversionPattern<IREE::Stream::TensorImportOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::TensorImportOp importOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (!llvm::isa<IREE::HAL::BufferType>(importOp.getSource().getType())) {
      return failure();
    }

    // TODO(benvanik): get a name for the tensor (argument name/etc).
    auto message = rewriter.getStringAttr("tensor");

    // Directly use the buffer.
    auto buffer = adaptor.getSource();
    rewriter.replaceOp(importOp, buffer);

    // Assert the storage is compatible with our expected device and usage.
    auto targetAllocator = lookupAllocatorFor(importOp, rewriter);
    auto resourceType =
        llvm::cast<IREE::Stream::ResourceType>(importOp.getResult().getType());
    if (failed(buildStorageAssertions(
            importOp.getLoc(), adaptor.getSource(), message, targetAllocator,
            adaptor.getResultSize(), resourceType, rewriter))) {
      return failure();
    }

    return success();
  }
};

struct TensorImportBufferViewOpPattern
    : public StreamConversionPattern<IREE::Stream::TensorImportOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::TensorImportOp importOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto sourceType = importOp.getSource().getType();
    if (!llvm::isa<IREE::HAL::BufferViewType>(sourceType) &&
        !llvm::isa<TensorType>(sourceType)) {
      return failure();
    }

    auto loc = importOp.getLoc();

    // TODO(benvanik): get a name for the tensor (argument name/etc).
    auto message = rewriter.getStringAttr("tensor");

    auto bufferView = adaptor.getSource();
    auto bufferType = rewriter.getType<IREE::HAL::BufferType>();
    auto bufferOp = rewriter.replaceOpWithNewOp<IREE::HAL::BufferViewBufferOp>(
        importOp, bufferType, bufferView);

    // Assert the storage is compatible with our expected device and usage.
    auto targetAllocator = lookupAllocatorFor(importOp, rewriter);
    auto resourceType =
        llvm::cast<IREE::Stream::ResourceType>(importOp.getResult().getType());
    if (failed(buildStorageAssertions(loc, bufferOp.getResult(), message,
                                      targetAllocator, adaptor.getResultSize(),
                                      resourceType, rewriter))) {
      return failure();
    }

    return success();
  }
};

struct TensorExportBufferOpPattern
    : public StreamConversionPattern<IREE::Stream::TensorExportOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::TensorExportOp exportOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (!llvm::isa<IREE::HAL::BufferType>(exportOp.getResult().getType())) {
      return failure();
    }
    rewriter.replaceOp(exportOp, adaptor.getSource());
    return success();
  }
};

struct TensorExportBufferViewOpPattern
    : public StreamConversionPattern<IREE::Stream::TensorExportOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::TensorExportOp exportOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto targetType = exportOp.getResult().getType();
    if (!llvm::isa<IREE::HAL::BufferViewType>(targetType) &&
        !llvm::isa<TensorType>(targetType)) {
      return failure();
    }

    auto loc = exportOp.getLoc();
    auto tensorType = llvm::cast<RankedTensorType>(adaptor.getSourceEncoding());
    auto dynamicDims = adaptor.getSourceEncodingDims();

    // NOTE: we should have verified supported encodings/types at entry into the
    // HAL pipeline.
    auto encodingType = rewriter.create<IREE::HAL::EncodingTypeOp>(
        loc, tensorType.getEncoding());
    auto elementType = rewriter.create<IREE::HAL::ElementTypeOp>(
        loc, tensorType.getElementType());

    // Flatten static + dynamic shape dimensions.
    SmallVector<Value> dims;
    unsigned dynamicIdx = 0;
    for (int64_t idx = 0; idx < tensorType.getRank(); ++idx) {
      if (tensorType.isDynamicDim(idx)) {
        dims.push_back(dynamicDims[dynamicIdx++]);
      } else {
        dims.push_back(rewriter.create<arith::ConstantIndexOp>(
            loc, tensorType.getDimSize(idx)));
      }
    }

    rewriter.replaceOpWithNewOp<IREE::HAL::BufferViewCreateOp>(
        exportOp, adaptor.getSource(),
        rewriter.create<arith::ConstantIndexOp>(loc, 0),
        adaptor.getSourceSize(), elementType, encodingType, dims);
    return success();
  }
};

struct TensorTraceOpPattern
    : public StreamConversionPattern<IREE::Stream::TensorTraceOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::TensorTraceOp traceOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    SmallVector<Value> bufferViews;
    auto resourceEncodingDims = adaptor.getResourceEncodingDims();
    for (auto [resource, resourceSize, resourceEncoding] : llvm::zip_equal(
             adaptor.getResources(), adaptor.getResourceSizes(),
             adaptor.getResourceEncodings().getAsRange<TypeAttr>())) {
      int64_t dynamicDimCount =
          cast<ShapedType>(resourceEncoding.getValue()).getNumDynamicDims();
      bufferViews.push_back(rewriter.create<IREE::Stream::TensorExportOp>(
          traceOp.getLoc(), rewriter.getType<IREE::HAL::BufferViewType>(),
          resource, resourceEncoding,
          resourceEncodingDims.take_front(dynamicDimCount), resourceSize,
          /*affinity=*/IREE::Stream::AffinityAttr{}));
      resourceEncodingDims = resourceEncodingDims.drop_front(dynamicDimCount);
    }
    rewriter.replaceOpWithNewOp<IREE::HAL::BufferViewTraceOp>(
        traceOp, traceOp.getKeyAttr(), bufferViews);
    return success();
  }
};

struct CmdFlushOpPattern
    : public StreamConversionPattern<IREE::Stream::CmdFlushOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::CmdFlushOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // TODO(benvanik): HAL command buffer op for flush.
    rewriter.eraseOp(op);
    return success();
  }
};

struct CmdInvalidateOpPattern
    : public StreamConversionPattern<IREE::Stream::CmdInvalidateOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::CmdInvalidateOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // TODO(benvanik): HAL command buffer op for invalidate.
    rewriter.eraseOp(op);
    return success();
  }
};

struct CmdDiscardOpPattern
    : public StreamConversionPattern<IREE::Stream::CmdDiscardOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::CmdDiscardOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // TODO(benvanik): HAL command buffer op for discard.
    rewriter.eraseOp(op);
    return success();
  }
};

struct CmdFillOpPattern
    : public StreamConversionPattern<IREE::Stream::CmdFillOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::CmdFillOp fillOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto commandBufferMapping = mapping->lookupCommandBufferFor(fillOp);
    auto targetBinding = commandBufferMapping.resolveBinding(
        fillOp.getLoc(), fillOp.getTarget(), adaptor.getTarget(),
        adaptor.getTargetOffset(), adaptor.getTargetLength(), rewriter);
    rewriter.replaceOpWithNewOp<IREE::HAL::CommandBufferFillBufferOp>(
        fillOp, commandBufferMapping.getHandle(), targetBinding.buffer,
        targetBinding.byteOffset, targetBinding.byteLength, adaptor.getValue());
    return success();
  }
};

struct CmdCopyOpPattern
    : public StreamConversionPattern<IREE::Stream::CmdCopyOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::CmdCopyOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto commandBufferMapping = mapping->lookupCommandBufferFor(op);
    auto sourceBinding = commandBufferMapping.resolveBinding(
        op.getLoc(), op.getSource(), adaptor.getSource(),
        adaptor.getSourceOffset(), adaptor.getLength(), rewriter);
    auto targetBinding = commandBufferMapping.resolveBinding(
        op.getLoc(), op.getTarget(), adaptor.getTarget(),
        adaptor.getTargetOffset(), adaptor.getLength(), rewriter);
    rewriter.replaceOpWithNewOp<IREE::HAL::CommandBufferCopyBufferOp>(
        op, commandBufferMapping.getHandle(), sourceBinding.buffer,
        sourceBinding.byteOffset, targetBinding.buffer,
        targetBinding.byteOffset, adaptor.getLength());
    return success();
  }
};

// NOTE: this relies on the enums being the same today. Ew.
static IREE::HAL::CollectiveAttr
convertCollectiveAttr(IREE::Stream::CollectiveAttr sourceAttr) {
  auto convertReductionOp =
      [](std::optional<IREE::Stream::CollectiveReductionOp> op)
      -> std::optional<IREE::HAL::CollectiveReductionOp> {
    if (!op.has_value())
      return std::nullopt;
    return static_cast<IREE::HAL::CollectiveReductionOp>(op.value());
  };
  return IREE::HAL::CollectiveAttr::get(
      sourceAttr.getContext(),
      static_cast<IREE::HAL::CollectiveKind>(sourceAttr.getKind()),
      convertReductionOp(sourceAttr.getReduction()),
      static_cast<IREE::HAL::CollectiveElementType>(
          sourceAttr.getElementType()));
}

struct CmdCollectiveOpPattern
    : public StreamConversionPattern<IREE::Stream::CmdCollectiveOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::CmdCollectiveOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto commandBufferMapping = mapping->lookupCommandBufferFor(op);

    IREE::HAL::BindingValue sendBinding;
    IREE::HAL::BindingValue recvBinding;
    switch (adaptor.getOp().getKind()) {
    default:
      assert(adaptor.getResources().size() == 2 && "should have verified");
      sendBinding = commandBufferMapping.resolveBinding(
          op.getLoc(), op.getResources()[0], adaptor.getResources()[0],
          adaptor.getResourceOffsets()[0], adaptor.getResourceLengths()[0],
          rewriter);
      recvBinding = commandBufferMapping.resolveBinding(
          op.getLoc(), op.getResources()[1], adaptor.getResources()[1],
          adaptor.getResourceOffsets()[1], adaptor.getResourceLengths()[1],
          rewriter);
      break;
    case IREE::Stream::CollectiveKind::Send:
      assert(adaptor.getResources().size() == 1 && "should have verified");
      sendBinding = commandBufferMapping.resolveBinding(
          op.getLoc(), op.getResources()[0], adaptor.getResources()[0],
          adaptor.getResourceOffsets()[0], adaptor.getResourceLengths()[0],
          rewriter);
      break;
    case IREE::Stream::CollectiveKind::Recv:
      assert(adaptor.getResources().size() == 1 && "should have verified");
      recvBinding = commandBufferMapping.resolveBinding(
          op.getLoc(), op.getResources()[0], adaptor.getResources()[0],
          adaptor.getResourceOffsets()[0], adaptor.getResourceLengths()[0],
          rewriter);
      break;
    }

    rewriter.replaceOpWithNewOp<IREE::HAL::CommandBufferCollectiveOp>(
        op, commandBufferMapping.getHandle(), adaptor.getChannel(),
        convertCollectiveAttr(adaptor.getOp()), adaptor.getElementCount(),
        adaptor.getParam(), sendBinding.buffer, sendBinding.byteOffset,
        sendBinding.byteLength, recvBinding.buffer, recvBinding.byteOffset,
        recvBinding.byteLength);
    return success();
  }
};

struct CmdDispatchOpPattern
    : public StreamConversionPattern<IREE::Stream::CmdDispatchOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::CmdDispatchOp dispatchOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = dispatchOp.getLoc();
    auto commandBufferMapping = mapping->lookupCommandBufferFor(dispatchOp);

    // TODO(multi-device): reusable command buffers done at the stream level may
    // make this difficult. For now we assume each stream region being lowered
    // has a singular affinity that may itself reference multiple devices in the
    // future but currently uniquely identifies a device.
    auto affinityAttr = IREE::Stream::AffinityAttr::lookupOrDefault(dispatchOp);

    // 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.
    Value device = rewriter.create<IREE::HAL::CommandBufferDeviceOp>(
        loc, rewriter.getType<IREE::HAL::DeviceType>(),
        commandBufferMapping.getHandle());

    // 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));
    });

    // If there is only one variant we can emit that directly without a
    // conditional check. The same result should occur later on but it saves
    // a lot of IR during generation if we know we can avoid it.
    if (caseExportOps.size() == 1) {
      auto [entryPointAttr, exportOp] = caseExportOps.front();
      rewriter.replaceOp(dispatchOp,
                         emitDispatchOp(loc, affinityAttr, device,
                                        commandBufferMapping, exportOp,
                                        entryPointAttr, dispatchOp, adaptor,
                                        rewriter));
    } else {
      // 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(device, rewriter);
          },
          rewriter);

      // Allow each variant to define how it is dispatched.
      auto switchOp = rewriter.create<scf::IndexSwitchOp>(
          loc, TypeRange{}, selectedIndex, caseIndices, caseIndices.size());
      for (size_t i = 0; i < caseExportOps.size(); ++i) {
        auto [entryPointAttr, exportOp] = caseExportOps[i];
        auto &caseBlock = switchOp.getCaseRegions()[i].emplaceBlock();
        auto caseBuilder = OpBuilder::atBlockBegin(&caseBlock);
        emitDispatchOp(loc, affinityAttr, device, commandBufferMapping,
                       exportOp, entryPointAttr, dispatchOp, adaptor,
                       caseBuilder);
        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.replaceOp(dispatchOp, switchOp);
    }

    return success();
  }

  Operation *emitDispatchOp(
      Location loc, IREE::Stream::AffinityAttr affinityAttr, Value device,
      CommandBufferConversionMapping &commandBufferMapping,
      IREE::HAL::ExecutableExportOp exportOp, SymbolRefAttr entryPointAttr,
      IREE::Stream::CmdDispatchOp dispatchOp, OpAdaptor adaptor,
      OpBuilder &builder) const {
    auto workgroupCount = exportOp.calculateWorkgroupCount(
        loc, device, adaptor.getWorkload(), builder);

    Value executable = builder.create<IREE::HAL::ExecutableLookupOp>(
        loc, builder.getType<IREE::HAL::ExecutableType>(), device,
        entryPointAttr.getRootReference().getValue());
    Value ordinal = builder.create<IREE::HAL::ExecutableExportOrdinalOp>(
        loc, builder.getIndexType(), entryPointAttr);

    auto layoutAttr = exportOp.getLayout();
    SmallVector<IREE::HAL::BindingValue> bindings;
    for (auto [i, bindingAttr] : llvm::enumerate(layoutAttr.getBindings())) {
      auto descriptorFlags = bindingAttr.getFlags();
      IREE::HAL::BindingValue binding;
      if (bitEnumContainsAll(descriptorFlags,
                             IREE::HAL::DescriptorFlags::Indirect)) {
        // Indirect binding resolved through the cached command buffer binding
        // table. The buffer recorded in the descriptor is a slot ordinal into
        // the binding table. Note that the range may be adjusted based on the
        // range bound to the slot in the table.
        auto resolvedBinding = commandBufferMapping.resolveBinding(
            loc, dispatchOp.getResources()[i], adaptor.getResources()[i],
            adaptor.getResourceOffsets()[i], adaptor.getResourceLengths()[i],
            builder);
        binding.buffer = resolvedBinding.buffer;
        binding.byteOffset = resolvedBinding.byteOffset;
        binding.byteLength = resolvedBinding.byteLength;
      } else {
        // Direct binding referencing the buffer and range provided on the op.
        binding.buffer = adaptor.getResources()[i];
        binding.byteOffset = adaptor.getResourceOffsets()[i];
        binding.byteLength = adaptor.getResourceLengths()[i];
      }
      bindings.push_back(binding);
    }

    auto flags = IREE::HAL::DispatchFlags::None;

    return builder.create<IREE::HAL::CommandBufferDispatchOp>(
        loc, commandBufferMapping.getHandle(), executable, ordinal,
        workgroupCount, adaptor.getUniformOperands(), bindings, flags);
  }
};

struct CmdFuncOpPattern
    : public StreamConversionPattern<IREE::Stream::CmdFuncOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::CmdFuncOp funcOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    SmallVector<DictionaryAttr> oldArgAttrs;
    funcOp.getAllArgAttrs(oldArgAttrs);
    SmallVector<Type> newArgTypes;
    SmallVector<Attribute> newArgAttrs;
    newArgTypes.push_back(rewriter.getType<IREE::HAL::CommandBufferType>());
    newArgAttrs.push_back(rewriter.getDictionaryAttr({})); // command buffer
    for (auto [i, oldType] : llvm::enumerate(funcOp.getArgumentTypes())) {
      if (isa<IREE::Stream::ResourceType>(oldType)) {
        // Resource converted into a (binding ordinal, buffer) pair.
        newArgTypes.push_back(rewriter.getIndexType());
        newArgAttrs.push_back(rewriter.getDictionaryAttr({}));
        newArgTypes.push_back(rewriter.getType<IREE::HAL::BufferType>());
        newArgAttrs.push_back(oldArgAttrs[i]);
      } else {
        // Primitive/other pass-through.
        // Support expansion by preserving the arg attr on the first expanded
        // type and filling in empty attrs for the remainder.
        size_t oldCount = newArgTypes.size();
        if (failed(getTypeConverter()->convertType(oldType, newArgTypes))) {
          return rewriter.notifyMatchFailure(funcOp,
                                             "failed to convert arg types");
        }
        size_t typeCount = newArgTypes.size() - oldCount;
        newArgAttrs.push_back(oldArgAttrs[i]);
        newArgAttrs.append(typeCount - 1, rewriter.getDictionaryAttr({}));
      }
    }
    SmallVector<Type> newResultTypes;
    if (failed(getTypeConverter()->convertTypes(funcOp.getResultTypes(),
                                                newResultTypes))) {
      return rewriter.notifyMatchFailure(funcOp,
                                         "failed to convert result types");
    }
    auto newOp = rewriter.replaceOpWithNewOp<IREE::Util::FuncOp>(
        funcOp, funcOp.getNameAttr(),
        rewriter.getFunctionType(newArgTypes, newResultTypes),
        /*tied_operands=*/ArrayAttr{}, funcOp.getSymVisibilityAttr(),
        rewriter.getArrayAttr(newArgAttrs), funcOp.getAllResultAttrs(),
        /*inlining_policy=*/IREE::Util::InliningPolicyAttrInterface{});
    newOp->setDialectAttrs(funcOp->getDialectAttrs());
    return success();
  }
};

struct CmdCallOpPattern
    : public StreamConversionPattern<IREE::Stream::CmdCallOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::CmdCallOp callOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto commandBufferMapping = mapping->lookupCommandBufferFor(callOp);

    // Memoized dummy values.
    Value zeroIndex;
    auto getZeroIndex = [&]() {
      if (!zeroIndex) {
        zeroIndex = rewriter.create<arith::ConstantIndexOp>(callOp.getLoc(), 0);
      }
      return zeroIndex;
    };
    Value nullBuffer;
    auto getNullBuffer = [&]() {
      if (!nullBuffer) {
        nullBuffer = rewriter.create<IREE::Util::NullOp>(
            callOp.getLoc(), rewriter.getType<IREE::HAL::BufferType>());
      }
      return nullBuffer;
    };

    // Always pass the command buffer as the first arg.
    SmallVector<Value> operands;
    operands.push_back(commandBufferMapping.getHandle());
    size_t resourceIndex = 0;
    for (auto [originalOperand, convertedOperand] : llvm::zip_equal(
             callOp.getResourceOperands(), adaptor.getResourceOperands())) {
      if (llvm::isa<IREE::Stream::ResourceType>(originalOperand.getType())) {
        // Resource type, pass binding index or buffer and offset/length.
        auto binding = commandBufferMapping.resolveBinding(
            callOp.getLoc(), originalOperand, convertedOperand,
            adaptor.getResourceOperandOffsets()[resourceIndex],
            adaptor.getResourceOperandLengths()[resourceIndex], rewriter);
        if (binding.buffer.getType().isIndex()) {
          operands.push_back(binding.buffer);
          operands.push_back(getNullBuffer());
        } else {
          operands.push_back(getZeroIndex());
          operands.push_back(binding.buffer);
        }
        operands.push_back(binding.byteOffset);
        operands.push_back(binding.byteLength);
        ++resourceIndex;
      } else {
        // Primitive/custom type.
        operands.push_back(convertedOperand);
      }
    }

    SmallVector<Type> resultTypes;
    for (auto result : callOp.getResults()) {
      SmallVector<Type> convertedTypes;
      if (failed(getTypeConverter()->convertType(result.getType(),
                                                 convertedTypes))) {
        return rewriter.notifyMatchFailure(callOp.getLoc(),
                                           "unconvertable result type");
      }
      llvm::append_range(resultTypes, convertedTypes);
    }

    rewriter.replaceOpWithNewOp<IREE::Util::CallOp>(
        callOp, resultTypes, callOp.getCallee(), operands,
        /*tied_operands=*/ArrayAttr{});
    return success();
  }
};

// Returns true if any primitive uniform value (i32, index, etc) captured within
// |op| (but not _by_ op) is a dynamic value (mutable global, calculated, etc).
// Returns false if all values are derived from constants or immutable globals.
static bool regionCapturesDynamicUniformValues(Operation *op) {
  auto isDynamicUniform = [](Value value) {
    if (value.getType().isIntOrIndexOrFloat()) {
      switch (IREE::HAL::categorizeValue(value)) {
      default:
      case IREE::HAL::ValueOrigin::Unknown:
      case IREE::HAL::ValueOrigin::MutableGlobal:
        return true;
      case IREE::HAL::ValueOrigin::LocalConstant:
      case IREE::HAL::ValueOrigin::ImmutableGlobal:
        return false;
      }
    }
    return false;
  };
  for (auto operand : op->getOperands()) {
    if (isDynamicUniform(operand)) {
      // Today this usually indicates a dynamic buffer size. We could perform
      // some tricks to adjust the size based on usage instead of requiring that
      // this size match however it's safer to treat dynamically sized buffers
      // as fully dynamic for now.
      return true;
    }
  }
  SetVector<Value> capturedValues;
  mlir::getUsedValuesDefinedAbove(op->getRegions(), capturedValues);
  for (auto capturedValue : capturedValues) {
    if (isDynamicUniform(capturedValue)) {
      return true;
    }
  }
  return false;
}

static void insertSerializationBarriers(Location loc, Block &block,
                                        Value commandBuffer,
                                        OpBuilder builder) {
  // TODO(benvanik): derive based on the type of operations that surround the
  // barriers. Can use deriveCommandCategories on the ranges to see what kind
  // of ops happen above and below, but really some analysis is required.
  auto sourceStage = IREE::HAL::ExecutionStageBitfield::CommandRetire |
                     IREE::HAL::ExecutionStageBitfield::Transfer |
                     IREE::HAL::ExecutionStageBitfield::Dispatch;
  auto targetStage = IREE::HAL::ExecutionStageBitfield::CommandIssue |
                     IREE::HAL::ExecutionStageBitfield::Transfer |
                     IREE::HAL::ExecutionStageBitfield::Dispatch;
  auto flags = IREE::HAL::ExecutionBarrierFlagBitfield::None;

  // Insert barriers after every op.
  // Note that we can't mutate the block while iterating it so we first grab
  // all the original ops.
  SmallVector<Operation *> serialOps;
  for (auto &op : block)
    serialOps.push_back(&op);
  for (auto *op : serialOps) {
    if (op->hasTrait<OpTrait::IsTerminator>())
      continue;
    builder.setInsertionPointAfter(op);
    builder.create<IREE::HAL::CommandBufferExecutionBarrierOp>(
        loc, commandBuffer, sourceStage, targetStage, flags);
  }
}

// Checks if |executeOp| contains only a single transfer operation and returns
// it. Non-transfer/dispatch operations like cache control will be ignored.
//
// Intended to match things like:
//   stream.cmd.execute ... {
//     stream.cmd.invalidate
//     stream.cmd.fill        <----- returned
//     stream.cmd.flush
//   }
// And not:
//   stream.cmd.execute ... {
//     stream.cmd.invalidate
//     stream.cmd.fill
//     stream.cmd.flush
//     stream.cmd.dispatch
//   }
static Operation *matchSingleTransferOp(IREE::Stream::CmdExecuteOp executeOp) {
  Operation *foundOp = nullptr;
  for (auto &block : executeOp.getBodyRegion()) {
    for (auto &op : block) {
      if (!TypeSwitch<Operation *, bool>(&op)
               // Ignore non-transfer/dispatch ops.
               .Case<IREE::Stream::CmdInvalidateOp, IREE::Stream::CmdFlushOp,
                     IREE::Stream::CmdDiscardOp, IREE::Stream::YieldOp>(
                   [&](auto metaOp) { return true; })
               .Case<IREE::Stream::CmdFillOp, IREE::Stream::CmdCopyOp>(
                   [&](auto transferOp) {
                     if (!foundOp) {
                       foundOp = &op; // first found
                       return true;
                     } else {
                       return false; // more than one transfer op
                     }
                   })
               // Dispatch/collective/etc fail the search.
               .Default([&](auto otherOp) { return false; })) {
        return nullptr;
      }
    }
  }
  return foundOp;
}

struct CmdExecuteOpPattern
    : public StreamConversionPattern<IREE::Stream::CmdExecuteOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::CmdExecuteOp executeOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = executeOp.getLoc();
    auto [device, queueAffinity] =
        lookupDeviceAndQueueAffinityFor(executeOp, rewriter);

    // If the command buffer only contains a single transfer command we may be
    // able to convert it to a queue operation instead. This will have
    // significantly less overhead than a command buffer especially if we are
    // not able to memoize it.
    if (auto *singleTransferOp = matchSingleTransferOp(executeOp)) {
      // Gather wait/signal fence, which are optional.
      Value waitFence =
          getOrCreateWaitFence(loc, adaptor.getAwaitTimepoint(), rewriter);
      Value signalFence = getOrCreateSignalFence(
          loc, device, executeOp.getResultTimepoint(), rewriter);

      // Replace the op with the queue operation.
      // Note that since we are matching an op nested within the region we have
      // to get the corresponding externally captured operand and lookup the
      // remapped value from the conversion state.
      //
      // Example:
      //   stream.cmd.execute ... with(%operand as %capture: !stream.resource)
      //     stream.cmd.fill ... %capture
      //  ->
      //   hal.device.queue.fill ... target(%operand : !hal.buffer)
      if (auto fillOp = dyn_cast<IREE::Stream::CmdFillOp>(*singleTransferOp)) {
        auto fillTargetBuffer = rewriter.getRemappedValue(
            executeOp.getClosureCapturedValue(fillOp.getTarget()));
        rewriter.create<IREE::HAL::DeviceQueueFillOp>(
            loc, device, queueAffinity, waitFence, signalFence,
            fillTargetBuffer, fillOp.getTargetOffset(),
            fillOp.getTargetLength(), fillOp.getValue(),
            /*flags=*/0);
      } else if (auto copyOp =
                     dyn_cast<IREE::Stream::CmdCopyOp>(*singleTransferOp)) {
        auto copySourceBuffer = rewriter.getRemappedValue(
            executeOp.getClosureCapturedValue(copyOp.getSource()));
        auto copyTargetBuffer = rewriter.getRemappedValue(
            executeOp.getClosureCapturedValue(copyOp.getTarget()));
        rewriter.create<IREE::HAL::DeviceQueueCopyOp>(
            loc, device, queueAffinity, waitFence, signalFence,
            copySourceBuffer, copyOp.getSourceOffset(), copyTargetBuffer,
            copyOp.getTargetOffset(), copyOp.getLength(),
            /*flags=*/0);
      }

      rewriter.replaceOp(executeOp, signalFence);
      return success();
    }

    // Until uniform buffers are implemented we can't reuse command buffers that
    // contain non-constant uniform values (i32, index, etc). We'll have a pass
    // that runs prior to conversion that creates new stream resources and
    // changes dispatches to use them for any dispatch we can - note that there
    // may still be some that slip through due to custom executables.
    const bool capturesDynamicUniformValues =
        clForceIndirectCommandBuffers
            ? false
            : regionCapturesDynamicUniformValues(executeOp);

    // Calculate the indirect buffer references used within the command buffer
    // by analyzing captured resources. This analysis will be used by subsequent
    // conversion to decide between embedding the direct buffer references or
    // indirect ones. We only do this if the execution region is reused.
    IndexSet indexSet(loc, rewriter);
    BindingTable bindingTable;
    if (!executeOp.getOnce() && !capturesDynamicUniformValues &&
        clIndirectCommandBuffers) {
      bindingTable = BindingTable(executeOp, adaptor.getResourceOperands(),
                                  adaptor.getResourceOperandSizes(), indexSet);
    }

    // If the execute op is one-shot or there's no indirect bindings then mark
    // the command buffer one-shot.
    IREE::HAL::CommandBufferModeBitfield modes =
        IREE::HAL::CommandBufferModeBitfield::None;
    if (!bindingTable.isSupported() || bindingTable.empty()) {
      modes = modes | IREE::HAL::CommandBufferModeBitfield::OneShot;
      if (!executeOp.getAwaitTimepoint()) {
        modes =
            modes | IREE::HAL::CommandBufferModeBitfield::AllowInlineExecution;
      }
      bindingTable = {};
    }

    // Cache the binding table values for use with the indirect execute.
    auto bindingTableValues = llvm::to_vector(bindingTable.getValues());

    // Derive the command buffer type based on the kind of operations present.
    // This can help the submission get routed to appropriate hardware queues
    // (like dedicated DMA controllers).
    auto commandCategories = deriveCommandCategories(executeOp.getBody());

    // Create, record, and finalize a command buffer at the current rewriter
    // insertion point. Returns the command buffer handle.
    auto recordCommandBuffer =
        [&](Value device, Value queueAffinity,
            ConversionPatternRewriter &rewriter) -> Value {
      // Create a new command buffer for recording.
      Value bindingTableCapacity =
          bindingTable.empty() ? Value{}
                               : rewriter.create<arith::ConstantIndexOp>(
                                     loc, bindingTable.size());
      Value commandBuffer = rewriter.create<IREE::HAL::CommandBufferCreateOp>(
          loc, rewriter.getType<IREE::HAL::CommandBufferType>(), device, modes,
          commandCategories, queueAffinity, bindingTableCapacity);
      mapping->mapCommandBuffer(executeOp, commandBuffer,
                                std::move(bindingTable));

      // Run through the execution region and serialize execution by inserting
      // barriers. Nested regions may elide barriers as needed.
      auto &bodyBlock = executeOp.getBody().front();
      insertSerializationBarriers(loc, bodyBlock, commandBuffer,
                                  OpBuilder::atBlockBegin(&bodyBlock));

      // Begin/end recording and inline the execution region between them.
      auto endOp = rewriter.create<IREE::HAL::CommandBufferFinalizeOp>(
          loc, commandBuffer);
      rewriter.inlineBlockBefore(&executeOp.getBody().front(), endOp,
                                 adaptor.getResourceOperands());

      // Return the command buffer handle.
      return commandBuffer;
    };

    // If reusable then we can memoize the command buffer by nesting it within
    // a memoization region and otherwise we inline the recording directly into
    // the original execution site.
    Value commandBuffer;
    if (!bitEnumContainsAll(modes,
                            IREE::HAL::CommandBufferModeBitfield::OneShot)) {
      auto memoizeOp = rewriter.create<IREE::HAL::DeviceMemoizeOp>(
          loc, rewriter.getType<IREE::HAL::CommandBufferType>(), device,
          queueAffinity);
      auto ip = rewriter.saveInsertionPoint();
      rewriter.setInsertionPointToStart(&memoizeOp.getBody().emplaceBlock());
      rewriter.create<IREE::HAL::ReturnOp>(
          loc, recordCommandBuffer(device, queueAffinity, rewriter));
      rewriter.restoreInsertionPoint(ip);
      commandBuffer = memoizeOp.getResult(0);
    } else {
      commandBuffer = recordCommandBuffer(device, queueAffinity, rewriter);
    }

    // Gather wait/signal fence, which are optional.
    Value waitFence =
        getOrCreateWaitFence(loc, adaptor.getAwaitTimepoint(), rewriter);
    Value signalFence = getOrCreateSignalFence(
        loc, device, executeOp.getResultTimepoint(), rewriter);

    // Queue execution.
    if (bindingTableValues.empty()) {
      rewriter.create<IREE::HAL::DeviceQueueExecuteOp>(
          loc, device, queueAffinity, waitFence, signalFence,
          ValueRange{commandBuffer});
    } else {
      rewriter.create<IREE::HAL::DeviceQueueExecuteIndirectOp>(
          loc, device, queueAffinity, waitFence, signalFence, commandBuffer,
          bindingTableValues);
    }

    rewriter.replaceOp(executeOp, signalFence);
    return success();
  }
};

struct CmdSerialOpPattern
    : public StreamConversionPattern<IREE::Stream::CmdSerialOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::CmdSerialOp serialOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto commandBufferMapping = mapping->lookupCommandBufferFor(serialOp);

    // Run through the execution region and serialize execution by inserting
    // barriers. Nested regions may elide barriers as needed.
    auto &bodyBlock = serialOp.getBody().front();
    insertSerializationBarriers(serialOp.getLoc(), bodyBlock,
                                commandBufferMapping.getHandle(),
                                OpBuilder::atBlockBegin(&bodyBlock));

    // Inline the serial execution region.
    rewriter.inlineBlockBefore(&serialOp.getBody().front(), serialOp);
    rewriter.eraseOp(serialOp);
    return success();
  }
};

struct CmdConcurrentOpPattern
    : public StreamConversionPattern<IREE::Stream::CmdConcurrentOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::CmdConcurrentOp concurrentOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Inline the concurrent execution region.
    // TODO(benvanik): split barriers (event set/wait) when nesting.
    rewriter.inlineBlockBefore(&concurrentOp.getBody().front(), concurrentOp);
    rewriter.eraseOp(concurrentOp);
    return success();
  }
};

struct TimepointImmediateOpPattern
    : public StreamConversionPattern<IREE::Stream::TimepointImmediateOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::TimepointImmediateOp immediateOp,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    rewriter.replaceOpWithNewOp<IREE::Util::NullOp>(
        immediateOp, rewriter.getType<IREE::HAL::FenceType>());
    return success();
  }
};

struct TimepointImportOpPattern
    : public StreamConversionPattern<IREE::Stream::TimepointImportOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::TimepointImportOp importOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Only handle imports from HAL semaphores _or_ fences.
    auto operands = adaptor.getOperands();
    if (operands.size() == 1 &&
        llvm::isa<IREE::HAL::FenceType>(operands[0].getType())) {
      rewriter.replaceOp(importOp, operands[0]);
      return success();
    } else {
      return rewriter.notifyMatchFailure(
          importOp, "only imports from HAL fences are supported");
    }
  }
};

struct TimepointExportOpPattern
    : public StreamConversionPattern<IREE::Stream::TimepointExportOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::TimepointExportOp exportOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Only handle exports into HAL fences.
    if (exportOp.getNumResults() != 1 ||
        !llvm::isa<IREE::HAL::FenceType>(exportOp.getResult(0).getType())) {
      return rewriter.notifyMatchFailure(
          exportOp, "only exports to HAL fences are supported");
    }
    rewriter.replaceOp(exportOp, adaptor.getAwaitTimepoint());
    return success();
  }
};

struct TimepointChainExternalOpPattern
    : public StreamConversionPattern<IREE::Stream::TimepointChainExternalOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::TimepointChainExternalOp exportOp,
                  OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Only handle exports into HAL fences.
    auto externalValues = exportOp.getExternalValues();
    if (externalValues.size() != 1 ||
        !llvm::isa<IREE::HAL::FenceType>(externalValues[0].getType())) {
      return rewriter.notifyMatchFailure(
          exportOp, "only exports to HAL fences are supported");
    }
    auto [device, queueAffinity] =
        lookupDeviceAndQueueAffinityFor(exportOp, rewriter);
    rewriter.replaceOpWithNewOp<IREE::HAL::DeviceQueueExecuteOp>(
        exportOp, device, queueAffinity,
        /*wait_fence=*/adaptor.getAwaitTimepoint(),
        /*signal_fence=*/externalValues[0], /*command_buffers=*/ValueRange{});
    return success();
  }
};

struct TimepointJoinOpPattern
    : public StreamConversionPattern<IREE::Stream::TimepointJoinOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::TimepointJoinOp joinOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    rewriter.replaceOpWithNewOp<IREE::HAL::FenceJoinOp>(
        joinOp, rewriter.getType<IREE::HAL::FenceType>(),
        adaptor.getAwaitTimepoints());
    return success();
  }
};

struct TimepointBarrierOpPattern
    : public StreamConversionPattern<IREE::Stream::TimepointBarrierOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::TimepointBarrierOp barrierOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Replace with a signaled fence.
    // NOTE: this assumes that if this op still exists the input resource is
    // already available. If it isn't then timepoint propagation should have
    // replaced the signal op with the producing timepoint.
    Value nullFence = rewriter.create<IREE::Util::NullOp>(
        barrierOp.getLoc(), rewriter.getType<IREE::HAL::FenceType>());
    rewriter.replaceOp(barrierOp, {adaptor.getResource(), nullFence});
    return success();
  }
};

struct TimepointAwaitOpPattern
    : public StreamConversionPattern<IREE::Stream::TimepointAwaitOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::TimepointAwaitOp awaitOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = awaitOp.getLoc();

    // Perform the blocking wait.
    Value timeoutMillis = rewriter.create<arith::ConstantIntOp>(loc, -1, 32);
    auto fenceOp = rewriter.create<IREE::HAL::FenceAwaitOp>(
        loc, rewriter.getI32Type(), timeoutMillis, adaptor.getAwaitTimepoint());
    rewriter.create<IREE::Util::StatusCheckOkOp>(loc, fenceOp.getStatus(),
                                                 "failed to wait on timepoint");

    // Pass along operands.
    rewriter.replaceOp(awaitOp, adaptor.getResourceOperands());
    return success();
  }
};

struct ChannelCreateOpPattern
    : public StreamConversionPattern<IREE::Stream::ChannelCreateOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::ChannelCreateOp createOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto [device, queueAffinity] =
        lookupDeviceAndQueueAffinityFor(createOp, rewriter);
    Value neg1I32;
    auto getDefault = [&]() {
      if (!neg1I32) {
        neg1I32 =
            rewriter.create<arith::ConstantIntOp>(createOp.getLoc(), -1, 32);
      }
      return neg1I32;
    };
    Value id = adaptor.getId();
    if (!id) {
      id = rewriter.create<IREE::Util::NullOp>(
          createOp.getLoc(), rewriter.getType<IREE::Util::BufferType>());
    }
    Value group =
        adaptor.getGroupAttr()
            ? rewriter
                  .create<IREE::Util::BufferConstantOp>(
                      createOp.getLoc(),
                      /*name=*/StringAttr{}, /*value=*/adaptor.getGroupAttr(),
                      /*alignment=*/IntegerAttr{}, /*mime_type=*/StringAttr{})
                  .getResult()
            : rewriter
                  .create<IREE::Util::NullOp>(
                      createOp.getLoc(),
                      rewriter.getType<IREE::Util::BufferType>())
                  .getResult();
    Value rank =
        adaptor.getRank()
            ? rewriter.create<arith::IndexCastOp>(
                  createOp.getLoc(), rewriter.getI32Type(), adaptor.getRank())
            : getDefault();
    Value count =
        adaptor.getRank()
            ? rewriter.create<arith::IndexCastOp>(
                  createOp.getLoc(), rewriter.getI32Type(), adaptor.getCount())
            : getDefault();
    rewriter.replaceOpWithNewOp<IREE::HAL::ChannelCreateOp>(
        createOp, rewriter.getType<IREE::HAL::ChannelType>(), device,
        queueAffinity, /*flags=*/rewriter.getI32IntegerAttr(0), id, group, rank,
        count);
    return success();
  }
};

struct ChannelSplitOpPattern
    : public StreamConversionPattern<IREE::Stream::ChannelSplitOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::ChannelSplitOp splitOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value color = rewriter.create<arith::IndexCastOp>(
        splitOp.getLoc(), rewriter.getI32Type(), adaptor.getColor());
    Value key = rewriter.create<arith::IndexCastOp>(
        splitOp.getLoc(), rewriter.getI32Type(), adaptor.getKey());
    rewriter.replaceOpWithNewOp<IREE::HAL::ChannelSplitOp>(
        splitOp, rewriter.getType<IREE::HAL::ChannelType>(),
        adaptor.getChannel(), color, key,
        /*flags=*/rewriter.getI32IntegerAttr(0));
    return success();
  }
};

struct ChannelRankOpPattern
    : public StreamConversionPattern<IREE::Stream::ChannelRankOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::ChannelRankOp rankOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto newOp = rewriter.create<IREE::HAL::ChannelRankAndCountOp>(
        rankOp.getLoc(), rewriter.getI32Type(), rewriter.getI32Type(),
        adaptor.getChannel());
    Value indexRank = rewriter.create<arith::IndexCastOp>(
        rankOp.getLoc(), rewriter.getIndexType(), newOp.getRank());
    rewriter.replaceOp(rankOp, indexRank);
    return success();
  }
};

struct ChannelCountOpPattern
    : public StreamConversionPattern<IREE::Stream::ChannelCountOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::ChannelCountOp countOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto newOp = rewriter.create<IREE::HAL::ChannelRankAndCountOp>(
        countOp.getLoc(), rewriter.getI32Type(), rewriter.getI32Type(),
        adaptor.getChannel());
    Value indexCount = rewriter.create<arith::IndexCastOp>(
        countOp.getLoc(), rewriter.getIndexType(), newOp.getCount());
    rewriter.replaceOp(countOp, indexCount);
    return success();
  }
};

struct ElideYieldOpPattern
    : public StreamConversionPattern<IREE::Stream::YieldOp> {
  using StreamConversionPattern::StreamConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Stream::YieldOp yieldOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    rewriter.eraseOp(yieldOp);
    return success();
  }
};

// Annoying we have to have this here, but there's no attribute converter
// equivalent we have access to so that we could do it in a generic way.
struct GlobalTimepointConversionPattern
    : public OpConversionPattern<IREE::Util::GlobalOp> {
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(IREE::Util::GlobalOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto initialValue = op.getInitialValue();
    if (!initialValue.has_value())
      return failure();
    if (!llvm::isa<IREE::Stream::TimepointAttr>(*initialValue))
      return failure();
    rewriter.modifyOpInPlace(op, [&]() { op.removeInitialValueAttr(); });
    return success();
  }
};

} // namespace

void populateStreamToHALPatterns(MLIRContext *context,
                                 ConversionTarget &conversionTarget,
                                 TypeConverter &typeConverter,
                                 RewritePatternSet &patterns) {
  conversionTarget.addIllegalDialect<IREE::Stream::StreamDialect>();

  typeConverter.addConversion(
      [=](IREE::Stream::ChannelType type, SmallVectorImpl<Type> &results) {
        // Collective channels are 1:1.
        results.push_back(IREE::HAL::ChannelType::get(context));
        return success();
      });

  typeConverter.addConversion(
      [=](IREE::Stream::FileType type, SmallVectorImpl<Type> &results) {
        results.push_back(IREE::HAL::FileType::get(context));
        return success();
      });

  typeConverter.addConversion(
      [=](IREE::Stream::ResourceType type, SmallVectorImpl<Type> &results) {
        // Resources are just buffers (no shape/encoding/etc).
        results.push_back(IREE::HAL::BufferType::get(context));
        return success();
      });

  typeConverter.addConversion(
      [=](IREE::Stream::TimepointType type, SmallVectorImpl<Type> &results) {
        results.push_back(IREE::HAL::FenceType::get(context));
        return success();
      });

  // Spooky action at a distance:
  patterns.insert<GlobalTimepointConversionPattern>(typeConverter, context);

  auto mapping = std::make_shared<StreamConversionMapping>();

  patterns.insert<ContextResolveOpPattern>(mapping, typeConverter, context);

  patterns.insert<ResourceAllocOpPattern, ResourceAllocaOpPattern,
                  ResourceDeallocaOpPattern, ResourceSizeOpPattern,
                  ResourceTryMapOpPattern, ResourceLoadOpPattern,
                  ResourceStoreOpPattern, ResourceSubviewOpPattern>(
      mapping, typeConverter, context);

  patterns.insert<FileConstantOpPattern, FileReadOpPattern, FileWriteOpPattern>(
      mapping, typeConverter, context);

  patterns.insert<TensorImportBufferOpPattern, TensorImportBufferViewOpPattern,
                  TensorExportBufferOpPattern, TensorExportBufferViewOpPattern,
                  TensorTraceOpPattern>(mapping, typeConverter, context);
  patterns
      .insert<CmdFlushOpPattern, CmdInvalidateOpPattern, CmdDiscardOpPattern,
              CmdFillOpPattern, CmdCopyOpPattern, CmdCollectiveOpPattern,
              CmdDispatchOpPattern, CmdFuncOpPattern, CmdCallOpPattern,
              CmdExecuteOpPattern, CmdSerialOpPattern, CmdConcurrentOpPattern>(
          mapping, typeConverter, context);
  patterns.insert<TimepointImmediateOpPattern, TimepointImportOpPattern,
                  TimepointExportOpPattern, TimepointChainExternalOpPattern,
                  TimepointJoinOpPattern, TimepointBarrierOpPattern,
                  TimepointAwaitOpPattern>(mapping, typeConverter, context);
  patterns.insert<ChannelCreateOpPattern, ChannelSplitOpPattern,
                  ChannelRankOpPattern, ChannelCountOpPattern>(
      mapping, typeConverter, context);
  patterns.insert<ElideYieldOpPattern>(mapping, typeConverter, context);
}

} // namespace mlir::iree_compiler
