// 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/Stream/IR/StreamOps.h"

#include "iree/compiler/Dialect/Encoding/IR/EncodingTypes.h"
#include "iree/compiler/Dialect/Util/IR/ClosureOpUtils.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
#include "llvm/ADT/BitVector.h"
#include "llvm/ADT/Repeated.h"
#include "llvm/ADT/SmallVectorExtras.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/Support/CommandLine.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/FunctionImplementation.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Transforms/RegionUtils.h"

namespace mlir::iree_compiler::IREE::Stream {

//===----------------------------------------------------------------------===//
// Op utilities used within the stream dialect
//===----------------------------------------------------------------------===//

// Verifies that a dispatch |op|'s |workload| matches that of the |exportOp|.
static LogicalResult
verifyDispatchWorkload(Operation *op, IREE::Stream::ExecutableExportOp exportOp,
                       ValueRange workload) {
  // If the target has a workgroup count computation function we can verify that
  // the workload here matches what is expected.
  if (!exportOp.getWorkgroupCount().empty()) {
    auto &workgroupCount = exportOp.getWorkgroupCount();
    auto explicitArgs = llvm::make_filter_range(
        workgroupCount.getArgumentTypes(), [](Type type) {
          return !type.hasTrait<
              mlir::OpTrait::IREE::Util::ImplicitlyCaptured>();
        });
    if (llvm::range_size(explicitArgs) != workload.size()) {
      return op->emitOpError()
             << "workload mismatch; entry point expects "
             << llvm::range_size(explicitArgs)
             << " arguments but dispatch provides " << workload.size();
    }
    for (auto [idx, expectedType, actualType] :
         llvm::enumerate(explicitArgs, workload.getTypes())) {
      if (expectedType != actualType) {
        return op->emitOpError()
               << "workload operand " << idx << " type mismatch; expected "
               << expectedType << " but passed " << actualType;
      }
    }
  }
  return success();
}

// Verifies the tied operand types are as the same as the result types.
static LogicalResult verifyTiedOperandEncodings(Operation *op,
                                                ArrayAttr operandEncodingsAttr,
                                                ArrayAttr resultEncodingsAttr) {
  auto tiedOp = dyn_cast<IREE::Util::TiedOpInterface>(op);
  if (!tiedOp) {
    return op->emitOpError()
           << "the op does not implement IREE::Util::TiedOpInterface";
  }

  ArrayRef<Attribute> operandEncodings = operandEncodingsAttr.getValue();
  unsigned tiedOperandBase = tiedOp.getTiedOperandsIndexAndLength().first;
  for (auto [idx, resultEncoding] :
       llvm::enumerate(resultEncodingsAttr.getValue())) {
    auto tiedOperand = tiedOp.getTiedResultOperandIndex(idx);
    if (!tiedOperand.has_value()) {
      continue;
    }
    auto operandIndex = tiedOperand.value() - tiedOperandBase;
    // The tied operand index comes from the tied_operands attribute and is not
    // otherwise range-checked for this op. An out-of-range value can be written
    // using the generic assembly form (e.g. `"stream.tensor.dispatch"(...)
    // {tied_operands = array<i64: 5>}`), which bypasses the custom-format
    // checks, so guard the encoding lookup below.
    if (operandIndex >= operandEncodings.size()) {
      return op->emitOpError()
             << "tied operand index " << operandIndex
             << " is out of range of the " << operandEncodings.size()
             << " operand encoding(s)";
    }
    if (operandEncodings[operandIndex] != resultEncoding) {
      return op->emitError()
             << "the " << operandIndex << "-th operandEncoding ("
             << operandEncodings[operandIndex]
             << ") does not match the resultEncoding (" << resultEncoding
             << ")";
    }
  }

  return success();
}

// Verifies that |dynamicDims| contains the appropriate number of dims for all
// the dynamic dimensions in |type|.
static LogicalResult verifyOpDynamicDims(Operation *op, TypeRange types,
                                         ValueRange dynamicDims) {
  unsigned requiredCount = 0;
  for (auto type : types) {
    if (auto shapedType = dyn_cast_if_present<ShapedType>(type)) {
      requiredCount += shapedType.getNumDynamicDims();
    }
  }
  if (dynamicDims.size() != requiredCount) {
    return op->emitOpError()
           << "type set has " << requiredCount
           << " dynamic dimensions but only " << dynamicDims.size()
           << " dimension values are attached";
  }
  return success();
}

// Verifies that |dynamicDims| contains the appropriate number of dims for all
// the dynamic dimensions in |type|.
static LogicalResult verifyOpDynamicDimsRange(Operation *op,
                                              ArrayAttr typesAttr,
                                              ValueRange dynamicDims) {
  unsigned requiredCount = 0;
  for (auto attr : typesAttr) {
    if (auto typeAttr = dyn_cast_if_present<TypeAttr>(attr)) {
      if (auto shapedType = dyn_cast<ShapedType>(typeAttr.getValue())) {
        requiredCount += shapedType.getNumDynamicDims();
      }
    }
  }
  if (dynamicDims.size() != requiredCount) {
    return op->emitOpError()
           << "type set has " << requiredCount
           << " dynamic dimensions but only " << dynamicDims.size()
           << " dimension values are attached";
  }
  return success();
}

// Verifies that |sizes| contains the appropriate number of sizes for all of the
// sized types in |values|.
static LogicalResult verifyOpValueSizes(Operation *op, ValueRange values,
                                        ValueRange sizes) {
  unsigned requiredCount = 0;
  for (auto value : values) {
    if (isa<IREE::Util::SizeAwareTypeInterface>(value.getType())) {
      ++requiredCount;
    }
  }
  if (sizes.size() != requiredCount) {
    return op->emitOpError() << "value set has " << requiredCount
                             << " dynamic dimensions but only " << sizes.size()
                             << " dimension values are attached";
  }
  return success();
}

// Verifies that all !stream.resources used within |region| are captured by
// the entry arguments to the region.
static LogicalResult verifyAllResourcesCaptured(Region &region) {
  SetVector<Value> availableResources;
  for (auto arg : region.front().getArguments()) {
    availableResources.insert(arg);
  }
  for (auto &op : region.front()) {
    for (auto result : op.getResults()) {
      availableResources.insert(result);
    }
    for (auto operand : op.getOperands()) {
      if (!operand) {
        continue;
      }
      if (!isa<IREE::Stream::ResourceType>(operand.getType())) {
        continue;
      }
      if (!availableResources.contains(operand)) {
        return op.emitOpError() << "used resource not listed in explicit "
                                   "captures (or produced internally)";
      }
    }
  }
  return success();
}

// Verifies that escaping !stream.resources have the sizes when they are
// yielded match the sizes declared on the parent op. This information is
// redundant but keeps analysis local and agnostic to the parent op structure
// which is useful for when we outline things.
static LogicalResult verifyEscapingResources(Region &region,
                                             ResultRange results,
                                             ValueRange resultSizes) {
  // Ensure yielded resources match the signature.
  for (auto yieldOp : region.getOps<IREE::Stream::YieldOp>()) {
    if (results.size() != yieldOp.getResourceOperands().size()) {
      return yieldOp.emitOpError()
             << "yield result count mismatch with parent op";
    }
    for (auto [outerValue, innerValue] :
         llvm::zip_equal(results, yieldOp.getResourceOperands())) {
      if (outerValue.getType() != innerValue.getType()) {
        return yieldOp.emitOpError()
               << "result type mismatch: expected " << outerValue.getType()
               << " but got " << innerValue.getType();
      }
    }
    for (auto [outerSize, innerSize] :
         llvm::zip_equal(resultSizes, yieldOp.getResourceOperandSizes())) {
      if (outerSize != innerSize) {
        return yieldOp.emitOpError() << "result size mismatch";
      }
    }
  }
  return success();
}

static void eraseStreamRegionResults(Region &region,
                                     ArrayRef<unsigned> excludedResultIndices) {
  for (auto &block : region.getBlocks()) {
    auto yieldOp = dyn_cast<IREE::Stream::YieldOp>(block.getTerminator());
    if (!yieldOp) {
      continue;
    }
    // HACK: there's no good way of updating the operand and size together today
    // - we should add a helper to the ClosureYieldOpInterface that checks for
    // size/shape aware traits and does this automatically.
    for (auto i : llvm::reverse(excludedResultIndices)) {
      unsigned resourceIndex = i;
      unsigned resourceSizeIndex =
          yieldOp.getResourceOperandsMutable().size() + i;
      yieldOp->eraseOperand(resourceSizeIndex);
      yieldOp->eraseOperand(resourceIndex);
    }
  }
}

// Checks if a timepoint transitively covers another timepoint through the
// await dependency chain. This performs a local dominance-based check without
// requiring global analysis.
static bool timepointCoversTimepoint(Value coveringTimepoint,
                                     Value coveredTimepoint) {
  if (coveringTimepoint == coveredTimepoint) {
    return true;
  }

  // Check if coveringTimepoint is defined by an op that awaits
  // coveredTimepoint.
  auto coveringOp = dyn_cast_or_null<IREE::Stream::TimelineOpInterface>(
      coveringTimepoint.getDefiningOp());
  if (!coveringOp) {
    return false;
  }

  // Get await timepoints and check if any match or transitively cover.
  SmallVector<Value> awaitTimepoints = coveringOp.getAwaitTimepoints();
  for (Value awaitTp : awaitTimepoints) {
    if (timepointCoversTimepoint(awaitTp, coveredTimepoint)) {
      return true;
    }
  }

  // Check stream.timepoint.join operations.
  if (auto joinOp = dyn_cast<IREE::Stream::TimepointJoinOp>(
          coveringTimepoint.getDefiningOp())) {
    for (Value joinedTp : joinOp.getAwaitTimepoints()) {
      if (timepointCoversTimepoint(joinedTp, coveredTimepoint)) {
        return true;
      }
    }
  }

  return false;
}

// Verifies that a timeline op's await clause covers all the resources it uses.
// For each resource operand, if it was produced by a timeline op with a
// timepoint, that timepoint should be in the await list or transitively
// covered.
static LogicalResult
verifyTimelineResourceDominance(IREE::Stream::TimelineOpInterface timelineOp,
                                ValueRange resourceOperands) {
  SmallVector<Value> awaitTimepoints = timelineOp.getAwaitTimepoints();
  for (Value resource : resourceOperands) {
    // Skip resources not defined by timeline ops.
    auto producerTimeline =
        dyn_cast_if_present<IREE::Stream::TimelineOpInterface>(
            resource.getDefiningOp());
    if (!producerTimeline) {
      continue;
    }
    Value producerTimepoint = producerTimeline.getResultTimepoint();
    if (!producerTimepoint) {
      continue;
    }

    // Check if any of the await timepoints cover the producer's timepoint.
    bool isCovered = false;
    for (Value awaitTp : awaitTimepoints) {
      if (timepointCoversTimepoint(awaitTp, producerTimepoint)) {
        isCovered = true;
        break;
      }
    }

    if (!isCovered) {
      return timelineOp->emitOpError()
             << "uses resource produced by timeline op but does not await its "
                "completion timepoint (resource produced at "
             << resource.getLoc() << ", timepoint not covered by await clause)";
    }
  }

  return success();
}

// Computes the value access bits starting from |rootValue|.
// Traverses the IR graph along tied ops but does not handle branches.
static IREE::Util::ValueAccess computeValueAccess(Value rootValue) {
  IREE::Util::ValueAccess access;
  DenseSet<Value> processedValues;
  SmallVector<Value> worklist;
  auto enqueueValue = [&](Value value) {
    if (processedValues.contains(value)) {
      return;
    }
    processedValues.insert(value);
    worklist.push_back(value);
  };
  enqueueValue(rootValue);
  while (!worklist.empty()) {
    Value value = worklist.back();
    worklist.pop_back();

    // Walk up the definition chain.
    if (auto definingOp = value.getDefiningOp()) {
      // Value is produced within the region and thus written.
      access.isWrite = true;
      if (auto tiedOp = dyn_cast<IREE::Util::TiedOpInterface>(definingOp)) {
        access.isRead = true;
        auto operand = tiedOp.getTiedResultOperand(value);
        if (operand) {
          // Value is tied back to another value; continue analyzing past it.
          enqueueValue(operand);
        } else {
          // Value contents are fully produced by this op.
          access.isDiscard = true;
        }
      } else if (isa<IREE::Stream::SubviewEffectOpInterface>(definingOp)) {
        // TODO(benvanik): actually query; for now assume *.
        access.isRead = true;
        access.isWrite = true;
      } else {
        // Value contents are fully produced by this op.
        access.isDiscard = true;
      }
    }

    // Walk down the use chain.
    for (auto user : value.getUsers()) {
      // Used by an op.
      access.isRead = true;
      if (auto tiedOp = dyn_cast<IREE::Util::TiedOpInterface>(user)) {
        auto tiedIndices = tiedOp.getTiedResultOperandIndices();
        for (int64_t tiedIndex : tiedIndices) {
          if (tiedIndex == IREE::Util::TiedOpInterface::kUntiedIndex) {
            continue;
          }
          auto operand = user->getOperand(tiedIndex);
          if (operand == value) {
            // Tied operand.
            access.isRead = true;
            access.isWrite = true;
            enqueueValue(operand);
          }
        }
      } else if (isa<IREE::Stream::SubviewEffectOpInterface>(user)) {
        // TODO(benvanik): actually query; for now assume *.
        access.isRead = true;
        access.isWrite = true;
      }
    }
  }
  return access;
}

//===----------------------------------------------------------------------===//
// custom<DispatchEntryPoints>($entry_points)
//===----------------------------------------------------------------------===//

static ParseResult parseDispatchEntryPoints(OpAsmParser &parser,
                                            ArrayAttr &entryPointAttrsArray) {
  SmallVector<Attribute> entryPointAttrs;
  if (succeeded(parser.parseOptionalLBrace())) {
    do {
      SymbolRefAttr entryPointAttr;
      if (failed(parser.parseAttribute(entryPointAttr))) {
        return failure();
      }
      entryPointAttrs.push_back(entryPointAttr);
    } while (succeeded(parser.parseOptionalComma()));
    if (failed(parser.parseRBrace())) {
      return failure();
    }
  } else {
    SymbolRefAttr entryPointAttr;
    if (failed(parser.parseAttribute(entryPointAttr))) {
      return failure();
    }
    entryPointAttrs.push_back(entryPointAttr);
  }
  entryPointAttrsArray = parser.getBuilder().getArrayAttr(entryPointAttrs);
  return success();
}

static void printDispatchEntryPoints(OpAsmPrinter &p, Operation *op,
                                     ArrayAttr entryPointAttrs) {
  if (entryPointAttrs.size() == 1) {
    p.printAttribute(entryPointAttrs.getValue().front());
  } else {
    p << '{';
    llvm::interleaveComma(entryPointAttrs, p.getStream(),
                          [&](Attribute attr) { p.printAttribute(attr); });
    p << '}';
  }
}

//===----------------------------------------------------------------------===//
// custom<EncodedResourceOperands>(
//     $resources, type($resources), $resource_sizes,
//     $resource_encodings, $resource_encoding_dims)
//===----------------------------------------------------------------------===//

static ParseResult parseEncodedResourceOperands(
    OpAsmParser &parser,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resources,
    SmallVectorImpl<Type> &resourceTypes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceSizes,
    ArrayAttr &resourceEncodings,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceEncodingDims) {
  SmallVector<Attribute> resourceEncodingAttrs;
  do {
    resources.emplace_back();
    TypeAttr resourceEncoding;
    if (failed(parser.parseOperand(resources.back())) ||
        failed(parser.parseColon()) ||
        failed(parser.parseAttribute(resourceEncoding))) {
      return failure();
    }
    resourceEncodingAttrs.push_back(resourceEncoding);
    if (int64_t dynamicDimCount =
            cast<ShapedType>(resourceEncoding.getValue()).getNumDynamicDims()) {
      if (failed(parser.parseOperandList(resourceEncodingDims, dynamicDimCount,
                                         AsmParser::Delimiter::Braces))) {
        return failure();
      }
    }
    resourceTypes.emplace_back();
    resourceSizes.emplace_back();
    if (failed(parser.parseKeyword("in")) ||
        failed(parseSizeAwareType(parser, resourceTypes.back(),
                                  resourceSizes.back()))) {
      return failure();
    }
  } while (succeeded(parser.parseOptionalComma()));
  resourceEncodings = parser.getBuilder().getArrayAttr(resourceEncodingAttrs);
  return success();
}

static void printEncodedResourceOperands(OpAsmPrinter &p, Operation *op,
                                         ValueRange resources,
                                         TypeRange resourceTypes,
                                         ValueRange resourceSizes,
                                         ArrayAttr resourceEncodings,
                                         ValueRange resourceEncodingDims) {
  p.increaseIndent();
  p.printNewline();
  llvm::interleave(
      llvm::zip_equal(resources, resourceTypes, resourceSizes,
                      resourceEncodings.getAsValueRange<TypeAttr>()),
      [&](auto it) {
        auto [resource, resourceType, resourceSize, resourceEncoding] = it;
        p << resource;
        p << " : ";
        p << resourceEncoding;
        if (int64_t dynamicDimCount =
                cast<ShapedType>(resourceEncoding).getNumDynamicDims()) {
          p << "{";
          llvm::interleaveComma(
              resourceEncodingDims.take_front(dynamicDimCount), p);
          resourceEncodingDims =
              resourceEncodingDims.drop_front(dynamicDimCount);
          p << "}";
        }
        p << " in ";
        p << resourceType;
        p << "{";
        p << resourceSize;
        p << "}";
      },
      [&]() {
        p << ",";
        p.printNewline();
      });
  p.decreaseIndent();
  p.printNewline();
}

//===----------------------------------------------------------------------===//
// custom<EncodedShapedTypeList>
//===----------------------------------------------------------------------===//
// encoding{%dim0, %dim1} in type{%size0}, type, type{%size1}

static ParseResult
parseShapedType(OpAsmParser &parser, Type &type,
                SmallVectorImpl<OpAsmParser::UnresolvedOperand> &dims) {
  if (failed(parser.parseType(type))) {
    return failure();
  }
  if (auto shapedType = dyn_cast<ShapedType>(type)) {
    if (!shapedType.hasStaticShape()) {
      SmallVector<OpAsmParser::UnresolvedOperand> dynamicDims;
      if (failed(parser.parseLBrace()) ||
          failed(parser.parseOperandList(dynamicDims,
                                         shapedType.getNumDynamicDims(),
                                         OpAsmParser::Delimiter::None)) ||
          failed(parser.parseRBrace())) {
        return failure();
      }
      dims.append(dynamicDims);
    }
  } else if (isa<IREE::Util::SizeAwareTypeInterface>(type)) {
    OpAsmParser::UnresolvedOperand size;
    if (failed(parser.parseLBrace()) || failed(parser.parseOperand(size)) ||
        failed(parser.parseRBrace())) {
      return failure();
    }
    dims.push_back(size);
  }
  return success();
}

static void printSizedType(OpAsmPrinter &p, Operation *op, Type type,
                           Value size) {
  p.printType(type);
  p << "{";
  p.printOperand(size);
  p << "}";
}

static OperandRange printShapedType(OpAsmPrinter &p, Operation *op, Type type,
                                    OperandRange dims) {
  p.printType(type);
  if (auto shapedType = dyn_cast<ShapedType>(type)) {
    if (!shapedType.hasStaticShape()) {
      if (dims.empty()) {
        p << "{<<INVALID>>}";
        return dims;
      }
      p << "{";
      llvm::interleaveComma(dims.take_front(shapedType.getNumDynamicDims()), p,
                            [&](Value value) { p.printOperand(value); });
      p << "}";
      dims = dims.drop_front(shapedType.getNumDynamicDims());
    }
  } else if (isa<IREE::Util::SizeAwareTypeInterface>(type)) {
    p << "{";
    p.printOperand(dims.front());
    p << "}";
    dims = dims.drop_front(1);
  }
  return dims;
}

static ParseResult parseEncodedShapedTypeList(
    OpAsmParser &parser, SmallVectorImpl<Type> &types,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &sizes,
    SmallVectorImpl<Type> &encodings,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &encodingDims) {
  do {
    Type type0;
    SmallVector<OpAsmParser::UnresolvedOperand> dims0;
    if (failed(parseShapedType(parser, type0, dims0))) {
      return failure();
    }
    if (succeeded(parser.parseOptionalKeyword("in"))) {
      Type type1;
      SmallVector<OpAsmParser::UnresolvedOperand> dims1;
      if (failed(parseShapedType(parser, type1, dims1))) {
        return failure();
      }
      types.push_back(type1);
      sizes.append(dims1);
      encodings.push_back(type0);
      encodingDims.append(dims0);
    } else {
      types.push_back(type0);
      sizes.append(dims0);
      encodings.push_back(IREE::Util::UnusedType::get(parser.getContext()));
    }
  } while (succeeded(parser.parseOptionalComma()));
  return success();
}

static void printEncodedShapedTypeList(OpAsmPrinter &p, Operation *op,
                                       TypeRange types, OperandRange sizes,
                                       ArrayAttr encodings,
                                       OperandRange encodingDims) {
  llvm::interleaveComma(
      llvm::zip_equal(types, encodings.getAsValueRange<TypeAttr>()), p,
      [&](std::tuple<Type, Type> it) {
        auto [type, encoding] = it;
        if (!isa<IREE::Util::UnusedType>(encoding)) {
          encodingDims = printShapedType(p, op, encoding, encodingDims);
          p << " in ";
        }
        sizes = printShapedType(p, op, type, sizes);
      });
}

//===----------------------------------------------------------------------===//
// custom<EncodedShapedResultList>
//===----------------------------------------------------------------------===//
// encoding{%dim0, %dim1} in type{%dim2}, type{%size}, %operand4
//
// Supported result formats:
//   type{%size}
//   %operand as type{%size}
//   encoding{%dim0, %dim1} in %operand4
//   encoding{%dim0, %dim1} in %operand4 as type{%size}

static ParseResult parseEncodedShapedResultList(
    OpAsmParser &parser, ArrayRef<OpAsmParser::UnresolvedOperand> operands,
    TypeRange operandTypes,
    ArrayRef<OpAsmParser::UnresolvedOperand> operandSizes,
    SmallVectorImpl<Type> &resultTypes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resultSizes,
    SmallVectorImpl<Type> &resultEncodingTypes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resultEncodingDims,
    ArrayAttr &tiedOperands) {
  SmallVector<int64_t> tiedOperandIndices;
  do {
    Type type0;
    SmallVector<OpAsmParser::UnresolvedOperand> dims0;
    auto typeResult = parser.parseOptionalType(type0);
    if (typeResult.has_value() && succeeded(typeResult.value())) {
      if (auto shapedType = dyn_cast<ShapedType>(type0)) {
        if (!shapedType.hasStaticShape()) {
          if (failed(parser.parseLBrace()) ||
              failed(parser.parseOperandList(dims0,
                                             shapedType.getNumDynamicDims(),
                                             OpAsmParser::Delimiter::None)) ||
              failed(parser.parseRBrace())) {
            return failure();
          }
        }
      } else if (auto sizedType =
                     dyn_cast<IREE::Util::SizeAwareTypeInterface>(type0)) {
        OpAsmParser::UnresolvedOperand size;
        if (failed(parser.parseLBrace()) || failed(parser.parseOperand(size)) ||
            failed(parser.parseRBrace())) {
          return failure();
        }
        dims0.push_back(size);
      }
    }

    // Type only:
    if (failed(parser.parseOptionalKeyword("in"))) {
      resultTypes.push_back(type0);
      resultSizes.append(dims0);
      resultEncodingTypes.push_back(
          IREE::Util::UnusedType::get(parser.getContext()));
      tiedOperandIndices.push_back(IREE::Util::TiedOpInterface::kUntiedIndex);
      continue;
    }

    // Check for optional tied result reference.
    OpAsmParser::UnresolvedOperand tiedResult;
    auto res = parser.parseOptionalOperand(tiedResult);
    Type resultType;
    int64_t tiedOperandIndex = IREE::Util::TiedOpInterface::kUntiedIndex;
    if (res.has_value() && succeeded(res.value())) {
      tiedOperandIndex = findTiedOperand(tiedResult, operands);
      if (tiedOperandIndex == IREE::Util::TiedOpInterface::kUntiedIndex) {
        return parser.emitError(tiedResult.location,
                                "tied operand not found for result reference ")
               << tiedResult.name;
      }
      if (succeeded(parser.parseOptionalKeyword("as"))) {
        // Type _may_ differ from the operand.
        if (failed(parser.parseType(resultType))) {
          return failure();
        }
      } else {
        // Use the operands type.
        resultType = operandTypes[tiedOperandIndex];
      }
    } else if (failed(parser.parseType(resultType))) {
      return failure();
    }

    // Parse optional type dimensions (usually resource size here).
    if (auto sizedType =
            dyn_cast<IREE::Util::SizeAwareTypeInterface>(resultType)) {
      OpAsmParser::UnresolvedOperand size;
      if (failed(parser.parseLBrace()) || failed(parser.parseOperand(size)) ||
          failed(parser.parseRBrace())) {
        return failure();
      }
      resultSizes.push_back(size);
    }

    resultTypes.push_back(resultType);
    resultEncodingTypes.push_back(type0);
    resultEncodingDims.append(dims0);
    tiedOperandIndices.push_back(tiedOperandIndex);
  } while (succeeded(parser.parseOptionalComma()));
  if (!tiedOperandIndices.empty()) {
    tiedOperands = parser.getBuilder().getIndexArrayAttr(tiedOperandIndices);
  }
  return success();
}

static void printEncodedShapedResultList(
    OpAsmPrinter &p, Operation *op, ValueRange operands, TypeRange operandTypes,
    OperandRange operandSizes, TypeRange resultTypes, OperandRange resultSizes,
    ArrayAttr resultEncodings, OperandRange resultEncodingDims,
    ArrayAttr tiedOperands) {
  auto tiedOp = dyn_cast<IREE::Util::TiedOpInterface>(op);
  for (unsigned i = 0; i < resultTypes.size(); ++i) {
    auto resultEncodingType =
        cast<TypeAttr>(resultEncodings.getValue()[i]).getValue();
    if (!isa<IREE::Util::UnusedType>(resultEncodingType)) {
      p.printType(resultEncodingType);
      if (auto shapedType = dyn_cast<ShapedType>(resultEncodingType)) {
        if (!shapedType.hasStaticShape()) {
          if (resultEncodingDims.empty()) {
            p << "{<<INVALID>>}";
            return;
          }
          p << "{";
          llvm::interleaveComma(
              resultEncodingDims.take_front(shapedType.getNumDynamicDims()), p,
              [&](Value value) { p.printOperand(value); });
          p << "}";
          resultEncodingDims =
              resultEncodingDims.drop_front(shapedType.getNumDynamicDims());
        }
      } else if (auto sizedType = dyn_cast<IREE::Util::SizeAwareTypeInterface>(
                     resultEncodingType)) {
        p << "{";
        p.printOperand(resultEncodingDims.front());
        p << "}";
        resultEncodingDims = resultEncodingDims.drop_front(1);
      }
      p << " in ";
    }
    auto resultType = resultTypes[i];
    auto tiedOperandIndex =
        tiedOp ? tiedOp.getTiedResultOperandIndex(i) : std::nullopt;
    bool printType = true;
    if (tiedOperandIndex.has_value()) {
      auto tiedOperand = op->getOperand(tiedOperandIndex.value());
      p.printOperand(tiedOperand);
      if (tiedOperand.getType() != resultType) {
        p << " as ";
      } else {
        // Type elided as it matches the operand.
        printType = false;
      }
    }
    if (printType) {
      p.printType(resultType);
    }
    if (auto sizedType =
            dyn_cast<IREE::Util::SizeAwareTypeInterface>(resultType)) {
      p << "{";
      p.printOperand(resultSizes.front());
      p << "}";
      resultSizes = resultSizes.drop_front(1);
    }
    if (i < resultTypes.size() - 1) {
      p << ", ";
    }
  }
}

//===----------------------------------------------------------------------===//
// custom<EncodedShapedFunctionType>
//===----------------------------------------------------------------------===//
// (type, encoding{%dim0, %dim1} in type{%size}, type) ->
//     (encoding{%dim} in type{%size}, %operand4)

static ParseResult parseEncodedShapedFunctionType(
    OpAsmParser &parser, ArrayRef<OpAsmParser::UnresolvedOperand> operands,
    SmallVectorImpl<Type> &operandTypes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operandSizes,
    ArrayAttr &operandEncodings,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operandEncodingDims,
    SmallVectorImpl<Type> &resultTypes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resultSizes,
    ArrayAttr &resultEncodings,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resultEncodingDims,
    ArrayAttr &tiedOperands) {
  SmallVector<Type> operandEncodingTypes;
  SmallVector<Type> resultEncodingTypes;
  if (failed(parser.parseLParen())) {
    return failure();
  }
  if (failed(parser.parseOptionalRParen())) {
    if (failed(parseEncodedShapedTypeList(parser, operandTypes, operandSizes,
                                          operandEncodingTypes,
                                          operandEncodingDims)) ||
        failed(parser.parseRParen())) {
      return failure();
    }
  }
  if (failed(parser.parseArrow())) {
    return failure();
  }
  if (succeeded(parser.parseOptionalLParen())) {
    if (succeeded(parser.parseOptionalRParen())) {
      // Empty list/no results `()`.
    } else {
      // One or more result types.
      if (failed(parseEncodedShapedResultList(
              parser, operands, operandTypes, operandSizes, resultTypes,
              resultSizes, resultEncodingTypes, resultEncodingDims,
              tiedOperands)) ||
          failed(parser.parseRParen())) {
        return failure();
      }
    }
  } else {
    // Single result with omitted `()`.
    if (failed(parseEncodedShapedResultList(
            parser, operands, operandTypes, operandSizes, resultTypes,
            resultSizes, resultEncodingTypes, resultEncodingDims,
            tiedOperands))) {
      return failure();
    }
  }
  operandEncodings = ArrayAttr::get(
      parser.getContext(),
      llvm::map_to_vector(operandEncodingTypes, [](Type type) -> Attribute {
        return type ? TypeAttr::get(type) : Attribute{};
      }));
  resultEncodings = ArrayAttr::get(
      parser.getContext(),
      llvm::map_to_vector(resultEncodingTypes, [](Type type) -> Attribute {
        return TypeAttr::get(type);
      }));
  return success();
}

static void printEncodedShapedFunctionType(
    OpAsmPrinter &p, Operation *op, ValueRange operands, TypeRange operandTypes,
    OperandRange operandSizes, ArrayAttr operandEncodings,
    OperandRange operandEncodingDims, TypeRange resultTypes,
    OperandRange resultSizes, ArrayAttr resultEncodings,
    OperandRange resultEncodingDims, ArrayAttr tiedOperands) {
  p << "(";
  printEncodedShapedTypeList(p, op, operandTypes, operandSizes,
                             operandEncodings, operandEncodingDims);
  p << ") -> ";
  if (resultTypes.size() != 1) {
    p << "(";
  }
  printEncodedShapedResultList(p, op, operands, operandTypes, operandSizes,
                               resultTypes, resultSizes, resultEncodings,
                               resultEncodingDims, tiedOperands);
  if (resultTypes.size() != 1) {
    p << ")";
  }
}

//===----------------------------------------------------------------------===//
// Shared parameter reference parsing/printing: %scope :: %key  or  %key
// (Overloads of the StringAttr version in StreamTypes.h — these handle
// operand-based scope/key for the parameter ops that now use !util.buffer.)
//===----------------------------------------------------------------------===//

// Parses an operand, then checks for `::`. If present, the first operand was
// the scope and the second is the key. Otherwise the first operand is the key
// and scope is unset.
static ParseResult
parseParameterReference(OpAsmParser &parser,
                        std::optional<OpAsmParser::UnresolvedOperand> &scope,
                        OpAsmParser::UnresolvedOperand &key) {
  OpAsmParser::UnresolvedOperand firstOperand;
  if (failed(parser.parseOperand(firstOperand))) {
    return failure();
  }
  if (succeeded(parser.parseOptionalColon())) {
    scope = firstOperand;
    if (failed(parser.parseColon()) || failed(parser.parseOperand(key))) {
      return failure();
    }
  } else {
    key = firstOperand;
  }
  return success();
}

static void printParameterReference(OpAsmPrinter &p, Operation *op, Value scope,
                                    Value key) {
  if (scope) {
    p.printOperand(scope);
    p << "::";
  }
  p.printOperand(key);
}

//===----------------------------------------------------------------------===//
// custom<ParameterLoadOperations>(
//     $source_scope, $source_keys, $source_offsets,
//     type($results), $result_sizes)
//===----------------------------------------------------------------------===//

static ParseResult parseParameterLoadOperations(
    OpAsmParser &parser,
    std::optional<OpAsmParser::UnresolvedOperand> &sourceScope,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &sourceKeys,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &sourceOffsets,
    SmallVectorImpl<Type> &resultTypes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resultSizes) {
  std::optional<OpAsmParser::UnresolvedOperand> firstScope;
  bool firstRow = true;
  do {
    std::optional<OpAsmParser::UnresolvedOperand> rowScope;
    OpAsmParser::UnresolvedOperand sourceKey;
    OpAsmParser::UnresolvedOperand sourceOffset;
    Type resultType;
    OpAsmParser::UnresolvedOperand resultSize;
    if (failed(parseParameterReference(parser, rowScope, sourceKey)) ||
        failed(parser.parseLSquare()) ||
        failed(parser.parseOperand(sourceOffset)) ||
        failed(parser.parseRSquare()) ||
        failed(parser.parseColonType(resultType)) ||
        failed(parser.parseLBrace()) ||
        failed(parser.parseOperand(resultSize)) ||
        failed(parser.parseRBrace())) {
      return failure();
    }
    if (firstRow) {
      firstScope = rowScope;
      sourceScope = rowScope;
      firstRow = false;
    } else if (rowScope.has_value() != firstScope.has_value() ||
               (rowScope.has_value() && rowScope->name != firstScope->name)) {
      return parser.emitError(parser.getCurrentLocation(),
                              "each operation must use the same scope");
    }
    sourceKeys.push_back(sourceKey);
    sourceOffsets.push_back(sourceOffset);
    resultTypes.push_back(resultType);
    resultSizes.push_back(resultSize);
  } while (succeeded(parser.parseOptionalComma()));
  return success();
}

static void
printParameterLoadOperations(OpAsmPrinter &p, Operation *op, Value sourceScope,
                             ValueRange sourceKeys, ValueRange sourceOffsets,
                             TypeRange resultTypes, ValueRange resultSizes) {
  p.increaseIndent();
  p.printNewline();
  llvm::interleave(
      llvm::zip_equal(sourceKeys, sourceOffsets, resultTypes, resultSizes),
      [&](std::tuple<Value, Value, Type, Value> it) {
        auto [sourceKey, sourceOffset, resultType, resultSize] = it;
        printParameterReference(p, op, sourceScope, sourceKey);
        p << "[";
        p.printOperand(sourceOffset);
        p << "] : ";
        p.printType(resultType);
        p << "{";
        p.printOperand(resultSize);
        p << "}";
      },
      [&]() {
        p << ',';
        p.printNewline();
      });
  p.decreaseIndent();
  p.printNewline();
}

//===----------------------------------------------------------------------===//
// custom<AsyncParameterGatherOperations>(
//     $source_scope, $source_keys, $source_offsets,
//     $target, type($target), $target_size, $target_offsets, $target_ends,
//     $target_lengths)
//===----------------------------------------------------------------------===//

static ParseResult parseAsyncParameterGatherOperations(
    OpAsmParser &parser,
    std::optional<OpAsmParser::UnresolvedOperand> &sourceScope,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &sourceKeys,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &sourceOffsets,
    OpAsmParser::UnresolvedOperand &target, Type &targetType,
    OpAsmParser::UnresolvedOperand &targetSize,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &targetOffsets,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &targetEnds,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &targetLengths) {
  std::optional<OpAsmParser::UnresolvedOperand> firstScope;
  bool firstRow = true;
  do {
    std::optional<OpAsmParser::UnresolvedOperand> rowScope;
    OpAsmParser::UnresolvedOperand sourceKey;
    OpAsmParser::UnresolvedOperand sourceOffset;
    OpAsmParser::UnresolvedOperand targetOffset;
    OpAsmParser::UnresolvedOperand targetEnd;
    OpAsmParser::UnresolvedOperand targetLength;
    OpAsmParser::UnresolvedOperand rowTarget;
    Type rowTargetType;
    OpAsmParser::UnresolvedOperand rowTargetSize;
    if (failed(parseParameterReference(parser, rowScope, sourceKey)) ||
        failed(parser.parseLSquare()) ||
        failed(parser.parseOperand(sourceOffset)) ||
        failed(parser.parseRSquare()) || failed(parser.parseArrow()) ||
        failed(parser.parseOperand(rowTarget)) ||
        failed(parser.parseLSquare()) ||
        failed(parser.parseOperand(targetOffset)) ||
        failed(parser.parseKeyword("to")) ||
        failed(parser.parseOperand(targetEnd)) ||
        failed(parser.parseKeyword("for")) ||
        failed(parser.parseOperand(targetLength)) ||
        failed(parser.parseRSquare()) ||
        failed(parser.parseColonType(rowTargetType)) ||
        failed(parser.parseLBrace()) ||
        failed(parser.parseOperand(rowTargetSize)) ||
        failed(parser.parseRBrace())) {
      return failure();
    }
    if (firstRow) {
      firstScope = rowScope;
      sourceScope = rowScope;
      target = rowTarget;
      targetType = rowTargetType;
      targetSize = rowTargetSize;
      firstRow = false;
    } else if (rowScope.has_value() != firstScope.has_value() ||
               (rowScope.has_value() && rowScope->name != firstScope->name) ||
               rowTarget.name != target.name || rowTargetType != targetType ||
               rowTargetSize.name != targetSize.name) {
      return parser.emitError(
          parser.getCurrentLocation(),
          "each operation must use the same scope and target resource");
    }
    sourceKeys.push_back(sourceKey);
    sourceOffsets.push_back(sourceOffset);
    targetOffsets.push_back(targetOffset);
    targetEnds.push_back(targetEnd);
    targetLengths.push_back(targetLength);
  } while (succeeded(parser.parseOptionalComma()));
  return success();
}

static void printAsyncParameterGatherOperations(
    OpAsmPrinter &p, Operation *op, Value sourceScope, ValueRange sourceKeys,
    ValueRange sourceOffsets, Value target, Type targetType, Value targetSize,
    ValueRange targetOffsets, ValueRange targetEnds, ValueRange targetLengths) {
  p.increaseIndent();
  p.printNewline();
  llvm::interleave(
      llvm::zip_equal(sourceKeys, sourceOffsets, targetOffsets, targetEnds,
                      targetLengths),
      [&](std::tuple<Value, Value, Value, Value, Value> it) {
        auto [sourceKey, sourceOffset, targetOffset, targetEnd, targetLength] =
            it;
        printParameterReference(p, op, sourceScope, sourceKey);
        p << "[";
        p.printOperand(sourceOffset);
        p << "] -> ";
        p.printOperand(target);
        p << "[";
        p.printOperand(targetOffset);
        p << " to ";
        p.printOperand(targetEnd);
        p << " for ";
        p.printOperand(targetLength);
        p << "] : ";
        p.printType(targetType);
        p << "{";
        p.printOperand(targetSize);
        p << "}";
      },
      [&]() {
        p << ',';
        p.printNewline();
      });
  p.decreaseIndent();
  p.printNewline();
}

//===----------------------------------------------------------------------===//
// custom<CmdParameterGatherOperations>(
//     $source_scope, $source_keys, $source_offsets,
//     $target, type($target), $target_size, $target_offsets, $target_lengths)
//===----------------------------------------------------------------------===//

static ParseResult parseCmdParameterGatherOperations(
    OpAsmParser &parser,
    std::optional<OpAsmParser::UnresolvedOperand> &sourceScope,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &sourceKeys,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &sourceOffsets,
    OpAsmParser::UnresolvedOperand &target, Type &targetType,
    OpAsmParser::UnresolvedOperand &targetSize,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &targetOffsets,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &targetLengths) {
  std::optional<OpAsmParser::UnresolvedOperand> firstScope;
  bool firstRow = true;
  do {
    std::optional<OpAsmParser::UnresolvedOperand> rowScope;
    OpAsmParser::UnresolvedOperand sourceKey;
    OpAsmParser::UnresolvedOperand sourceOffset;
    OpAsmParser::UnresolvedOperand targetOffset;
    OpAsmParser::UnresolvedOperand targetLength;
    OpAsmParser::UnresolvedOperand rowTarget;
    Type rowTargetType;
    OpAsmParser::UnresolvedOperand rowTargetSize;
    if (failed(parseParameterReference(parser, rowScope, sourceKey)) ||
        failed(parser.parseLSquare()) ||
        failed(parser.parseOperand(sourceOffset)) ||
        failed(parser.parseRSquare()) || failed(parser.parseArrow()) ||
        failed(parser.parseOperand(rowTarget)) ||
        failed(parser.parseLSquare()) ||
        failed(parser.parseOperand(targetOffset)) ||
        failed(parser.parseKeyword("for")) ||
        failed(parser.parseOperand(targetLength)) ||
        failed(parser.parseRSquare()) ||
        failed(parser.parseColonType(rowTargetType)) ||
        failed(parser.parseLBrace()) ||
        failed(parser.parseOperand(rowTargetSize)) ||
        failed(parser.parseRBrace())) {
      return failure();
    }
    if (firstRow) {
      firstScope = rowScope;
      sourceScope = rowScope;
      target = rowTarget;
      targetType = rowTargetType;
      targetSize = rowTargetSize;
      firstRow = false;
    } else if (rowScope.has_value() != firstScope.has_value() ||
               (rowScope.has_value() && rowScope->name != firstScope->name) ||
               rowTarget.name != target.name || rowTargetType != targetType ||
               rowTargetSize.name != targetSize.name) {
      return parser.emitError(
          parser.getCurrentLocation(),
          "each operation must use the same scope and target resource");
    }
    sourceKeys.push_back(sourceKey);
    sourceOffsets.push_back(sourceOffset);
    targetOffsets.push_back(targetOffset);
    targetLengths.push_back(targetLength);
  } while (succeeded(parser.parseOptionalComma()));
  return success();
}

static void printCmdParameterGatherOperations(
    OpAsmPrinter &p, Operation *op, Value sourceScope, ValueRange sourceKeys,
    ValueRange sourceOffsets, Value target, Type targetType, Value targetSize,
    ValueRange targetOffsets, ValueRange targetLengths) {
  p.increaseIndent();
  p.printNewline();
  llvm::interleave(
      llvm::zip_equal(sourceKeys, sourceOffsets, targetOffsets, targetLengths),
      [&](std::tuple<Value, Value, Value, Value> it) {
        auto [sourceKey, sourceOffset, targetOffset, targetLength] = it;
        printParameterReference(p, op, sourceScope, sourceKey);
        p << "[";
        p.printOperand(sourceOffset);
        p << "] -> ";
        p.printOperand(target);
        p << "[";
        p.printOperand(targetOffset);
        p << " for ";
        p.printOperand(targetLength);
        p << "] : ";
        p.printType(targetType);
        p << "{";
        p.printOperand(targetSize);
        p << "}";
      },
      [&]() {
        p << ',';
        p.printNewline();
      });
  p.decreaseIndent();
  p.printNewline();
}

//===----------------------------------------------------------------------===//
// custom<AsyncParameterScatterOperations>(
//     $source, type($source), $source_size, $source_offsets, $source_ends,
//     $source_lengths, $target_scope, $target_keys, $target_offsets)
//===----------------------------------------------------------------------===//

static ParseResult parseAsyncParameterScatterOperations(
    OpAsmParser &parser, OpAsmParser::UnresolvedOperand &source,
    Type &sourceType, OpAsmParser::UnresolvedOperand &sourceSize,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &sourceOffsets,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &sourceEnds,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &sourceLengths,
    std::optional<OpAsmParser::UnresolvedOperand> &targetScope,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &targetKeys,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &targetOffsets) {
  std::optional<OpAsmParser::UnresolvedOperand> firstScope;
  bool firstRow = true;
  do {
    OpAsmParser::UnresolvedOperand rowSource;
    Type rowSourceType;
    OpAsmParser::UnresolvedOperand rowSourceSize;
    OpAsmParser::UnresolvedOperand sourceOffset;
    OpAsmParser::UnresolvedOperand sourceEnd;
    OpAsmParser::UnresolvedOperand sourceLength;
    std::optional<OpAsmParser::UnresolvedOperand> rowScope;
    OpAsmParser::UnresolvedOperand targetKey;
    OpAsmParser::UnresolvedOperand targetOffset;
    if (failed(parser.parseOperand(rowSource)) ||
        failed(parser.parseLSquare()) ||
        failed(parser.parseOperand(sourceOffset)) ||
        failed(parser.parseKeyword("to")) ||
        failed(parser.parseOperand(sourceEnd)) ||
        failed(parser.parseKeyword("for")) ||
        failed(parser.parseOperand(sourceLength)) ||
        failed(parser.parseRSquare()) ||
        failed(parser.parseColonType(rowSourceType)) ||
        failed(parser.parseLBrace()) ||
        failed(parser.parseOperand(rowSourceSize)) ||
        failed(parser.parseRBrace()) || failed(parser.parseArrow()) ||
        failed(parseParameterReference(parser, rowScope, targetKey)) ||
        failed(parser.parseLSquare()) ||
        failed(parser.parseOperand(targetOffset)) ||
        failed(parser.parseRSquare())) {
      return failure();
    }
    if (firstRow) {
      source = rowSource;
      sourceType = rowSourceType;
      sourceSize = rowSourceSize;
      firstScope = rowScope;
      targetScope = rowScope;
      firstRow = false;
    } else if (rowSource.name != source.name || rowSourceType != sourceType ||
               rowSourceSize.name != sourceSize.name ||
               rowScope.has_value() != firstScope.has_value() ||
               (rowScope.has_value() && rowScope->name != firstScope->name)) {
      return parser.emitError(
          parser.getCurrentLocation(),
          "each operation must use the same source resource and scope");
    }
    sourceOffsets.push_back(sourceOffset);
    sourceEnds.push_back(sourceEnd);
    sourceLengths.push_back(sourceLength);
    targetKeys.push_back(targetKey);
    targetOffsets.push_back(targetOffset);
  } while (succeeded(parser.parseOptionalComma()));
  return success();
}

static void printAsyncParameterScatterOperations(
    OpAsmPrinter &p, Operation *op, Value source, Type sourceType,
    Value sourceSize, ValueRange sourceOffsets, ValueRange sourceEnds,
    ValueRange sourceLengths, Value targetScope, ValueRange targetKeys,
    ValueRange targetOffsets) {
  p.increaseIndent();
  p.printNewline();
  llvm::interleave(
      llvm::zip_equal(sourceOffsets, sourceEnds, sourceLengths, targetKeys,
                      targetOffsets),
      [&](std::tuple<Value, Value, Value, Value, Value> it) {
        auto [sourceOffset, sourceEnd, sourceLength, targetKey, targetOffset] =
            it;
        p.printOperand(source);
        p << "[";
        p.printOperand(sourceOffset);
        p << " to ";
        p.printOperand(sourceEnd);
        p << " for ";
        p.printOperand(sourceLength);
        p << "] : ";
        p.printType(sourceType);
        p << "{";
        p.printOperand(sourceSize);
        p << "} -> ";
        printParameterReference(p, op, targetScope, targetKey);
        p << "[";
        p.printOperand(targetOffset);
        p << "]";
      },
      [&]() {
        p << ',';
        p.printNewline();
      });
  p.decreaseIndent();
  p.printNewline();
}

//===----------------------------------------------------------------------===//
// custom<CmdParameterScatterOperations>(
//     $source, type($source), $source_size, $source_offsets, $source_lengths,
//     $target_scope, $target_keys, $target_offsets)
//===----------------------------------------------------------------------===//

static ParseResult parseCmdParameterScatterOperations(
    OpAsmParser &parser, OpAsmParser::UnresolvedOperand &source,
    Type &sourceType, OpAsmParser::UnresolvedOperand &sourceSize,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &sourceOffsets,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &sourceLengths,
    std::optional<OpAsmParser::UnresolvedOperand> &targetScope,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &targetKeys,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &targetOffsets) {
  std::optional<OpAsmParser::UnresolvedOperand> firstScope;
  bool firstRow = true;
  do {
    OpAsmParser::UnresolvedOperand rowSource;
    Type rowSourceType;
    OpAsmParser::UnresolvedOperand rowSourceSize;
    OpAsmParser::UnresolvedOperand sourceOffset;
    OpAsmParser::UnresolvedOperand sourceLength;
    std::optional<OpAsmParser::UnresolvedOperand> rowScope;
    OpAsmParser::UnresolvedOperand targetKey;
    OpAsmParser::UnresolvedOperand targetOffset;
    if (failed(parser.parseOperand(rowSource)) ||
        failed(parser.parseLSquare()) ||
        failed(parser.parseOperand(sourceOffset)) ||
        failed(parser.parseKeyword("for")) ||
        failed(parser.parseOperand(sourceLength)) ||
        failed(parser.parseRSquare()) ||
        failed(parser.parseColonType(rowSourceType)) ||
        failed(parser.parseLBrace()) ||
        failed(parser.parseOperand(rowSourceSize)) ||
        failed(parser.parseRBrace()) || failed(parser.parseArrow()) ||
        failed(parseParameterReference(parser, rowScope, targetKey)) ||
        failed(parser.parseLSquare()) ||
        failed(parser.parseOperand(targetOffset)) ||
        failed(parser.parseRSquare())) {
      return failure();
    }
    if (firstRow) {
      source = rowSource;
      sourceType = rowSourceType;
      sourceSize = rowSourceSize;
      firstScope = rowScope;
      targetScope = rowScope;
      firstRow = false;
    } else if (rowSource.name != source.name || rowSourceType != sourceType ||
               rowSourceSize.name != sourceSize.name ||
               rowScope.has_value() != firstScope.has_value() ||
               (rowScope.has_value() && rowScope->name != firstScope->name)) {
      return parser.emitError(
          parser.getCurrentLocation(),
          "each operation must use the same source resource and scope");
    }
    sourceOffsets.push_back(sourceOffset);
    sourceLengths.push_back(sourceLength);
    targetKeys.push_back(targetKey);
    targetOffsets.push_back(targetOffset);
  } while (succeeded(parser.parseOptionalComma()));
  return success();
}

static void printCmdParameterScatterOperations(
    OpAsmPrinter &p, Operation *op, Value source, Type sourceType,
    Value sourceSize, ValueRange sourceOffsets, ValueRange sourceLengths,
    Value targetScope, ValueRange targetKeys, ValueRange targetOffsets) {
  p.increaseIndent();
  p.printNewline();
  llvm::interleave(
      llvm::zip_equal(sourceOffsets, sourceLengths, targetKeys, targetOffsets),
      [&](std::tuple<Value, Value, Value, Value> it) {
        auto [sourceOffset, sourceLength, targetKey, targetOffset] = it;
        p.printOperand(source);
        p << "[";
        p.printOperand(sourceOffset);
        p << " for ";
        p.printOperand(sourceLength);
        p << "] : ";
        p.printType(sourceType);
        p << "{";
        p.printOperand(sourceSize);
        p << "} -> ";
        printParameterReference(p, op, targetScope, targetKey);
        p << "[";
        p.printOperand(targetOffset);
        p << "]";
      },
      [&]() {
        p << ',';
        p.printNewline();
      });
  p.decreaseIndent();
  p.printNewline();
}

//===----------------------------------------------------------------------===//
// custom<ResourceRegion>($operands, type($operands), $operand_sizes,
//                        type($results), $result_sizes,
//                        $tied_operands, $body)
//===----------------------------------------------------------------------===//

static ParseResult parseResourceRegion(
    OpAsmParser &parser,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operands,
    SmallVectorImpl<Type> &operandTypes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operandSizes,
    SmallVectorImpl<Type> &resultTypes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resultSizes,
    ArrayAttr &tiedOperands, Region &body) {
  SmallVector<OpAsmParser::Argument, 16> regionArgs;
  if (failed(parser.parseLParen())) {
    return failure();
  }
  if (failed(parser.parseOptionalRParen())) {
    do {
      // Reserve entries in the lists.
      operands.emplace_back();
      operandTypes.emplace_back();
      operandSizes.emplace_back();
      regionArgs.emplace_back();
      if (failed(parser.parseOperand(operands.back())) ||
          failed(parser.parseKeyword("as")) ||
          failed(parser.parseArgument(regionArgs.back())) ||
          failed(parser.parseColon()) ||
          failed(parseSizeAwareType(parser, operandTypes.back(),
                                    operandSizes.back()))) {
        return failure();
      }
    } while (succeeded(parser.parseOptionalComma()));
    if (failed(parser.parseRParen())) {
      return failure();
    }
  }

  if (succeeded(parser.parseOptionalArrow())) {
    if (succeeded(parser.parseOptionalLParen())) {
      if (succeeded(parser.parseOptionalRParen())) {
        // -> ()
      } else if (failed(parseShapedResultList(parser, operands, operandTypes,
                                              operandSizes, resultTypes,
                                              resultSizes, tiedOperands)) ||
                 failed(parser.parseRParen())) {
        return failure();
      }
    } else {
      if (failed(parseShapedResultList(parser, operands, operandTypes,
                                       operandSizes, resultTypes, resultSizes,
                                       tiedOperands))) {
        return failure();
      }
    }
  }

  for (auto [iterArg, type] : llvm::zip_equal(regionArgs, operandTypes)) {
    iterArg.type = type;
  }
  return parser.parseRegion(body, regionArgs);
}

static void printResourceRegion(OpAsmPrinter &p, Operation *op,
                                ValueRange operands, TypeRange operandTypes,
                                ValueRange operandSizes, TypeRange resultTypes,
                                ValueRange resultSizes, ArrayAttr tiedOperands,
                                Region &body) {
  p << "(";
  llvm::interleaveComma(
      llvm::zip_equal(operands, body.getArguments()), p, [&](auto it) {
        auto operand = std::get<0>(it);
        auto arg = std::get<1>(it);
        p << operand;
        p << " as ";
        p << arg;
        p << ": ";
        p << arg.getType();
        if (isa<IREE::Util::SizeAwareTypeInterface>(arg.getType())) {
          p << "{" << operandSizes.front() << "}";
          operandSizes = operandSizes.drop_front(1);
        }
      });
  p << ")";
  if (!resultTypes.empty()) {
    p << " -> ";
    if (resultTypes.size() != 1) {
      p << "(";
    }
    printShapedResultList(p, op, operands, operandTypes, operandSizes,
                          resultTypes, resultSizes, tiedOperands);
    if (resultTypes.size() != 1) {
      p << ")";
    }
  }
  p << " ";
  p.printRegion(body, /*printEntryBlockArgs=*/false,
                /*printBlockTerminators=*/true);
}

//===----------------------------------------------------------------------===//
// custom<ExplicitResourceRegion>($operands, type($operands), $operand_sizes,
//                                $body)
//===----------------------------------------------------------------------===//

static ParseResult parseExplicitResourceRegion(
    OpAsmParser &parser,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operands,
    SmallVectorImpl<Type> &operandTypes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &operandSizes,
    Region &body) {
  SmallVector<OpAsmParser::Argument, 16> regionArgs;
  if (failed(parser.parseLParen())) {
    return failure();
  }
  if (failed(parser.parseOptionalRParen())) {
    do {
      // Reserve entries in the lists.
      operands.emplace_back();
      operandTypes.emplace_back();
      operandSizes.emplace_back();
      regionArgs.emplace_back();
      if (failed(parser.parseOperand(operands.back())) ||
          failed(parser.parseKeyword("as")) ||
          failed(parser.parseArgument(regionArgs.back())) ||
          failed(parser.parseColon()) ||
          failed(parseSizeAwareType(parser, operandTypes.back(),
                                    operandSizes.back()))) {
        return failure();
      }
    } while (succeeded(parser.parseOptionalComma()));
    if (failed(parser.parseRParen())) {
      return failure();
    }
  }
  for (auto [iterArg, type] : llvm::zip_equal(regionArgs, operandTypes)) {
    iterArg.type = type;
  }
  if (failed(parser.parseRegion(body, regionArgs))) {
    return failure();
  }
  // HACK: I can't figure out how to make this work with the default parsing -
  // it doesn't call this like it should.
  IREE::Stream::CmdExecuteOp::ensureTerminator(
      body, parser.getBuilder(),
      parser.getEncodedSourceLoc(parser.getCurrentLocation()));
  return success();
}

static void printExplicitResourceRegion(OpAsmPrinter &p, Operation *op,
                                        ValueRange operands,
                                        TypeRange operandTypes,
                                        ValueRange operandSizes, Region &body) {
  p << "(";
  llvm::interleaveComma(
      llvm::zip_equal(operands, body.getArguments()), p, [&](auto it) {
        auto operand = std::get<0>(it);
        auto arg = std::get<1>(it);
        p << operand;
        p << " as ";
        p << arg;
        p << ": ";
        p << arg.getType();
        if (isa<IREE::Util::SizeAwareTypeInterface>(arg.getType())) {
          p << "{" << operandSizes.front() << "}";
          operandSizes = operandSizes.drop_front(1);
        }
      });
  p << ") ";
  p.printRegion(body, /*printEntryBlockArgs=*/false,
                /*printBlockTerminators=*/false);
}

//===----------------------------------------------------------------------===//
// custom<PackSliceRanges>($lifetime_intervals,
//                         $dynamic_slice_sizes,
//                         type($packed_offsets))
//===----------------------------------------------------------------------===//

static ParseResult parsePackSliceRanges(
    OpAsmParser &parser, ArrayAttr &lifetimeIntervals,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &dynamicSliceSizes,
    SmallVectorImpl<Type> &packedOffsetTypes) {
  auto indexType = parser.getBuilder().getIndexType();
  SmallVector<Attribute> lifetimeRangeValues;
  do {
    if (failed(parser.parseOptionalLSquare())) {
      break;
    }
    IntegerAttr lifetimeStart;
    IntegerAttr lifetimeEnd;
    OpAsmParser::UnresolvedOperand dynamicSliceSize;
    if (failed(parser.parseAttribute(lifetimeStart, indexType)) ||
        failed(parser.parseComma()) ||
        failed(parser.parseAttribute(lifetimeEnd, indexType)) ||
        failed(parser.parseRSquare()) || failed(parser.parseEqual()) ||
        failed(parser.parseOperand(dynamicSliceSize))) {
      return failure();
    }
    lifetimeRangeValues.push_back(lifetimeStart);
    lifetimeRangeValues.push_back(lifetimeEnd);
    dynamicSliceSizes.push_back(dynamicSliceSize);
    packedOffsetTypes.push_back(indexType);
  } while (succeeded(parser.parseOptionalComma()));
  lifetimeIntervals = parser.getBuilder().getArrayAttr(lifetimeRangeValues);
  return success();
}

static void printPackSliceRanges(OpAsmPrinter &p, Operation *op,
                                 ArrayAttr lifetimeIntervals,
                                 ValueRange dynamicSliceSizes,
                                 TypeRange packedOffsetTypes) {
  if (packedOffsetTypes.empty()) {
    return;
  }
  for (unsigned i = 0; i < packedOffsetTypes.size(); ++i) {
    auto lifetimeStart = lifetimeIntervals[i * 2];
    auto lifetimeEnd = lifetimeIntervals[i * 2 + 1];
    auto sliceSize = dynamicSliceSizes[i];
    p.printNewline();
    p << "  [";
    p.printAttributeWithoutType(lifetimeStart);
    p << ", ";
    p.printAttributeWithoutType(lifetimeEnd);
    p << "] = ";
    p.printOperand(sliceSize);
    if (i < packedOffsetTypes.size() - 1) {
      p << ",";
    }
  }
  p.printNewline();
}

//===----------------------------------------------------------------------===//
// custom<ConstantValueList>(type($results),
//                           $result_sizes,
//                           $values)
//===----------------------------------------------------------------------===//
//   !stream.resource<constant>{%sz} = #value,
//   !stream.resource<constant>{%sz} = #value

static ParseResult parseConstantValueList(
    OpAsmParser &parser, SmallVectorImpl<Type> &resultTypes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resultSizes,
    ArrayAttr &values) {
  SmallVector<Attribute> valueAttrs;
  do {
    Type resultType;
    OpAsmParser::UnresolvedOperand resultSize;
    Attribute valueAttr;
    if (failed(parseSizeAwareType(parser, resultType, resultSize)) ||
        failed(parser.parseEqual()) ||
        failed(parser.parseAttribute(valueAttr))) {
      return failure();
    }
    resultTypes.push_back(resultType);
    resultSizes.push_back(resultSize);
    valueAttrs.push_back(valueAttr);
  } while (succeeded(parser.parseOptionalComma()));
  values = parser.getBuilder().getArrayAttr(valueAttrs);
  return success();
}

static void printConstantValueList(OpAsmPrinter &p, Operation *op,
                                   TypeRange resultTypes,
                                   ValueRange resultSizes, ArrayAttr values) {
  if (resultTypes.empty()) {
    return;
  }
  for (unsigned i = 0; i < resultTypes.size(); ++i) {
    p.printNewline();
    p << "  ";
    printSizeAwareType(p, op, resultTypes[i], resultSizes[i]);
    p << " = ";
    p.printAttribute(values[i]);
    if (i < resultTypes.size() - 1) {
      p << ",";
    }
  }
}

//===----------------------------------------------------------------------===//
// custom<WorkgroupCountRegion>($body)
//===----------------------------------------------------------------------===//

static ParseResult parseWorkgroupCountRegion(OpAsmParser &parser,
                                             Region &body) {
  if (failed(parser.parseOptionalKeyword("workgroups"))) {
    // Omitted.
    return success();
  }

  SmallVector<OpAsmParser::Argument> args;
  if (failed(parser.parseArgumentList(args, AsmParser::Delimiter::Paren,
                                      /*allowType=*/true,
                                      /*allowAttrs=*/true))) {
    return failure();
  }

  // Return types must be 3 dimensions (workgroup count XYZ).
  SmallVector<Type> returnTypes;
  if (failed(parser.parseArrowTypeList(returnTypes))) {
    return failure();
  }
  if (returnTypes.size() != 3 ||
      !llvm::all_of(returnTypes, [](Type type) { return type.isIndex(); })) {
    return parser.emitError(parser.getCurrentLocation())
           << "workgroup count region must return the XYZ dimension counts";
  }

  // Parse region contents.
  if (failed(parser.parseRegion(body, args, /*enableNameShadowing=*/false))) {
    return failure();
  }

  // Verify the return types match.
  for (auto returnOp : body.getOps<IREE::Stream::ReturnOp>()) {
    for (auto [returnType, operandType] :
         llvm::zip_equal(returnTypes, returnOp.getOperandTypes())) {
      if (returnType != operandType) {
        return returnOp.emitOpError()
               << "operands do not match expected region return types";
      }
    }
  }

  return success();
}

static void printWorkgroupCountRegion(OpAsmPrinter &p, Operation *op,
                                      Region &body) {
  if (body.empty()) {
    return;
  }
  p << "workgroups(";
  auto args = body.getArguments();
  for (unsigned i = 0; i < args.size(); ++i) {
    if (i > 0) {
      p << ", ";
    }
    p.printRegionArgument(args[i]);
  }
  p << ")";
  Type indexType = IndexType::get(body.getContext());
  p.printArrowTypeList(TypeRange{indexType, indexType, indexType});
  p << " ";
  p.printRegion(body, /*printEntryBlockArgs=*/false,
                /*printBlockTerminators=*/true);
}

//===----------------------------------------------------------------------===//
// stream.resource.alloc
//===----------------------------------------------------------------------===//

// static
std::pair<IREE::Stream::ResourceAllocOp, SmallVector<Value>>
ResourceAllocOp::createSuballocations(
    Type resourceType, ArrayRef<Location> locs, ValueRange storageSizes,
    bool uninitialized, AffinityAttr affinityAttr, OpBuilder &builder) {
  assert(locs.size() == storageSizes.size() &&
         "expect locs and storageSizes to match");
  if (locs.empty()) {
    return {};
  }
  if (locs.size() == 1) {
    auto allocOp = IREE::Stream::ResourceAllocOp::create(
        builder, locs.front(), resourceType, storageSizes.front(),
        uninitialized, affinityAttr);
    return {allocOp, {allocOp.getResult()}};
  }
  auto fusedLoc = builder.getFusedLoc(locs);

  // NOTE: this is risky: we are assuming right now that all of the
  // allocations will fit within the constraints of the system. This is not
  // guaranteed: a very low maximum buffer range may lead to packed slabs
  // that are not fully addressable. For now we are processing models with
  // small enough workloads and our target devices are relatively lax on
  // things so long as we stay under UINT32_MAX boundaries.

  // Compute total size and the offsets of all suballocated resources via the
  // pack op.
  auto packOp = IREE::Stream::ResourcePackOp::create(
      builder, fusedLoc, /*offset=*/nullptr, storageSizes, affinityAttr);

  // Create the new alloca based on the total required size.
  auto allocOp = IREE::Stream::ResourceAllocOp::create(
      builder, fusedLoc, resourceType, packOp.getTotalLength(), uninitialized,
      affinityAttr);
  auto slab = allocOp.getResult();
  auto slabSize = packOp.getTotalLength();

  // Create subviews for all of the suballocated resources.
  SmallVector<Value> results;
  for (auto [loc, subviewOffset, subviewLength] :
       llvm::zip_equal(locs, packOp.getPackedOffsets(), storageSizes)) {
    results.push_back(
        IREE::Stream::ResourceSubviewOp::create(builder, loc, slab, slabSize,
                                                subviewOffset, subviewLength)
            .getResult());
  }
  return {allocOp, results};
}

//===----------------------------------------------------------------------===//
// stream.resource.alloca
//===----------------------------------------------------------------------===//

// static
std::pair<IREE::Stream::ResourceAllocaOp, SmallVector<Value>>
ResourceAllocaOp::createSuballocations(Type timepointType, Type resourceType,
                                       ArrayRef<Location> locs,
                                       ValueRange storageSizes,
                                       Value awaitTimepoint,
                                       AffinityAttr affinityAttr,
                                       OpBuilder &builder) {
  assert(locs.size() == storageSizes.size() &&
         "expect locs and storageSizes to match");
  if (locs.empty()) {
    return {};
  }
  if (locs.size() == 1) {
    auto allocaOp = IREE::Stream::ResourceAllocaOp::create(
        builder, locs.front(), resourceType, timepointType,
        storageSizes.front(),
        /*indeterminate_lifetime=*/UnitAttr{}, awaitTimepoint, affinityAttr);
    return {allocaOp, {allocaOp.getResult()}};
  }
  auto fusedLoc = builder.getFusedLoc(locs);

  // NOTE: this is risky: we are assuming right now that all of the
  // allocations will fit within the constraints of the system. This is not
  // guaranteed: a very low maximum buffer range may lead to packed slabs
  // that are not fully addressable. For now we are processing models with
  // small enough workloads and our target devices are relatively lax on
  // things so long as we stay under UINT32_MAX boundaries. If a user starts
  // hitting this the solution is to do in-place outputs such that we don't
  // need to allocate them; when possible that's always going to be better than
  // leaving them to the IREE compiled program to deal with.

  // All slices are 0-0 (overlapping).
  size_t sliceCount = locs.size();
  SmallVector<int64_t> lifetimeIntervals(sliceCount * 2, 0);

  // Compute total size and the offsets of all suballocated resources via the
  // pack op.
  auto indexType = builder.getIndexType();
  llvm::Repeated<Type> packedOffsetTypes(sliceCount, indexType);
  auto packOp = IREE::Stream::ResourcePackOp::create(
      builder, fusedLoc, indexType, packedOffsetTypes, /*offset=*/nullptr,
      builder.getIndexArrayAttr(lifetimeIntervals), storageSizes, affinityAttr);

  // Create the new alloca based on the total required size.
  auto allocaOp = IREE::Stream::ResourceAllocaOp::create(
      builder, fusedLoc, resourceType, timepointType, packOp.getTotalLength(),
      /*indeterminate_lifetime=*/UnitAttr{}, awaitTimepoint, affinityAttr);
  auto slab = allocaOp.getResult();
  auto slabSize = packOp.getTotalLength();

  // Create subviews for all of the suballocated resources.
  SmallVector<Value> results;
  for (auto [loc, subviewOffset, subviewLength] :
       llvm::zip_equal(locs, packOp.getPackedOffsets(), storageSizes)) {
    results.push_back(
        IREE::Stream::ResourceSubviewOp::create(builder, loc, slab, slabSize,
                                                subviewOffset, subviewLength)
            .getResult());
  }
  return {allocaOp, results};
}

//===----------------------------------------------------------------------===//
// stream.resource.retain
//===----------------------------------------------------------------------===//

//===----------------------------------------------------------------------===//
// stream.resource.release
//===----------------------------------------------------------------------===//

void ResourceReleaseOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "is_terminal");
}

//===----------------------------------------------------------------------===//
// stream.resource.is_terminal
//===----------------------------------------------------------------------===//

void ResourceIsTerminalOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "is_terminal");
}

//===----------------------------------------------------------------------===//
// stream.resource.try_map
//===----------------------------------------------------------------------===//

LogicalResult ResourceTryMapOp::verify() {
  ResourceTryMapOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.resource.load
//===----------------------------------------------------------------------===//

LogicalResult ResourceLoadOp::verify() {
  ResourceLoadOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.resource.store
//===----------------------------------------------------------------------===//

LogicalResult ResourceStoreOp::verify() {
  ResourceStoreOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getTarget(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.resource.pack
//===----------------------------------------------------------------------===//

void ResourcePackOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  // TODO(benvanik): figure out if we can get the names to coalesce when there
  // are multiple results. Ideally we'd have `%total_length, %offsets:123` but
  // unfortunately all get splatted out and create 10k+ char lines that are a
  // pain to read.
  // setNameFn(total_length(), "total_length");
  // for (auto packedOffset : llvm::enumerate(packed_offsets())) {
  // setNameFn(packedOffset.value(),
  //           "offset" + std::to_string(packedOffset.index()));
  // }
}

void ResourcePackOp::build(OpBuilder &builder, OperationState &state,
                           Value offset, ValueRange valueSizes,
                           IREE::Stream::AffinityAttr affinity) {
  // All slices are 0-0 (overlapping).
  size_t sliceCount = valueSizes.size();
  SmallVector<int64_t> lifetimeIntervals(sliceCount * 2, 0);
  auto indexType = builder.getIndexType();
  llvm::Repeated<Type> indexTypes(sliceCount, indexType);
  build(builder, state, indexType, indexTypes, offset,
        builder.getIndexArrayAttr(lifetimeIntervals), valueSizes, affinity);
}

LogicalResult ResourcePackOp::verify() {
  ResourcePackOp op = *this;
  size_t sliceCount = op.getPackedOffsets().size();
  if (op.getLifetimeIntervals().size() != sliceCount * 2) {
    return op.emitOpError() << "requires a [start, end] range for each slice";
  }
  if (op.getDynamicSliceSizes().size() != sliceCount) {
    return op.emitOpError() << "requires a size for each slice";
  }
  return success();
}

SmallVector<ResourcePackOp::Slice> ResourcePackOp::getSlices() {
  auto intervalPairs = getLifetimeIntervals().getValue();
  auto sizes = getDynamicSliceSizes();
  auto offsets = getPackedOffsets();
  SmallVector<ResourcePackOp::Slice> slices(offsets.size());
  for (size_t i = 0; i < offsets.size(); ++i) {
    int64_t start = cast<IntegerAttr>(intervalPairs[i * 2 + 0]).getInt();
    int64_t end = cast<IntegerAttr>(intervalPairs[i * 2 + 1]).getInt();
    slices[i] = {start, end, sizes[i], offsets[i]};
  }
  return slices;
}

//===----------------------------------------------------------------------===//
// stream.resource.constants
//===----------------------------------------------------------------------===//

LogicalResult ResourceConstantsOp::verify() {
  ResourceConstantsOp op = *this;
  size_t count = op.getResults().size();
  if (op.getResultSizes().size() != count || op.getValues().size() != count) {
    return op.emitOpError() << "mismatched constant/result counts";
  }
  if (!llvm::all_equal(op.getResults().getTypes())) {
    return op.emitOpError() << "all results must be of the same type";
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.resource.subview
//===----------------------------------------------------------------------===//

LogicalResult ResourceSubviewOp::verify() {
  ResourceSubviewOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  return success();
}

bool ResourceSubviewOp::isMetadata() { return true; }

Value ResourceSubviewOp::getViewSource() { return getSource(); }

Value ResourceSubviewOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getSource());
}

::std::optional<unsigned>
ResourceSubviewOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // source
}

SmallVector<int64_t> ResourceSubviewOp::getTiedResultOperandIndices() {
  return {0}; // source
}

// static
IREE::Stream::ResourceSubviewOp ResourceSubviewOp::findSubviewOp(Value value) {
  while (value) {
    auto *definingOp = value.getDefiningOp();
    if (!definingOp) {
      // Defined as a block argument - stop walk.
      break;
    } else if (isa<IREE::Stream::TimelineOpInterface>(definingOp)) {
      // Don't traverse timeline ops as time travel isn't possible (yet).
      break;
    } else if (auto subviewOp =
                   dyn_cast<IREE::Stream::ResourceSubviewOp>(definingOp)) {
      // Found!
      return subviewOp;
    } else if (auto tiedOp =
                   dyn_cast<IREE::Util::TiedOpInterface>(definingOp)) {
      // Continue walking up through the tied operand.
      value = tiedOp.getTiedResultOperand(value);
    } else {
      break;
    }
  }
  return {};
}

//===----------------------------------------------------------------------===//
// stream.resource.transients
//===----------------------------------------------------------------------===//

bool ResourceTransientsOp::pinsValueAffinity() {
  return getAffinity().has_value();
}

Value ResourceTransientsOp::getTiedResult(unsigned resultIndex) {
  if (resultIndex == 0) {
    return IREE::Util::TiedOpInterface::findTiedBaseValue(getResource());
  }
  return {}; // timepoint result is not tied
}

::std::optional<unsigned>
ResourceTransientsOp::getTiedResultOperandIndex(unsigned resultIndex) {
  if (resultIndex == 0) {
    return {0}; // resource
  }
  return std::nullopt; // timepoint result is not tied
}

SmallVector<int64_t> ResourceTransientsOp::getTiedResultOperandIndices() {
  return {0, -1}; // resource is tied, timepoint is not
}

namespace {

struct FoldConsecutiveResourceTransientsOps
    : OpRewritePattern<ResourceTransientsOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(ResourceTransientsOp op,
                                PatternRewriter &rewriter) const override {
    auto resourceOp = op.getResource().getDefiningOp<ResourceTransientsOp>();
    if (!resourceOp) {
      return failure();
    }
    // Replace the resource and dims with those from the source transients op.
    // The outer storage wins (most recent specification).
    rewriter.modifyOpInPlace(op, [&]() {
      op.getResourceMutable().assign(resourceOp.getResource());
      op.getResourceSizeMutable().assign(resourceOp.getResourceSize());
    });
    return success();
  }
};

} // namespace

void ResourceTransientsOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.insert<FoldConsecutiveResourceTransientsOps>(context);
}

//===----------------------------------------------------------------------===//
// stream.file.constant
//===----------------------------------------------------------------------===//

void FileConstantOp::getAsmResultNames(mlir::OpAsmSetValueNameFn setNameFn) {
  setNameFn(getResult(), "file");
}

IREE::Util::SubrangeOperand
FileConstantOp::getSubrangeOperand(unsigned operandIndex) {
  if (operandIndex == 0) {
    return IREE::Util::SubrangeOperand{getSource(), getSourceSize(),
                                       getSourceOffset(), getSourceLength()};
  } else {
    assert(false && "only source is a subrange");
    return {};
  }
}

void FileConstantOp::setSubrangeOperand(unsigned operandIndex,
                                        IREE::Util::SubrangeOperand operand) {
  assert(operandIndex == 0 && "only source is a subrange");
  getSourceMutable().assign(operand.resource);
  getSourceSizeMutable().assign(operand.resourceSize);
  getSourceOffsetMutable().assign(operand.offset);
  getSourceLengthMutable().assign(operand.length);
}

//===----------------------------------------------------------------------===//
// stream.file.read
//===----------------------------------------------------------------------===//

LogicalResult FileReadOp::verify() {
  FileReadOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getTarget(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.file.write
//===----------------------------------------------------------------------===//

LogicalResult FileWriteOp::verify() {
  FileWriteOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.tensor.import
//===----------------------------------------------------------------------===//

LogicalResult TensorImportOp::verify() {
  TensorImportOp op = *this;
  if (failed(verifyOpDynamicDims(op, op.getResultEncoding(),
                                 op.getResultEncodingDims())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  return success();
}

bool TensorImportOp::pinsValueAffinity() { return getAffinity().has_value(); }

Value TensorImportOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getSource());
}

::std::optional<unsigned>
TensorImportOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // source
}

SmallVector<int64_t> TensorImportOp::getTiedResultOperandIndices() {
  return {0}; // source
}

//===----------------------------------------------------------------------===//
// stream.tensor.export
//===----------------------------------------------------------------------===//

LogicalResult TensorExportOp::verify() {
  TensorExportOp op = *this;
  if (failed(verifyOpDynamicDims(op, op.getSourceEncoding(),
                                 op.getSourceEncodingDims())) ||
      failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize()))) {
    return failure();
  }
  return success();
}

bool TensorExportOp::pinsValueAffinity() { return getAffinity().has_value(); }

Value TensorExportOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getSource());
}

::std::optional<unsigned>
TensorExportOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // source
}

SmallVector<int64_t> TensorExportOp::getTiedResultOperandIndices() {
  return {0}; // source
}

//===----------------------------------------------------------------------===//
// stream.tensor.sizeof
//===----------------------------------------------------------------------===//

LogicalResult TensorSizeOfOp::verify() {
  TensorSizeOfOp op = *this;
  if (failed(verifyOpDynamicDims(op, op.getEncoding(), op.getEncodingDims()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.tensor.empty
//===----------------------------------------------------------------------===//

void TensorEmptyOp::getAsmResultNames(mlir::OpAsmSetValueNameFn setNameFn) {
  setNameFn(getResult(), "empty");
}

LogicalResult TensorEmptyOp::verify() {
  TensorEmptyOp op = *this;
  if (failed(verifyOpDynamicDims(op, op.getResultEncoding(),
                                 op.getResultEncodingDims()))) {
    return failure();
  }
  return success();
}

bool TensorEmptyOp::isMetadata() { return true; }

bool TensorEmptyOp::preferCloneToConsumers() { return true; }

//===----------------------------------------------------------------------===//
// stream.tensor.constant
//===----------------------------------------------------------------------===//

void TensorConstantOp::getAsmResultNames(mlir::OpAsmSetValueNameFn setNameFn) {
  setNameFn(getResult(), "cst");
}

LogicalResult TensorConstantOp::verify() {
  TensorConstantOp op = *this;
  if (failed(verifyOpDynamicDims(op, op.getResultEncoding(),
                                 op.getResultEncodingDims()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.tensor.splat
//===----------------------------------------------------------------------===//

LogicalResult TensorSplatOp::verify() {
  TensorSplatOp op = *this;
  if (failed(verifyOpDynamicDims(op, op.getResultEncoding(),
                                 op.getResultEncodingDims())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  return success();
}

bool TensorSplatOp::preferCloneToConsumers() { return true; }

//===----------------------------------------------------------------------===//
// stream.tensor.clone
//===----------------------------------------------------------------------===//

LogicalResult TensorCloneOp::verify() {
  TensorCloneOp op = *this;
  // Clones can't change encodings but they can change shape and element type
  // information.
  auto sourceEncoding = cast<RankedTensorType>(op.getSourceEncoding());
  auto resultEncoding = cast<RankedTensorType>(op.getResultEncoding());
  if (!IREE::Encoding::SerializableAttr::areCompatible(
          sourceEncoding.getEncoding(), resultEncoding.getEncoding())) {
    return op.emitOpError() << "clones changing tensor encoding from "
                            << sourceEncoding.getEncoding() << " to "
                            << resultEncoding.getEncoding() << "; not allowed";
  }
  if (failed(verifyOpDynamicDims(op, op.getSourceEncoding(),
                                 op.getSourceEncodingDims())) ||
      failed(verifyOpDynamicDims(op, op.getResultEncoding(),
                                 op.getResultEncodingDims())) ||
      failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  return success();
}

bool TensorCloneOp::preferCloneToConsumers() { return true; }

//===----------------------------------------------------------------------===//
// stream.tensor.encode
//===----------------------------------------------------------------------===//

LogicalResult TensorEncodeOp::verify() {
  TensorEncodeOp op = *this;
  if (failed(verifyOpDynamicDims(op, op.getSourceEncoding(),
                                 op.getSourceEncodingDims())) ||
      failed(verifyOpDynamicDims(op, op.getResultEncoding(),
                                 op.getResultEncodingDims())) ||
      failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.tensor.slice
//===----------------------------------------------------------------------===//

LogicalResult TensorSliceOp::verify() {
  TensorSliceOp op = *this;
  if (failed(verifyOpDynamicDims(op, op.getSourceEncoding(),
                                 op.getSourceEncodingDims())) ||
      failed(verifyOpDynamicDims(op, op.getResultEncoding(),
                                 op.getResultEncodingDims())) ||
      failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  auto sourceType = cast<ShapedType>(op.getSourceEncoding());
  if (op.getStartIndices().size() != sourceType.getRank() ||
      op.getLengths().size() != sourceType.getRank()) {
    return op.emitOpError() << "start_indices/lengths rank mismatch";
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.tensor.update
//===----------------------------------------------------------------------===//

LogicalResult TensorUpdateOp::verify() {
  TensorUpdateOp op = *this;
  if (failed(verifyOpDynamicDims(op, op.getUpdateEncoding(),
                                 op.getUpdateEncodingDims())) ||
      failed(verifyOpDynamicDims(op, op.getTargetEncoding(),
                                 op.getTargetEncodingDims())) ||
      failed(verifyOpValueSizes(op, op.getUpdate(), op.getUpdateSize())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

Value TensorUpdateOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getTarget());
}

::std::optional<unsigned>
TensorUpdateOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // target
}

SmallVector<int64_t> TensorUpdateOp::getTiedResultOperandIndices() {
  return {0}; // target
}

//===----------------------------------------------------------------------===//
// stream.tensor.fill
//===----------------------------------------------------------------------===//

LogicalResult TensorFillOp::verify() {
  TensorFillOp op = *this;
  if (failed(verifyOpDynamicDims(op, op.getTargetEncoding(),
                                 op.getTargetEncodingDims())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

Value TensorFillOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getTarget());
}

::std::optional<unsigned>
TensorFillOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // target
}

SmallVector<int64_t> TensorFillOp::getTiedResultOperandIndices() {
  return {0}; // target
}

//===----------------------------------------------------------------------===//
// stream.tensor.load
//===----------------------------------------------------------------------===//

LogicalResult TensorLoadOp::verify() {
  TensorLoadOp op = *this;
  if (failed(verifyOpDynamicDims(op, op.getSourceEncoding(),
                                 op.getSourceEncodingDims())) ||
      failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize()))) {
    return failure();
  }
  auto sourceType = cast<ShapedType>(op.getSourceEncoding());
  if (op.getIndices().size() != sourceType.getRank()) {
    return op.emitOpError() << "indices rank mismatch";
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.tensor.store
//===----------------------------------------------------------------------===//

LogicalResult TensorStoreOp::verify() {
  TensorStoreOp op = *this;
  if (failed(verifyOpDynamicDims(op, op.getTargetEncoding(),
                                 op.getTargetEncodingDims())) ||
      failed(verifyOpValueSizes(op, op.getTarget(), op.getTargetSize()))) {
    return failure();
  }
  auto targetType = cast<ShapedType>(op.getTargetEncoding());
  if (op.getIndices().size() != targetType.getRank()) {
    return op.emitOpError() << "indices rank mismatch";
  }
  return success();
}

Value TensorStoreOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getTarget());
}

::std::optional<unsigned>
TensorStoreOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // target
}

SmallVector<int64_t> TensorStoreOp::getTiedResultOperandIndices() {
  return {0}; // target
}

//===----------------------------------------------------------------------===//
// stream.tensor.trace
//===----------------------------------------------------------------------===//

LogicalResult TensorTraceOp::verify() {
  TensorTraceOp op = *this;
  if (op.getResources().size() != op.getResourceEncodings().size() ||
      op.getResources().size() != op.getResourceSizes().size()) {
    return op.emitOpError(
        "each resource needs a matching resource encoding and size "
        "(array length mismatch)");
  }
  auto resourceEncodingDims = op.getResourceEncodingDims();
  for (auto [resource, resourceSize, resourceEncoding] :
       llvm::zip_equal(op.getResources(), op.getResourceSizes(),
                       op.getResourceEncodings().getAsValueRange<TypeAttr>())) {
    int64_t dynamicDimCount =
        cast<ShapedType>(resourceEncoding).getNumDynamicDims();
    if (failed(verifyOpDynamicDims(
            op, resourceEncoding,
            resourceEncodingDims.take_front(dynamicDimCount))) ||
        failed(verifyOpValueSizes(op, resource, resourceSize))) {
      return failure();
    }
    resourceEncodingDims = resourceEncodingDims.drop_front(dynamicDimCount);
  }
  return success();
}

ValueRange TensorTraceOp::getOperandDynamicDims(unsigned idx) {
  auto resourceEncodings = getResourceEncodings().getValue();
  auto resourceEncodingDims = getResourceEncodingDims();
  for (unsigned i = 0; i <= idx; ++i) {
    auto resourceEncoding = cast<TypeAttr>(resourceEncodings[i]).getValue();
    int64_t dynamicDimCount =
        cast<ShapedType>(resourceEncoding).getNumDynamicDims();
    if (i == idx) {
      return resourceEncodingDims.take_front(dynamicDimCount);
    }
    resourceEncodingDims = resourceEncodingDims.drop_front(dynamicDimCount);
  }
  return ValueRange{};
}

ValueRange TensorTraceOp::getResultDynamicDims(unsigned idx) {
  return ValueRange{};
}

//===----------------------------------------------------------------------===//
// stream.tensor.dispatch
//===----------------------------------------------------------------------===//

LogicalResult TensorDispatchOp::verify() {
  TensorDispatchOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getMixedOperands(),
                                op.getOperandSizes())) ||
      failed(verifyOpValueSizes(op, op.getResults(), op.getResultSizes()))) {
    return failure();
  }
  if (failed(verifyOpDynamicDimsRange(op, op.getOperandEncodings(),
                                      op.getOperandEncodingDims())) ||
      failed(verifyOpDynamicDimsRange(op, op.getResultEncodings(),
                                      op.getResultEncodingDims()))) {
    return failure();
  }
  if (failed(verifyTiedOperandEncodings(op, op.getOperandEncodings(),
                                        op.getResultEncodings()))) {
    return failure();
  }
  return success();
}

static LogicalResult
verifyDispatchSymbolUses(Operation *op, ArrayAttr entryPointsAttr,
                         ValueRange workload,
                         SymbolTableCollection &symbolTable) {
  auto entryPointAttrs = entryPointsAttr.getAsRange<SymbolRefAttr>();
  if (entryPointAttrs.empty()) {
    return op->emitOpError() << "at least one entry point must be defined";
  }
  for (auto entryPointAttr : entryPointAttrs) {
    auto exportOp =
        symbolTable.lookupNearestSymbolFrom<IREE::Stream::ExecutableExportOp>(
            op, entryPointAttr);
    if (!exportOp) {
      // TODO(benvanik): there are a lot of tests that are assuming this is not
      // verified. We'll need to go add dummy executables for all of them. Today
      // we just bail on the verifier if the symbol isn't found.
      //
      // Should be:
      //   return op->emitOpError() << "undefined entry point: " <<
      //   entry_point();
      return success();
    }

    // Verify that the workload parameters captured match the target export.
    if (failed(verifyDispatchWorkload(op, exportOp, workload))) {
      return failure();
    }

    // TODO(benvanik): verify that the target function has matching operands.
  }
  return success();
}

LogicalResult
TensorDispatchOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  return verifyDispatchSymbolUses(getOperation(), getEntryPointsAttr(),
                                  getWorkload(), symbolTable);
}

std::pair<unsigned, unsigned>
TensorDispatchOp::getTiedOperandsIndexAndLength() {
  return getODSOperandIndexAndLength(1); // $operands
}

//===----------------------------------------------------------------------===//
// stream.tensor.parameter.load
//===----------------------------------------------------------------------===//

LogicalResult TensorParameterLoadOp::verify() {
  TensorParameterLoadOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.tensor.parameter.write
//===----------------------------------------------------------------------===//

LogicalResult TensorParameterWriteOp::verify() {
  TensorParameterWriteOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize()))) {
    return failure();
  }
  return success();
}

Value TensorParameterWriteOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getSource());
}

::std::optional<unsigned>
TensorParameterWriteOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // result tied to source
}

SmallVector<int64_t> TensorParameterWriteOp::getTiedResultOperandIndices() {
  return {0}; // result tied to source
}

//===----------------------------------------------------------------------===//
// stream.async.alloca
//===----------------------------------------------------------------------===//

bool AsyncAllocaOp::isMetadata() { return true; }

bool AsyncAllocaOp::preferCloneToConsumers() { return true; }

//===----------------------------------------------------------------------===//
// stream.async.constant
//===----------------------------------------------------------------------===//

bool AsyncConstantOp::isMetadata() { return true; }

void AsyncConstantOp::getAsmResultNames(mlir::OpAsmSetValueNameFn setNameFn) {
  setNameFn(getResult(), "cst");
}

LogicalResult AsyncConstantOp::verify() {
  AsyncConstantOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  return success();
}

void AsyncConstantOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  ranges.push_back({ResourceAccessBitfield::Write, getResult(), Value{},
                    getResultSize(), getResultSize()});
}

//===----------------------------------------------------------------------===//
// stream.async.splat
//===----------------------------------------------------------------------===//

void AsyncSplatOp::build(OpBuilder &builder, OperationState &state,
                         Type resultType, Value value, Value resultSize,
                         Attribute affinity, Value awaitTimepoint) {
  state.addTypes(resultType);
  if (awaitTimepoint) {
    state.addOperands(awaitTimepoint);
  }
  state.addOperands(value);
  state.addOperands(resultSize);
  if (affinity) {
    state.addAttribute("affinity", affinity);
  }
}

LogicalResult AsyncSplatOp::verify() {
  AsyncSplatOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  return success();
}

bool AsyncSplatOp::preferCloneToConsumers() { return true; }

void AsyncSplatOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  ranges.push_back({ResourceAccessBitfield::Write, getResult(), Value{},
                    getResultSize(), getResultSize()});
}

//===----------------------------------------------------------------------===//
// stream.async.clone
//===----------------------------------------------------------------------===//

LogicalResult AsyncCloneOp::verify() {
  AsyncCloneOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  return success();
}

bool AsyncCloneOp::preferCloneToConsumers() { return true; }

void AsyncCloneOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  ranges.push_back({ResourceAccessBitfield::Read, getSource(), Value{},
                    getSourceSize(), getSourceSize()});
  ranges.push_back({ResourceAccessBitfield::Write, getResult(), Value{},
                    getResultSize(), getResultSize()});
}

//===----------------------------------------------------------------------===//
// stream.async.slice
//===----------------------------------------------------------------------===//

LogicalResult AsyncSliceOp::verify() {
  AsyncSliceOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  return success();
}

void AsyncSliceOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  ranges.push_back({ResourceAccessBitfield::Read, getSource(),
                    getSourceOffset(), getSourceEnd(), getResultSize()});
  ranges.push_back({ResourceAccessBitfield::Write, getResult(), Value{},
                    getResultSize(), getResultSize()});
}

//===----------------------------------------------------------------------===//
// stream.async.fill
//===----------------------------------------------------------------------===//

LogicalResult AsyncFillOp::verify() {
  AsyncFillOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getResult(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

Value AsyncFillOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getTarget());
}

::std::optional<unsigned>
AsyncFillOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // target
}

SmallVector<int64_t> AsyncFillOp::getTiedResultOperandIndices() {
  return {0}; // target
}

void AsyncFillOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  ranges.push_back({ResourceAccessBitfield::Write, getTarget(),
                    getTargetOffset(), getTargetEnd(), getTargetLength()});
  ranges.push_back({ResourceAccessBitfield::Write, getResult(),
                    getTargetOffset(), getTargetEnd(), getTargetLength()});
}

//===----------------------------------------------------------------------===//
// stream.async.update
//===----------------------------------------------------------------------===//

LogicalResult AsyncUpdateOp::verify() {
  AsyncUpdateOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getUpdate(), op.getUpdateSize())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

Value AsyncUpdateOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getTarget());
}

::std::optional<unsigned>
AsyncUpdateOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // target
}

SmallVector<int64_t> AsyncUpdateOp::getTiedResultOperandIndices() {
  return {0}; // target
}

void AsyncUpdateOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  ranges.push_back({ResourceAccessBitfield::Read, getUpdate(), Value{},
                    getUpdateSize(), getUpdateSize()});
  ranges.push_back({ResourceAccessBitfield::Write, getTarget(),
                    getTargetOffset(), getTargetEnd(), getUpdateSize()});
  ranges.push_back({ResourceAccessBitfield::Write, getResult(),
                    getTargetOffset(), getTargetEnd(), getUpdateSize()});
}

//===----------------------------------------------------------------------===//
// stream.async.copy
//===----------------------------------------------------------------------===//

LogicalResult AsyncCopyOp::verify() {
  AsyncCopyOp op = *this;
  // TODO(ezhulenev): We should reject copy operations when we know that buffers
  // overlap. This will be verified at run time by command buffer validation,
  // but it would be better to reject invalid IR early.
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

Value AsyncCopyOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getTarget());
}

::std::optional<unsigned>
AsyncCopyOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // target
}

SmallVector<int64_t> AsyncCopyOp::getTiedResultOperandIndices() {
  return {0}; // target
}

void AsyncCopyOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  ranges.push_back({ResourceAccessBitfield::Read, getSource(),
                    getSourceOffset(), getSourceEnd(), getLength()});
  ranges.push_back({ResourceAccessBitfield::Write, getTarget(),
                    getTargetOffset(), getTargetEnd(), getLength()});
  ranges.push_back({ResourceAccessBitfield::Write, getResult(),
                    getTargetOffset(), getTargetEnd(), getLength()});
}

//===----------------------------------------------------------------------===//
// stream.async.collective
//===----------------------------------------------------------------------===//

static const char *getCollectiveParamKeyword(Attribute opAttr) {
  auto attr = cast<IREE::Stream::CollectiveAttr>(opAttr);
  switch (attr.getKind()) {
  case IREE::Stream::CollectiveKind::Broadcast:
    return "source";
  case IREE::Stream::CollectiveKind::Reduce:
    return "target";
  case IREE::Stream::CollectiveKind::Send:
    return "target";
  case IREE::Stream::CollectiveKind::Recv:
    return "source";
  case IREE::Stream::CollectiveKind::SendRecv:
    return "source_target_pair";
  default:
    return nullptr;
  }
}

static ParseResult parseCollectiveParam(
    OpAsmParser &parser, Attribute opAttr,
    std::optional<OpAsmParser::UnresolvedOperand> &optionalParamValue) {
  const char *keyword = getCollectiveParamKeyword(opAttr);
  if (!keyword) {
    return success(); // optional
  }
  OpAsmParser::UnresolvedOperand paramValue;
  if (failed(parser.parseKeyword(keyword)) || failed(parser.parseLParen()) ||
      failed(parser.parseOperand(paramValue)) || failed(parser.parseRParen())) {
    return failure();
  }
  optionalParamValue = paramValue;
  return success();
}

static void printCollectiveParam(OpAsmPrinter &p, Operation *op,
                                 Attribute opAttr, Value paramValue) {
  const char *keyword = getCollectiveParamKeyword(opAttr);
  if (paramValue) {
    assert(keyword && "collective op must have a param keyword");
    p << keyword << "(";
    p.printOperand(paramValue);
    p << ") ";
  }
}

LogicalResult AsyncCollectiveOp::verify() {
  AsyncCollectiveOp op = *this;

  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getTargetSize()))) {
    return failure();
  }

  const bool hasParam = !!op.getParam();
  const char *paramKeyword = getCollectiveParamKeyword(op.getOp());
  if (paramKeyword) {
    // Param required.
    if (!hasParam) {
      return op.emitOpError() << "collective operation requires a "
                              << paramKeyword << " parameter but none present";
    }
  } else {
    // No param required.
    if (hasParam) {
      return op.emitOpError() << "collective operation does not require a "
                                 "parameter but one present";
    }
  }

  return success();
}

Value AsyncCollectiveOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getTarget());
}

::std::optional<unsigned>
AsyncCollectiveOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // target
}

SmallVector<int64_t> AsyncCollectiveOp::getTiedResultOperandIndices() {
  return {0}; // target
}

void AsyncCollectiveOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  ranges.push_back({ResourceAccessBitfield::Read, getSource(),
                    getSourceOffset(), getSourceEnd(), getSourceLength()});
  ranges.push_back({ResourceAccessBitfield::Write, getTarget(),
                    getTargetOffset(), getTargetEnd(), getTargetLength()});
  ranges.push_back({ResourceAccessBitfield::Write, getResult(),
                    getTargetOffset(), getTargetEnd(), getTargetLength()});
}

//===----------------------------------------------------------------------===//
// stream.async.barrier
//===----------------------------------------------------------------------===//

bool AsyncBarrierOp::isMetadata() { return true; }

LogicalResult AsyncBarrierOp::verify() { return success(); }

Value AsyncBarrierOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getSource());
}

::std::optional<unsigned>
AsyncBarrierOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // source
}

SmallVector<int64_t> AsyncBarrierOp::getTiedResultOperandIndices() {
  return {0}; // source
}

//===----------------------------------------------------------------------===//
// stream.async.transfer
//===----------------------------------------------------------------------===//

LogicalResult AsyncTransferOp::verify() {
  AsyncTransferOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  return success();
}

IREE::Stream::AffinityAttr AsyncTransferOp::getAffinityAttr() {
  if (auto executionAffinityAttr = getExecutionAffinityAttr()) {
    return executionAffinityAttr;
  }
  return getDefaultExecutionAffinityAttr();
}

IREE::Stream::AffinityAttr AsyncTransferOp::getDefaultExecutionAffinityAttr() {
  auto sourceType = cast<IREE::Stream::ResourceType>(getSource().getType());
  auto resultType = cast<IREE::Stream::ResourceType>(getResult().getType());
  if (sourceType.getLifetime() == IREE::Stream::Lifetime::Staging &&
      resultType.getLifetime() == IREE::Stream::Lifetime::Staging) {
    // TODO(multi-device): figure out how to model staging->staging transfers.
    return getSourceAffinityAttr();
  } else if (sourceType.getLifetime() == IREE::Stream::Lifetime::External ||
             sourceType.getLifetime() == IREE::Stream::Lifetime::Staging) {
    // If source is staging then the op should execute on the consumer.
    return getResultAffinityAttr();
  } else if (resultType.getLifetime() == IREE::Stream::Lifetime::External ||
             resultType.getLifetime() == IREE::Stream::Lifetime::Staging) {
    // If result is staging then the op should execute on the producer.
    return getSourceAffinityAttr();
  } else {
    // Default to source affinity (today we transfer from sources).
    return getSourceAffinityAttr();
  }
}

void AsyncTransferOp::setAffinityAttr(IREE::Stream::AffinityAttr value) {
  if (value) {
    setExecutionAffinityAttr(value);
  } else {
    removeExecutionAffinityAttr();
  }
}

IREE::Stream::AffinityAttr AsyncTransferOp::getResultAffinityAttr() {
  return getTargetAffinityAttr();
}

void AsyncTransferOp::removeAffinityAttrs() {
  removeSourceAffinityAttr();
  removeTargetAffinityAttr();
  removeExecutionAffinityAttr();
}

AsyncTransferOp AsyncTransferOp::create(OpBuilder &builder, Location location,
                                        Type type, Value source,
                                        Value sourceSize, Value resultSize,
                                        AffinityAttr sourceAffinity,
                                        AffinityAttr targetAffinity) {
  return create(builder, location, type, source, sourceSize, resultSize,
                sourceAffinity, targetAffinity, nullptr);
}

void AsyncTransferOp::build(OpBuilder &builder, OperationState &state,
                            Type type, Value source, Value sourceSize,
                            Value resultSize, AffinityAttr sourceAffinity,
                            AffinityAttr targetAffinity) {
  build(builder, state, type, source, sourceSize, resultSize, sourceAffinity,
        targetAffinity, nullptr);
}

void AsyncTransferOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  ranges.push_back({ResourceAccessBitfield::Read, getSource(), Value{},
                    getSourceSize(), getSourceSize()});
  ranges.push_back({ResourceAccessBitfield::Write, getResult(), Value{},
                    getResultSize(), getResultSize()});
}

//===----------------------------------------------------------------------===//
// stream.async.cast
//===----------------------------------------------------------------------===//

Value AsyncCastOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getSource());
}

::std::optional<unsigned>
AsyncCastOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // source
}

SmallVector<int64_t> AsyncCastOp::getTiedResultOperandIndices() {
  return {0}; // source
}

LogicalResult AsyncCastOp::verify() {
  AsyncCastOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize())) ||
      failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  if (getSourceSize() != getResultSize()) {
    return emitOpError("source and result sizes must be equal (tied storage)");
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.async.load
//===----------------------------------------------------------------------===//

LogicalResult AsyncLoadOp::verify() {
  AsyncLoadOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.async.store
//===----------------------------------------------------------------------===//

LogicalResult AsyncStoreOp::verify() {
  AsyncStoreOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getTarget(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

Value AsyncStoreOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getTarget());
}

::std::optional<unsigned>
AsyncStoreOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0}; // target
}

SmallVector<int64_t> AsyncStoreOp::getTiedResultOperandIndices() {
  return {0}; // target
}

//===----------------------------------------------------------------------===//
// stream.async.dispatch
//===----------------------------------------------------------------------===//

void AsyncDispatchOp::build(OpBuilder &builder, OperationState &state,
                            ExecutableExportOp exportOp, ValueRange workload,
                            TypeRange resultTypes, ValueRange operands,
                            ValueRange operandSizes, ValueRange operandOffsets,
                            ValueRange operandEnds, ValueRange operandLengths,
                            ValueRange resultSizes,
                            ArrayRef<int64_t> tiedOperands,
                            AffinityAttr affinityAttr) {
  StringRef executableOpSymName =
      exportOp->getParentOp()
          ->getAttrOfType<StringAttr>(SymbolTable::getSymbolAttrName())
          .getValue();
  auto entryPoint =
      SymbolRefAttr::get(builder.getContext(), executableOpSymName,
                         {SymbolRefAttr::get(exportOp)});
  build(builder, state, resultTypes, workload,
        builder.getArrayAttr({entryPoint}), operands, operandSizes,
        operandOffsets, operandEnds, operandLengths, resultSizes,
        cast<ArrayAttr>(builder.getIndexArrayAttr(tiedOperands)), affinityAttr);
}

static ParseResult parseDispatchOperands(
    OpAsmParser &parser,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceOperands,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceOffsets,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceEnds,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceLengths) {
  if (failed(parser.parseLParen())) {
    return failure();
  }
  // Handle the case of no operands specially.
  if (succeeded(parser.parseOptionalRParen())) {
    return success();
  }
  do {
    // All entries at least have an %operand.
    resourceOperands.emplace_back();
    if (failed(parser.parseOperand(resourceOperands.back()))) {
      return failure();
    }
    // Resources have a range.
    if (succeeded(parser.parseOptionalLSquare())) {
      resourceOffsets.emplace_back();
      resourceEnds.emplace_back();
      resourceLengths.emplace_back();
      if (failed(parser.parseOperand(resourceOffsets.back())) ||
          failed(parser.parseKeyword("to")) ||
          failed(parser.parseOperand(resourceEnds.back())) ||
          failed(parser.parseKeyword("for")) ||
          failed(parser.parseOperand(resourceLengths.back())) ||
          failed(parser.parseRSquare())) {
        return failure();
      }
    }
  } while (succeeded(parser.parseOptionalComma()));
  if (failed(parser.parseRParen())) {
    return failure();
  }
  return success();
}

static void printDispatchOperands(OpAsmPrinter &p, Operation *op,
                                  ValueRange resourceOperands,
                                  ValueRange resourceOffsets,
                                  ValueRange resourceEnds,
                                  ValueRange resourceLengths) {
  p << "(";
  unsigned resourceIndex = 0;
  llvm::interleaveComma(resourceOperands, p, [&](Value operand) {
    p.printOperand(operand);
    if (isa<IREE::Stream::ResourceType>(operand.getType())) {
      p << "[";
      p.printOperand(resourceOffsets[resourceIndex]);
      p << " to ";
      p.printOperand(resourceEnds[resourceIndex]);
      p << " for ";
      p.printOperand(resourceLengths[resourceIndex]);
      p << "]";
      ++resourceIndex;
    }
  });
  p << ")";
}

LogicalResult AsyncDispatchOp::verify() {
  AsyncDispatchOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getResourceOperands(),
                                op.getResourceOperandSizes())) ||
      failed(verifyOpValueSizes(op, op.getResults(), op.getResultSizes()))) {
    return failure();
  }
  unsigned requiredRangeCount = 0;
  for (auto value : op.getResourceOperands()) {
    if (isa<IREE::Stream::ResourceType>(value.getType())) {
      ++requiredRangeCount;
    }
  }
  unsigned presentRangeCount = op.getResourceOperandOffsets().size();
  if (op.getResourceOperandEnds().size() != presentRangeCount ||
      op.getResourceOperandLengths().size() != presentRangeCount) {
    return op->emitOpError() << "mismatch on resource range "
                                "offsets/ends/lengths; counts must match";
  }
  if (presentRangeCount != requiredRangeCount) {
    return op->emitOpError() << "expects " << requiredRangeCount
                             << " resource range operand sets but "
                             << presentRangeCount << " are present";
  }
  return success();
}

LogicalResult
AsyncDispatchOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  return verifyDispatchSymbolUses(getOperation(), getEntryPointsAttr(),
                                  getWorkload(), symbolTable);
}

std::pair<unsigned, unsigned> AsyncDispatchOp::getTiedOperandsIndexAndLength() {
  return getODSOperandIndexAndLength(1); // $operands
}

void AsyncDispatchOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  unsigned rangeIndex = 0;
  unsigned tiedOperandBase = getTiedOperandsIndexAndLength().first;
  for (auto [operandIndex, operand] : llvm::enumerate(getResourceOperands())) {
    if (!isa<IREE::Stream::ResourceType>(operand.getType())) {
      continue;
    }
    ResourceAccessBitfield access = ResourceAccessBitfield::Read;
    auto tiedResults = getOperandTiedResults(tiedOperandBase + operandIndex);
    if (!tiedResults.empty()) {
      access = access | ResourceAccessBitfield::Write;
    }
    Value start = getResourceOperandOffsets()[rangeIndex];
    Value end = getResourceOperandEnds()[rangeIndex];
    Value length = getResourceOperandLengths()[rangeIndex];
    ++rangeIndex;
    ranges.push_back({access, operand, start, end, length});
    for (auto result : tiedResults) {
      ranges.push_back({access, result, start, end, length});
    }
  }
  for (auto [i, result, resultSize] :
       llvm::zip_equal(llvm::seq<unsigned>(0, getResults().size()),
                       getResults(), getResultSizes())) {
    if (getTiedResultOperandIndex(i).has_value()) {
      // Already covered above.
      continue;
    }
    ranges.push_back({ResourceAccessBitfield::Write, result, Value{},
                      resultSize, resultSize});
  }
}

bool AsyncDispatchOp::preferCloneToConsumers() {
  // If the dispatch does not consume any resources then it is effectively a
  // slow splat and should be treated like one.
  const bool consumesAny = llvm::any_of(
      getResourceOperands(), +[](Value operand) {
        return isa<IREE::Stream::AffinityTypeInterface>(operand.getType());
      });
  return !consumesAny;
}

//===----------------------------------------------------------------------===//
// stream.async.func
//===----------------------------------------------------------------------===//

AsyncFuncOp AsyncFuncOp::create(Location location, StringRef name,
                                FunctionType type,
                                ArrayRef<int64_t> tiedOperands,
                                ArrayRef<DictionaryAttr> argAttrs,
                                ArrayRef<DictionaryAttr> resAttrs) {
  OpBuilder builder(location->getContext());
  OperationState state(location, getOperationName());
  AsyncFuncOp::build(builder, state, name, type,
                     builder.getIndexArrayAttr(tiedOperands), argAttrs,
                     resAttrs);
  return cast<AsyncFuncOp>(Operation::create(state));
}

void AsyncFuncOp::build(OpBuilder &builder, OperationState &state,
                        StringRef name, FunctionType type,
                        ArrayAttr tiedOperands,
                        ArrayRef<DictionaryAttr> argAttrs,
                        ArrayRef<DictionaryAttr> resAttrs) {
  state.addAttribute(SymbolTable::getSymbolAttrName(),
                     builder.getStringAttr(name));
  state.addAttribute(SymbolTable::getVisibilityAttrName(),
                     builder.getStringAttr("private"));
  state.addAttribute("function_type", TypeAttr::get(type));
  if (tiedOperands) {
    state.addAttribute(IREE::Util::TiedOpInterface::getStorageAttrName(),
                       tiedOperands);
  }
  state.addRegion();
  if (!argAttrs.empty() || !resAttrs.empty()) {
    assert(argAttrs.empty() || (type.getNumInputs() == argAttrs.size()));
    assert(resAttrs.empty() || (type.getNumResults() == resAttrs.size()));
    call_interface_impl::addArgAndResultAttrs(
        builder, state, argAttrs, resAttrs, builder.getStringAttr("arg_attrs"),
        builder.getStringAttr("res_attrs"));
  }
}

bool AsyncFuncOp::isResultTied(int resultIndex) {
  auto tiedOperandsAttr = getTiedOperandsAttr();
  if (!tiedOperandsAttr) {
    return false;
  }
  auto indexAttr = dyn_cast_if_present<IntegerAttr>(
      tiedOperandsAttr.getValue()[resultIndex]);
  if (!indexAttr) {
    return false;
  }
  return indexAttr.getInt() != IREE::Util::TiedOpInterface::kUntiedIndex;
}

//===----------------------------------------------------------------------===//
// stream.async.call
//===----------------------------------------------------------------------===//

LogicalResult AsyncCallOp::verify() {
  AsyncCallOp op = *this;

  if (failed(verifyOpValueSizes(op, op.getResourceOperands(),
                                op.getResourceOperandSizes())) ||
      failed(verifyOpValueSizes(op, op.getResults(), op.getResultSizes()))) {
    return failure();
  }

  unsigned requiredRangeCount = 0;
  for (auto value : op.getResourceOperands()) {
    if (isa<IREE::Stream::ResourceType>(value.getType())) {
      ++requiredRangeCount;
    }
  }

  unsigned presentRangeCount = op.getResourceOperandOffsets().size();
  if (op.getResourceOperandEnds().size() != presentRangeCount ||
      op.getResourceOperandLengths().size() != presentRangeCount) {
    return op->emitOpError() << "mismatch on resource range "
                                "offsets/ends/lengths; counts must match";
  }
  if (presentRangeCount != requiredRangeCount) {
    return op->emitOpError() << "expects " << requiredRangeCount
                             << " resource range operand sets but "
                             << presentRangeCount << " are present";
  }

  // TODO(benvanik): support non-resource returns.
  // This would require fixing stream.async.execute and stream.async.concurrent
  // to be able to return non-resource types as well and adjust partitioning
  // to set them up as return values. For now we just avoid this.
  for (auto resultType : op.getResultTypes()) {
    if (!isa<IREE::Stream::ResourceType>(resultType)) {
      return op->emitOpError() << "non-resource return values are not yet "
                                  "supported on async calls";
    }
  }

  return success();
}

LogicalResult
AsyncCallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  Operation *op = getOperation();
  auto calleeOp =
      symbolTable.lookupNearestSymbolFrom<IREE::Stream::AsyncFuncOp>(
          op, getCalleeAttr());
  if (!calleeOp) {
    return op->emitOpError() << "undefined external call: " << getCallee();
  }

  // NOTE: we allow the func to have broader lifetimes (`*`) than the calls.
  auto callType = getCalleeType();
  auto calleeType = calleeOp.getFunctionType();
  if (calleeType.getNumInputs() != callType.getNumInputs() ||
      calleeType.getNumResults() != callType.getNumResults()) {
    return emitOpError("function type mismatch; expected ")
           << callType << " but callee is " << calleeType;
  }
  // auto typesCompatible = [](Type actual, Type expected) {
  auto typesCompatible = [](Type callee, Type call) {
    if (callee == call) {
      return true;
    }
    auto calleeResource = dyn_cast<IREE::Stream::ResourceType>(callee);
    auto callResource = dyn_cast<IREE::Stream::ResourceType>(call);
    if (calleeResource && callResource) {
      if (calleeResource.getLifetime() == IREE::Stream::Lifetime::Unknown) {
        // Allow anything to match with an unknown lifetime on the async.func.
        return true;
      }
      // Lifetime is specified on the func and the call doesn't match.
      return false;
    }
    return false;
  };
  for (auto [calleeArg, expectedArg] :
       llvm::zip_equal(calleeType.getInputs(), callType.getInputs())) {
    if (!typesCompatible(calleeArg, expectedArg)) {
      return emitOpError("function argument type mismatch; expected ")
             << expectedArg << " but callee provides " << calleeArg;
    }
  }
  for (auto [calleeResult, expectedResult] :
       llvm::zip_equal(calleeType.getResults(), callType.getResults())) {
    if (!typesCompatible(calleeResult, expectedResult)) {
      return emitOpError("function result type mismatch; expected ")
             << expectedResult << " but callee provides " << calleeResult;
    }
  }

  return success();
}

FunctionType AsyncCallOp::getCalleeType() {
  auto operandTypes = llvm::map_to_vector(
      getArgOperands(), [](Value arg) { return arg.getType(); });
  return FunctionType::get(getContext(), operandTypes, getResultTypes());
}

std::pair<unsigned, unsigned> AsyncCallOp::getTiedOperandsIndexAndLength() {
  return getODSOperandIndexAndLength(0); // $operands
}

void AsyncCallOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  unsigned rangeIndex = 0;
  unsigned tiedOperandBase = getTiedOperandsIndexAndLength().first;
  for (auto [operandIndex, operand] : llvm::enumerate(getResourceOperands())) {
    if (!isa<IREE::Stream::ResourceType>(operand.getType())) {
      continue;
    }
    ResourceAccessBitfield access = ResourceAccessBitfield::Read;
    auto tiedResults = getOperandTiedResults(tiedOperandBase + operandIndex);
    if (!tiedResults.empty()) {
      access = access | ResourceAccessBitfield::Write;
    }
    Value start = getResourceOperandOffsets()[rangeIndex];
    Value end = getResourceOperandEnds()[rangeIndex];
    Value length = getResourceOperandLengths()[rangeIndex];
    ++rangeIndex;
    ranges.push_back({access, operand, start, end, length});
    for (auto result : tiedResults) {
      ranges.push_back({access, result, start, end, length});
    }
  }
  for (auto [i, result, resultSize] :
       llvm::zip_equal(llvm::seq<unsigned>(0, getResults().size()),
                       getResults(), getResultSizes())) {
    if (getTiedResultOperandIndex(i).has_value()) {
      // Already covered above.
      continue;
    }
    ranges.push_back({ResourceAccessBitfield::Write, result, Value{},
                      resultSize, resultSize});
  }
}

//===----------------------------------------------------------------------===//
// stream.async.execute
//===----------------------------------------------------------------------===//

void AsyncExecuteOp::build(OpBuilder &builder, OperationState &state,
                           TypeRange resultTypes, ValueRange resultSizes,
                           Value awaitTimepoint, ValueRange operands,
                           ValueRange operandSizes,
                           ArrayRef<int64_t> tiedOperands,
                           ArrayRef<NamedAttribute> attributes) {
  state.addTypes(resultTypes);
  state.addTypes(IREE::Stream::TimepointType::get(builder.getContext()));
  state.addOperands(operands);
  state.addOperands(operandSizes);
  state.addOperands(resultSizes);
  if (awaitTimepoint) {
    state.addOperands(awaitTimepoint);
  }
  state.addAttributes(attributes);
  state.attributes.erase(IREE::Util::TiedOpInterface::getStorageAttrName());
  state.addAttribute(IREE::Util::TiedOpInterface::getStorageAttrName(),
                     builder.getIndexArrayAttr(tiedOperands));
  state.attributes.erase(getOperandSegmentSizeAttr());
  state.addAttribute(getOperandSegmentSizeAttr(),
                     builder.getDenseI32ArrayAttr({
                         static_cast<int32_t>(operands.size()),
                         static_cast<int32_t>(operandSizes.size()),
                         static_cast<int32_t>(resultSizes.size()),
                         awaitTimepoint ? 1 : 0,
                     }));
  state.addRegion();
}

LogicalResult AsyncExecuteOp::verify() {
  AsyncExecuteOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getResourceOperands(),
                                op.getResourceOperandSizes())) ||
      failed(verifyOpValueSizes(op, op.getResults(), op.getResultSizes()))) {
    return failure();
  }
  if (failed(verifyAllResourcesCaptured(op.getBody())) ||
      failed(verifyEscapingResources(op.getBody(), op.getResults(),
                                     op.getResultSizes()))) {
    return failure();
  }
  return success();
}

std::pair<unsigned, unsigned> AsyncExecuteOp::getTiedResultsIndexAndLength() {
  return {0, getResults().size()};
}

OperandRange
AsyncExecuteOp::getEntrySuccessorOperands(RegionSuccessor successor) {
  assert(successor.getSuccessor() == &getBody() && "invalid region index");
  return getResourceOperands();
}

ValueRange AsyncExecuteOp::getSuccessorInputs(RegionSuccessor successor) {
  return successor.isOperation() ? ValueRange(getResults())
                                 : ValueRange(getBodyRegion().getArguments());
}

void AsyncExecuteOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  // Unconditional control flow into the region and back to the parent, so
  // return the correct RegionSuccessor purely based on the index being None or
  // 0.
  if (!point.isParent()) {
    regions.push_back(RegionSuccessor(getOperation()));
  } else {
    regions.push_back(RegionSuccessor(&getBody()));
  }
}

// Gets the async access ranges for the generic stream execution op capturing
// resources.
template <typename Op>
static void
getExecutionAsyncAccessRanges(Op op,
                              SmallVectorImpl<AsyncAccessRange> &ranges) {
  unsigned tiedOperandBase = op.getTiedOperandsIndexAndLength().first;
  for (auto [i, operand, operandSize] : llvm::zip_equal(
           llvm::seq<unsigned>(0, op.getResourceOperands().size()),
           op.getResourceOperands(), op.getResourceOperandSizes())) {
    if (!isa<IREE::Stream::ResourceType>(operand.getType())) {
      continue;
    }
    ResourceAccessBitfield access = ResourceAccessBitfield::Read;
    auto tiedResults = op.getOperandTiedResults(tiedOperandBase + i);
    if (!tiedResults.empty()) {
      access = access | ResourceAccessBitfield::Write;
    }
    ranges.push_back({access, operand, Value{}, operandSize, operandSize});
    for (auto result : tiedResults) {
      ranges.push_back({access, result, Value{}, operandSize, operandSize});
    }
  }
  for (auto [i, result, resultSize] :
       llvm::zip_equal(llvm::seq<unsigned>(0, op.getResults().size()),
                       op.getResults(), op.getResultSizes())) {
    if (op.getTiedResultOperandIndex(i).has_value()) {
      // Already covered above.
      continue;
    }
    ranges.push_back({ResourceAccessBitfield::Write, result, Value{},
                      resultSize, resultSize});
  }
}

void AsyncExecuteOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  getExecutionAsyncAccessRanges(*this, ranges);
}

Operation::operand_range AsyncExecuteOp::getClosureOperands() {
  return getResourceOperands();
}

Operation::result_range AsyncExecuteOp::getClosureResults() {
  return getResults();
}

bool AsyncExecuteOp::canClosureContainOp(Operation *op) { return false; }

IREE::Util::ValueAccess
AsyncExecuteOp::getOperandAccess(unsigned operandIndex) {
  auto arg = getBody().getArgument(operandIndex);
  return computeValueAccess(arg);
}

IREE::Util::ValueAccess AsyncExecuteOp::getResultAccess(unsigned resultIndex) {
  auto yieldOp = cast<YieldOp>(getBody().getBlocks().front().getTerminator());
  return computeValueAccess(yieldOp.getOperand(resultIndex));
}

IREE::Util::ClosureOpInterface
AsyncExecuteOp::cloneReplacementExcludingOperandsAndResults(
    ArrayRef<unsigned> excludedOperandIndices,
    ArrayRef<unsigned> excludedResultIndices, PatternRewriter &rewriter) {
  auto newResultTypes = llvm::map_to_vector(
      getResults(), [](auto value) { return value.getType(); });
  auto newResultSizes = llvm::to_vector(getResultSizes());
  auto newOperandsValues = llvm::to_vector(getResourceOperands());
  auto newOperandSizes = llvm::to_vector(getResourceOperandSizes());
  IREE::Util::excludeClosureOperandsAndResults(
      newOperandsValues, newOperandSizes, excludedOperandIndices,
      newResultTypes, newResultSizes, excludedResultIndices);

  auto newTiedOperandIndices = llvm::to_vector(getTiedResultOperandIndices());
  IREE::Util::excludeTiedOperandAndResultIndices(
      excludedOperandIndices, excludedResultIndices, newTiedOperandIndices);
  assert(getTiedOperandsIndexAndLength().first == 0 &&
         "operands must be the first ODS group");

  auto newOp = AsyncExecuteOp::create(
      rewriter, getLoc(), newResultTypes, newResultSizes, getAwaitTimepoint(),
      newOperandsValues, newOperandSizes, newTiedOperandIndices,
      getOperation()->getAttrs());
  auto &newBody = newOp.getClosureBodyRegion();
  newBody.takeBody(getClosureBodyRegion());
  eraseStreamRegionResults(newBody, excludedResultIndices);

  auto &block = newBody.front();
  BitVector eraseIndices(block.getNumArguments());
  for (auto i : excludedOperandIndices) {
    eraseIndices.set(i);
  }
  block.eraseArguments(eraseIndices);
  return newOp;
}

//===----------------------------------------------------------------------===//
// stream.async.concurrent
//===----------------------------------------------------------------------===//

void AsyncConcurrentOp::build(OpBuilder &builder, OperationState &state,
                              TypeRange resultTypes, ValueRange resultSizes,
                              ValueRange operands, ValueRange operandSizes,
                              ArrayRef<int64_t> tiedOperands,
                              ArrayRef<NamedAttribute> attributes) {
  state.addTypes(resultTypes);
  state.addOperands(operands);
  state.addOperands(operandSizes);
  state.addOperands(resultSizes);
  state.addAttributes(attributes);
  state.attributes.erase(IREE::Util::TiedOpInterface::getStorageAttrName());
  state.addAttribute(IREE::Util::TiedOpInterface::getStorageAttrName(),
                     builder.getIndexArrayAttr(tiedOperands));
  state.attributes.erase(getOperandSegmentSizeAttr());
  state.addAttribute(getOperandSegmentSizeAttr(),
                     builder.getDenseI32ArrayAttr({
                         static_cast<int32_t>(operands.size()),
                         static_cast<int32_t>(operandSizes.size()),
                         static_cast<int32_t>(resultSizes.size()),
                     }));
  state.addRegion();
}

LogicalResult AsyncConcurrentOp::verify() {
  AsyncConcurrentOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getResourceOperands(),
                                op.getResourceOperandSizes())) ||
      failed(verifyOpValueSizes(op, op.getResults(), op.getResultSizes()))) {
    return failure();
  }
  if (failed(verifyAllResourcesCaptured(op.getBody())) ||
      failed(verifyEscapingResources(op.getBody(), op.getResults(),
                                     op.getResultSizes()))) {
    return failure();
  }
  return success();
}

OperandRange
AsyncConcurrentOp::getEntrySuccessorOperands(RegionSuccessor successor) {
  assert(successor.getSuccessor() == &getBody() && "invalid region index");
  return getResourceOperands();
}

ValueRange AsyncConcurrentOp::getSuccessorInputs(RegionSuccessor successor) {
  return successor.isOperation() ? ValueRange(getResults())
                                 : ValueRange(getBodyRegion().getArguments());
}

void AsyncConcurrentOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  // Unconditional control flow into the region and back to the parent, so
  // return the correct RegionSuccessor purely based on the index being None or
  // 0.
  if (!point.isParent()) {
    regions.push_back(RegionSuccessor(getOperation()));
  } else {
    regions.push_back(RegionSuccessor(&getBody()));
  }
}

void AsyncConcurrentOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  getExecutionAsyncAccessRanges(*this, ranges);
}

Operation::operand_range AsyncConcurrentOp::getClosureOperands() {
  return getResourceOperands();
}

Operation::result_range AsyncConcurrentOp::getClosureResults() {
  return getResults();
}

bool AsyncConcurrentOp::canClosureContainOp(Operation *op) { return false; }

IREE::Util::ValueAccess
AsyncConcurrentOp::getOperandAccess(unsigned operandIndex) {
  auto arg = getBody().getArgument(operandIndex);
  return computeValueAccess(arg);
}

IREE::Util::ValueAccess
AsyncConcurrentOp::getResultAccess(unsigned resultIndex) {
  auto yieldOp = cast<YieldOp>(getBody().getBlocks().front().getTerminator());
  return computeValueAccess(yieldOp.getOperand(resultIndex));
}

IREE::Util::ClosureOpInterface
AsyncConcurrentOp::cloneReplacementExcludingOperandsAndResults(
    ArrayRef<unsigned> excludedOperandIndices,
    ArrayRef<unsigned> excludedResultIndices, PatternRewriter &rewriter) {
  auto newResultTypes = llvm::to_vector(getResultTypes());
  auto newResultSizes = llvm::to_vector(getResultSizes());
  auto newOperandsValues = llvm::to_vector(getResourceOperands());
  auto newOperandSizes = llvm::to_vector(getResourceOperandSizes());
  IREE::Util::excludeClosureOperandsAndResults(
      newOperandsValues, newOperandSizes, excludedOperandIndices,
      newResultTypes, newResultSizes, excludedResultIndices);

  auto newTiedOperandIndices = llvm::to_vector(getTiedResultOperandIndices());
  IREE::Util::excludeTiedOperandAndResultIndices(
      excludedOperandIndices, excludedResultIndices, newTiedOperandIndices);
  assert(getTiedOperandsIndexAndLength().first == 0 &&
         "operands must be the first ODS group");

  auto newOp = AsyncConcurrentOp::create(
      rewriter, getLoc(), newResultTypes, newResultSizes, newOperandsValues,
      newOperandSizes, newTiedOperandIndices, getOperation()->getAttrs());
  auto &newBody = newOp.getClosureBodyRegion();
  newBody.takeBody(getClosureBodyRegion());
  eraseStreamRegionResults(newBody, excludedResultIndices);
  auto &block = newBody.front();
  BitVector eraseIndices(block.getNumArguments());
  for (auto i : excludedOperandIndices) {
    eraseIndices.set(i);
  }
  block.eraseArguments(eraseIndices);
  return newOp;
}

//===----------------------------------------------------------------------===//
// stream.async.parameter.load
//===----------------------------------------------------------------------===//

LogicalResult AsyncParameterLoadOp::verify() {
  AsyncParameterLoadOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getResult(), op.getResultSize()))) {
    return failure();
  }
  return success();
}

void AsyncParameterLoadOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  ranges.push_back({ResourceAccessBitfield::Write, getResult(), Value{},
                    getResultSize(), getResultSize()});
}

//===----------------------------------------------------------------------===//
// stream.async.parameter.read
//===----------------------------------------------------------------------===//

LogicalResult AsyncParameterReadOp::verify() {
  AsyncParameterReadOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getTarget(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

Value AsyncParameterReadOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getTarget());
}

::std::optional<unsigned>
AsyncParameterReadOp::getTiedResultOperandIndex(unsigned resultIndex) {
  if (resultIndex == 0) {
    return {0}; // result tied to target
  }
  return std::nullopt; // result_timepoint not tied
}

SmallVector<int64_t> AsyncParameterReadOp::getTiedResultOperandIndices() {
  return {
      0,
      IREE::Util::TiedOpInterface::kUntiedIndex}; // result tied to target,
                                                  // result_timepoint not tied
}

void AsyncParameterReadOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  ranges.push_back({ResourceAccessBitfield::Write, getTarget(),
                    getTargetOffset(), getTargetEnd(), getTargetLength()});
  ranges.push_back({ResourceAccessBitfield::Write, getResult(),
                    getTargetOffset(), getTargetEnd(), getTargetLength()});
}

//===----------------------------------------------------------------------===//
// stream.async.parameter.write
//===----------------------------------------------------------------------===//

LogicalResult AsyncParameterWriteOp::verify() {
  AsyncParameterWriteOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize()))) {
    return failure();
  }
  return success();
}

Value AsyncParameterWriteOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getSource());
}

::std::optional<unsigned>
AsyncParameterWriteOp::getTiedResultOperandIndex(unsigned resultIndex) {
  if (resultIndex == 0) {
    return {0}; // result tied to source
  }
  return std::nullopt; // result_timepoint not tied
}

SmallVector<int64_t> AsyncParameterWriteOp::getTiedResultOperandIndices() {
  return {
      0,
      IREE::Util::TiedOpInterface::kUntiedIndex}; // result tied to source,
                                                  // result_timepoint not tied
}

void AsyncParameterWriteOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  ranges.push_back({ResourceAccessBitfield::Read, getSource(),
                    getSourceOffset(), getSourceEnd(), getSourceLength()});
  ranges.push_back({ResourceAccessBitfield::Read, getResult(),
                    getSourceOffset(), getSourceEnd(), getSourceLength()});
}

//===----------------------------------------------------------------------===//
// stream.async.parameter.gather
//===----------------------------------------------------------------------===//

LogicalResult AsyncParameterGatherOp::inferReturnTypes(
    MLIRContext *context, std::optional<Location> location, ValueRange operands,
    DictionaryAttr attributes, PropertyRef properties, RegionRange regions,
    SmallVectorImpl<Type> &inferredReturnTypes) {
  // The result type matches the target operand type.
  AsyncParameterGatherOpAdaptor adaptor(operands, attributes, properties,
                                        regions);
  inferredReturnTypes.push_back(adaptor.getTarget().getType());
  inferredReturnTypes.push_back(IREE::Stream::TimepointType::get(context));
  return success();
}

LogicalResult AsyncParameterGatherOp::verify() {
  AsyncParameterGatherOp op = *this;
  size_t expectedCount = op.getSourceKeys().size();
  if (op.getSourceOffsets().size() != expectedCount ||
      op.getTargetOffsets().size() != expectedCount ||
      op.getTargetEnds().size() != expectedCount ||
      op.getTargetLengths().size() != expectedCount) {
    return op.emitOpError()
           << "requires that the source keys, source offsets, target offsets, "
              "target ends, and lengths are all 1:1";
  }
  if (failed(verifyOpValueSizes(op, op.getTarget(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

Value AsyncParameterGatherOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getTarget());
}

::std::optional<unsigned>
AsyncParameterGatherOp::getTiedResultOperandIndex(unsigned resultIndex) {
  if (resultIndex == 0) {
    return {
        getSourceOffsets()
            .size()}; // result tied to target (after variadic source_offsets)
  }
  return std::nullopt; // result_timepoint not tied
}

SmallVector<int64_t> AsyncParameterGatherOp::getTiedResultOperandIndices() {
  return {
      static_cast<int64_t>(getSourceOffsets().size()),
      IREE::Util::TiedOpInterface::kUntiedIndex}; // result tied to target,
                                                  // result_timepoint not tied
}

void AsyncParameterGatherOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  // Gather writes to the entire target resource across all gather operations.
  // We conservatively mark the entire target as written. We could return
  // the exact ranges (adding more) but I don't currently know the performance
  // implications of that (this may be part of an n^2 algorithm).
  ranges.push_back({ResourceAccessBitfield::Write, getTarget(), Value{},
                    getTargetSize(), getTargetSize()});
  ranges.push_back({ResourceAccessBitfield::Write, getResult(), Value{},
                    getTargetSize(), getTargetSize()});
}

//===----------------------------------------------------------------------===//
// stream.async.parameter.scatter
//===----------------------------------------------------------------------===//

LogicalResult AsyncParameterScatterOp::inferReturnTypes(
    MLIRContext *context, std::optional<Location> location, ValueRange operands,
    DictionaryAttr attributes, PropertyRef properties, RegionRange regions,
    SmallVectorImpl<Type> &inferredReturnTypes) {
  // The result type matches the source operand type.
  AsyncParameterScatterOpAdaptor adaptor(operands, attributes, properties,
                                         regions);
  inferredReturnTypes.push_back(adaptor.getSource().getType());
  inferredReturnTypes.push_back(IREE::Stream::TimepointType::get(context));
  return success();
}

LogicalResult AsyncParameterScatterOp::verify() {
  AsyncParameterScatterOp op = *this;
  size_t expectedCount = op.getTargetKeys().size();
  if (op.getSourceOffsets().size() != expectedCount ||
      op.getSourceEnds().size() != expectedCount ||
      op.getSourceLengths().size() != expectedCount ||
      op.getTargetOffsets().size() != expectedCount) {
    return op.emitOpError()
           << "requires that the source offsets, source ends, lengths, "
              "target keys, and target offsets are all 1:1";
  }
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize()))) {
    return failure();
  }
  return success();
}

Value AsyncParameterScatterOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getSource());
}

::std::optional<unsigned>
AsyncParameterScatterOp::getTiedResultOperandIndex(unsigned resultIndex) {
  if (resultIndex == 0) {
    return {0}; // result tied to source
  }
  return std::nullopt; // result_timepoint not tied
}

SmallVector<int64_t> AsyncParameterScatterOp::getTiedResultOperandIndices() {
  return {
      0,
      IREE::Util::TiedOpInterface::kUntiedIndex}; // result tied to source,
                                                  // result_timepoint not tied
}

void AsyncParameterScatterOp::getAsyncAccessRanges(
    SmallVectorImpl<AsyncAccessRange> &ranges) {
  // Scatter reads from the entire source resource across all scatter
  // operations. We conservatively mark the entire source as read. We could
  // return the exact ranges (adding more) but I don't currently know the
  // performance implications of that (this may be part of an n^2 algorithm).
  ranges.push_back({ResourceAccessBitfield::Read, getSource(), Value{},
                    getSourceSize(), getSourceSize()});
  ranges.push_back({ResourceAccessBitfield::Read, getResult(), Value{},
                    getSourceSize(), getSourceSize()});
}

//===----------------------------------------------------------------------===//
// stream.cmd.flush
//===----------------------------------------------------------------------===//

LogicalResult CmdFlushOp::verify() {
  CmdFlushOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getTarget(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.cmd.invalidate
//===----------------------------------------------------------------------===//

LogicalResult CmdInvalidateOp::verify() {
  CmdInvalidateOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getTarget(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.cmd.discard
//===----------------------------------------------------------------------===//

LogicalResult CmdDiscardOp::verify() {
  CmdDiscardOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getTarget(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.cmd.fill
//===----------------------------------------------------------------------===//

LogicalResult CmdFillOp::verify() {
  CmdFillOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getTarget(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.cmd.copy
//===----------------------------------------------------------------------===//

LogicalResult CmdCopyOp::verify() {
  CmdCopyOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize())) ||
      failed(verifyOpValueSizes(op, op.getTarget(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.cmd.collective
//===----------------------------------------------------------------------===//

LogicalResult CmdCollectiveOp::verify() {
  CmdCollectiveOp op = *this;
  size_t resourceCount = op.getResources().size();
  if (op.getResourceSizes().size() != resourceCount ||
      op.getResourceOffsets().size() != resourceCount ||
      op.getResourceLengths().size() != resourceCount ||
      op.getResourceAccesses().size() != resourceCount) {
    return op->emitOpError() << "with " << resourceCount
                             << " resources has mismatched associated ranges";
  }

  size_t requiredCount = 0;
  IREE::Stream::ResourceAccessBitfield requiredAccess[2] = {
      IREE::Stream::ResourceAccessBitfield::None,
      IREE::Stream::ResourceAccessBitfield::None,
  };
  switch (getOp().getKind()) {
  default:
    requiredCount = 2; // send & recv
    requiredAccess[0] = IREE::Stream::ResourceAccessBitfield::Read;
    requiredAccess[1] = IREE::Stream::ResourceAccessBitfield::Write;
    break;
  case IREE::Stream::CollectiveKind::Send:
    requiredCount = 1; // send
    requiredAccess[0] = IREE::Stream::ResourceAccessBitfield::Read;
    break;
  case IREE::Stream::CollectiveKind::Recv:
    requiredCount = 1; // recv
    requiredAccess[0] = IREE::Stream::ResourceAccessBitfield::Write;
    break;
  }
  if (resourceCount != requiredCount) {
    return op->emitOpError()
           << "requires " << requiredCount << " resources but " << resourceCount
           << " provided";
  }
  for (size_t i = 0; i < requiredCount; ++i) {
    auto declaredAccess = cast<IREE::Stream::ResourceAccessBitfieldAttr>(
                              op.getResourceAccesses()[i])
                              .getValue();
    if (!bitEnumContainsAll(declaredAccess, requiredAccess[i])) {
      return op->emitOpError()
             << "resource " << i << " requires access "
             << stringifyEnum(requiredAccess[i]) << " but is declared as "
             << stringifyEnum(declaredAccess);
    }
  }

  return success();
}

//===----------------------------------------------------------------------===//
// stream.cmd.dispatch
//===----------------------------------------------------------------------===//

LogicalResult CmdDispatchOp::verify() {
  CmdDispatchOp op = *this;
  size_t resourceCount = op.getResources().size();
  if (op.getResourceSizes().size() != resourceCount ||
      op.getResourceOffsets().size() != resourceCount ||
      op.getResourceLengths().size() != resourceCount ||
      op.getResourceAccesses().size() != resourceCount) {
    return op->emitOpError() << "dispatch with " << resourceCount
                             << " resources has mismatched associated ranges";
  }
  return success();
}

LogicalResult
CmdDispatchOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  Operation *op = getOperation();
  auto entryPointRefs = getEntryPointRefs();
  if (entryPointRefs.empty()) {
    return emitOpError() << "at least one entry point must be defined";
  }
  for (auto entryPointAttr : entryPointRefs) {
    auto exportOp =
        symbolTable.lookupNearestSymbolFrom<IREE::Stream::ExecutableExportOp>(
            op, entryPointAttr);
    if (!exportOp) {
      // TODO(benvanik): there are a lot of tests that are assuming this is not
      // verified. We'll need to go add dummy executables for all of them. Today
      // we just bail on the verifier if the symbol isn't found.
      //
      // Should be:
      //   return op->emitOpError() << "undefined entry point: " <<
      //   entry_point();
      return success();
    }

    // Verify that the workload parameters captured match the target export.
    if (failed(verifyDispatchWorkload(op, exportOp, getWorkload()))) {
      return failure();
    }

    // Verify that the exported function, if present, matches the signature of
    // the dispatch.
    auto funcOp = exportOp.lookupFunctionRef();
    if (!funcOp) {
      return success();
    }

    TypeRange uniformTypes = getUniformOperands().getTypes();
    int64_t numResources = getResources().size();

    SmallVector<Type> uniformEntryPointTypes;
    int64_t bindingCounts = 0;
    for (auto entryPointArg : funcOp.getArgumentTypes()) {
      if (isa<IREE::Stream::BindingType>(entryPointArg)) {
        bindingCounts++;
      } else {
        uniformEntryPointTypes.push_back(entryPointArg);
      }
    }
    if (uniformTypes.size() != uniformEntryPointTypes.size()) {
      return emitOpError("function type mismatch; expected ")
             << uniformTypes.size()
             << " uniform arguments on exported function, but has "
             << uniformEntryPointTypes.size();
    }
    if (numResources != bindingCounts) {
      return emitOpError("function type mismatch; expected ")
             << numResources
             << " binding arguments on exported function, but has "
             << bindingCounts;
    }
    for (auto [expectedType, actualType] :
         llvm::zip_equal(uniformTypes, uniformEntryPointTypes)) {
      if (expectedType != actualType) {
        return emitOpError("uniform dispatch argument type mismatch: expected ")
               << expectedType << " but got " << actualType;
      }
    }
  }
  return success();
}

static ParseResult parseDispatchResources(
    OpAsmParser &parser,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resources,
    SmallVectorImpl<Type> &resourceTypes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceSizes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceOffsets,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceLengths,
    ArrayAttr &resourceAccesses) {
  SmallVector<Attribute> accessAttrs;
  do {
    // Reserve entries in the lists.
    resources.emplace_back();
    resourceTypes.emplace_back();
    resourceSizes.emplace_back();
    resourceOffsets.emplace_back();
    resourceLengths.emplace_back();
    StringRef accessStr;
    if (failed(parser.parseKeyword(&accessStr)) ||
        failed(parser.parseOperand(resources.back())) ||
        failed(parser.parseLSquare()) ||
        failed(parser.parseOperand(resourceOffsets.back())) ||
        failed(parser.parseKeyword("for")) ||
        failed(parser.parseOperand(resourceLengths.back())) ||
        failed(parser.parseRSquare()) || failed(parser.parseColon()) ||
        failed(parseSizeAwareType(parser, resourceTypes.back(),
                                  resourceSizes.back()))) {
      return failure();
    }
    IREE::Stream::ResourceAccessBitfield accessBits =
        IREE::Stream::ResourceAccessBitfield::None;
    if (accessStr == "ro") {
      accessBits = IREE::Stream::ResourceAccessBitfield::Read;
    } else if (accessStr == "wo") {
      accessBits = IREE::Stream::ResourceAccessBitfield::Write;
    } else if (accessStr == "rw") {
      accessBits = IREE::Stream::ResourceAccessBitfield::Read |
                   IREE::Stream::ResourceAccessBitfield::Write;
    }
    accessAttrs.push_back(IREE::Stream::ResourceAccessBitfieldAttr::get(
        parser.getBuilder().getContext(), accessBits));
  } while (succeeded(parser.parseOptionalComma()));
  resourceAccesses = parser.getBuilder().getArrayAttr(accessAttrs);
  return success();
}

static void
printDispatchResources(OpAsmPrinter &p, Operation *op, ValueRange resources,
                       TypeRange resourceTypes, ValueRange resourceSizes,
                       ValueRange resourceOffsets, ValueRange resourceLengths,
                       ArrayAttr resourceAccesses) {
  for (size_t i = 0; i < resources.size(); ++i) {
    auto resource = resources[i];
    auto resourceType = resourceTypes[i];
    auto resourceSize = resourceSizes[i];
    auto resourceOffset = resourceOffsets[i];
    auto resourceLength = resourceLengths[i];
    auto resourceAccess =
        cast<IREE::Stream::ResourceAccessBitfieldAttr>(resourceAccesses[i])
            .getValue();
    p.printNewline();
    p << "  ";
    if (bitEnumContainsAll(resourceAccess,
                           IREE::Stream::ResourceAccessBitfield::Read |
                               IREE::Stream::ResourceAccessBitfield::Write)) {
      p << "rw";
    } else if (bitEnumContainsAll(resourceAccess,
                                  IREE::Stream::ResourceAccessBitfield::Read)) {
      p << "ro";
    } else if (bitEnumContainsAll(
                   resourceAccess,
                   IREE::Stream::ResourceAccessBitfield::Write)) {
      p << "wo";
    }
    p << ' ';
    p.printOperand(resource);
    p << "[";
    p.printOperand(resourceOffset);
    p << " for ";
    p.printOperand(resourceLength);
    p << "] : ";
    printSizeAwareType(p, op, resourceType, resourceSize);
    if (i < resources.size() - 1) {
      p << ",";
    }
  }
}

// This is sloppy because the function has interleaved bindings and operands;
// if we had our own op we could just reuse the map we have for operands.
// static
SmallVector<unsigned>
CmdDispatchOp::makeOperandToArgMap(mlir::FunctionOpInterface funcOp) {
  unsigned operandCount =
      llvm::count_if(funcOp.getArgumentTypes(), [](Type type) {
        return !isa<IREE::Stream::BindingType>(type);
      });
  SmallVector<unsigned> map(operandCount);
  unsigned operandIdx = 0;
  for (auto it : llvm::enumerate(funcOp.getArgumentTypes())) {
    unsigned argIdx = it.index();
    auto argType = it.value();
    if (!isa<IREE::Stream::BindingType>(argType)) {
      map[operandIdx++] = argIdx;
    }
  }
  return map;
}

// static
SmallVector<unsigned>
CmdDispatchOp::makeResourceToArgMap(mlir::FunctionOpInterface funcOp) {
  unsigned operandCount = llvm::count_if(
      funcOp.getArgumentTypes(), llvm::IsaPred<IREE::Stream::BindingType>);
  SmallVector<unsigned> map(operandCount);
  size_t operandIdx = 0;
  for (auto [argIdx, argType] : llvm::enumerate(funcOp.getArgumentTypes())) {
    if (isa<IREE::Stream::BindingType>(argType)) {
      map[operandIdx++] = argIdx;
    }
  }
  return map;
}

//===----------------------------------------------------------------------===//
// stream.cmd.func
//===----------------------------------------------------------------------===//

CmdFuncOp CmdFuncOp::create(Location location, StringRef name,
                            FunctionType type,
                            ArrayRef<DictionaryAttr> argAttrs,
                            ArrayRef<DictionaryAttr> resAttrs) {
  OpBuilder builder(location->getContext());
  OperationState state(location, getOperationName());
  CmdFuncOp::build(builder, state, name, type, argAttrs, resAttrs);
  return cast<CmdFuncOp>(Operation::create(state));
}

void CmdFuncOp::build(OpBuilder &builder, OperationState &state, StringRef name,
                      FunctionType type, ArrayRef<DictionaryAttr> argAttrs,
                      ArrayRef<DictionaryAttr> resAttrs) {
  state.addAttribute(SymbolTable::getSymbolAttrName(),
                     builder.getStringAttr(name));
  state.addAttribute(SymbolTable::getVisibilityAttrName(),
                     builder.getStringAttr("private"));
  state.addAttribute("function_type", TypeAttr::get(type));
  state.addRegion();
  if (!argAttrs.empty() || !resAttrs.empty()) {
    assert(type.getNumInputs() == argAttrs.size());
    assert(type.getNumResults() == resAttrs.size());
    call_interface_impl::addArgAndResultAttrs(
        builder, state, argAttrs, resAttrs, builder.getStringAttr("arg_attrs"),
        builder.getStringAttr("res_attrs"));
  }
}

//===----------------------------------------------------------------------===//
// custom<DispatchFunctionSignature>
//===----------------------------------------------------------------------===//
// (%arg0: type {some.attr = 54 : index}, %arg1: type) -> (type, %arg1 as type)
// (%arg0[%arg1 for %arg2]: !stream.resource<*>, ...)

static ParseResult parseDispatchFunctionArgumentList(
    OpAsmParser &parser, SmallVectorImpl<OpAsmParser::UnresolvedOperand> &args,
    SmallVectorImpl<Type> &types, ArrayAttr &attrs) {
  auto indexType = parser.getBuilder().getIndexType();
  auto emptyDictionaryAttr = parser.getBuilder().getDictionaryAttr({});
  SmallVector<Attribute> argAttrsVec;
  do {
    OpAsmParser::UnresolvedOperand arg;
    if (failed(parser.parseOperand(arg))) {
      return failure();
    }
    bool hasOffsetLength = false;
    OpAsmParser::UnresolvedOperand offsetArg;
    OpAsmParser::UnresolvedOperand lengthArg;
    if (succeeded(parser.parseOptionalLSquare())) {
      // %offset for %length]
      if (failed(parser.parseOperand(offsetArg)) ||
          failed(parser.parseKeyword("for")) ||
          failed(parser.parseOperand(lengthArg)) ||
          failed(parser.parseRSquare())) {
        return failure();
      }
      hasOffsetLength = true;
    }
    Type type;
    NamedAttrList attrsVec;
    if (failed(parser.parseColonType(type)) ||
        failed(parser.parseOptionalAttrDict(attrsVec))) {
      return failure();
    }
    args.push_back(arg);
    types.push_back(type);
    argAttrsVec.push_back(parser.getBuilder().getDictionaryAttr(attrsVec));
    if (hasOffsetLength) {
      args.push_back(offsetArg);
      args.push_back(lengthArg);
      types.push_back(indexType);
      types.push_back(indexType);
      argAttrsVec.push_back(emptyDictionaryAttr);
      argAttrsVec.push_back(emptyDictionaryAttr);
    }
  } while (succeeded(parser.parseOptionalComma()));
  if (!argAttrsVec.empty()) {
    attrs = parser.getBuilder().getArrayAttr(argAttrsVec);
  }
  return success();
}

static ParseResult
parseDispatchFunctionResultList(OpAsmParser &parser,
                                SmallVectorImpl<Type> &resultTypes,
                                ArrayAttr &resultAttrs) {
  SmallVector<Attribute> resultAttrsVec;
  SmallVector<int64_t> tiedOperandIndices;
  do {
    Type type;
    if (failed(parser.parseType(type))) {
      return failure();
    }
    NamedAttrList attrs;
    if (failed(parser.parseOptionalAttrDict(attrs))) {
      return failure();
    }
    resultTypes.push_back(type);
    resultAttrsVec.push_back(parser.getBuilder().getDictionaryAttr(attrs));
  } while (succeeded(parser.parseOptionalComma()));
  if (!resultAttrsVec.empty()) {
    resultAttrs = parser.getBuilder().getArrayAttr(resultAttrsVec);
  }
  return success();
}

static void printDispatchFunctionResultList(OpAsmPrinter &p, Operation *op,
                                            TypeRange resultTypes,
                                            ArrayAttr resultAttrs) {
  for (unsigned i = 0; i < resultTypes.size(); ++i) {
    auto resultType = resultTypes[i];
    p.printType(resultType);
    if (resultAttrs) {
      auto attrs =
          dyn_cast_if_present<DictionaryAttr>(resultAttrs.getValue()[i]);
      if (attrs && !attrs.empty()) {
        p.printOptionalAttrDict(attrs.getValue());
      }
    }
    if (i < resultTypes.size() - 1) {
      p << ", ";
    }
  }
}

ParseResult parseDispatchFunctionSignature(OpAsmParser &parser,
                                           TypeAttr &functionTypeAttr,
                                           ArrayAttr &argAttrs,
                                           ArrayAttr &resultAttrs) {
  SmallVector<OpAsmParser::UnresolvedOperand> args;
  SmallVector<Type> argTypes;
  SmallVector<Type> resultTypes;
  if (failed(parser.parseLParen())) {
    return failure();
  }
  if (failed(parser.parseOptionalRParen())) {
    if (failed(parseDispatchFunctionArgumentList(parser, args, argTypes,
                                                 argAttrs)) ||
        failed(parser.parseRParen())) {
      return failure();
    }
  }
  if (succeeded(parser.parseOptionalArrow())) {
    if (succeeded(parser.parseOptionalLParen())) {
      if (failed(parseDispatchFunctionResultList(parser, resultTypes,
                                                 resultAttrs)) ||
          failed(parser.parseRParen())) {
        return failure();
      }
    } else {
      if (failed(parseDispatchFunctionResultList(parser, resultTypes,
                                                 resultAttrs))) {
        return failure();
      }
    }
  }
  functionTypeAttr = TypeAttr::get(
      FunctionType::get(parser.getContext(), argTypes, resultTypes));
  return success();
}

void printDispatchFunctionSignature(OpAsmPrinter &p, Operation *op,
                                    TypeAttr functionTypeAttr,
                                    ArrayAttr argAttrs, ArrayAttr resultAttrs) {
  auto functionType = cast<FunctionType>(functionTypeAttr.getValue());
  p << "(";
  for (size_t argIndex = 0; argIndex < functionType.getNumInputs();) {
    if (argIndex) {
      p << ", ";
    }
    int baseArgIndex = argIndex;
    auto type = functionType.getInput(baseArgIndex);
    p << "%arg";
    p << (baseArgIndex + 0);
    if (isa<IREE::Stream::ResourceType>(type)) {
      p << "[%arg" << (baseArgIndex + 1) << " for %arg" << (baseArgIndex + 2)
        << "]";
      argIndex += 3; // <resource, offset, length>
    } else {
      argIndex += 1; // unmodified arg
    }
    p << ": ";
    p.printType(type);
    if (argAttrs) {
      auto attrs = dyn_cast_if_present<DictionaryAttr>(
          argAttrs.getValue()[baseArgIndex]);
      if (attrs && !attrs.empty()) {
        p.printOptionalAttrDict(attrs.getValue());
      }
    }
  }
  p << ")";
  auto resultTypes = functionType.getResults();
  if (!resultTypes.empty()) {
    p << " -> ";
    if (resultTypes.size() != 1) {
      p << "(";
    }
    printDispatchFunctionResultList(p, op, resultTypes, resultAttrs);
    if (resultTypes.size() != 1) {
      p << ")";
    }
  }
}

//===----------------------------------------------------------------------===//
// stream.cmd.call
//===----------------------------------------------------------------------===//

LogicalResult CmdCallOp::verify() {
  CmdCallOp op = *this;

  if (failed(verifyOpValueSizes(op, op.getResourceOperands(),
                                op.getResourceOperandSizes())) ||
      failed(verifyOpValueSizes(op, op.getResults(), op.getResultSizes()))) {
    return failure();
  }

  unsigned resourceCount = 0;
  for (auto value : op.getResourceOperands()) {
    if (isa<IREE::Stream::ResourceType>(value.getType())) {
      ++resourceCount;
    }
  }

  unsigned rangeCount = op.getResourceOperandOffsets().size();
  if (op.getResourceOperandLengths().size() != rangeCount) {
    return op->emitOpError() << "mismatch on resource range "
                                "offsets/lengths; counts must match";
  }
  if (rangeCount != resourceCount) {
    return op->emitOpError()
           << "expects " << resourceCount << " resource range operand sets but "
           << rangeCount << " are present";
  }

  if (op.getResourceOperandAccessesAttr().size() != resourceCount) {
    return op->emitOpError()
           << "expects " << resourceCount << " resource access specifiers but "
           << op.getResourceOperandAccessesAttr().size() << " are present";
  }

  return success();
}

LogicalResult CmdCallOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
  Operation *op = getOperation();
  auto calleeOp = symbolTable.lookupNearestSymbolFrom(op, getCalleeAttr());
  if (!calleeOp) {
    return op->emitOpError() << "undefined external call: " << getCallee();
  }

  // TODO(benvanik): verification against the target callee.

  return success();
}

static ParseResult parseCmdCallOperands(
    OpAsmParser &parser,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceOperands,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceOffsets,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &resourceLengths,
    ArrayAttr &resourceAccesses) {
  if (failed(parser.parseLParen())) {
    return failure();
  }
  // Handle the case of no operands specially.
  if (succeeded(parser.parseOptionalRParen())) {
    return success();
  }
  SmallVector<Attribute> accessAttrs;
  do {
    StringRef accessStr;
    if (succeeded(
            parser.parseOptionalKeyword(&accessStr, {"ro", "rw", "wo"}))) {
      // A resource operand that'll have an offset/length associated.
      IREE::Stream::ResourceAccessBitfield accessBits =
          IREE::Stream::ResourceAccessBitfield::None;
      if (accessStr == "ro") {
        accessBits = IREE::Stream::ResourceAccessBitfield::Read;
      } else if (accessStr == "wo") {
        accessBits = IREE::Stream::ResourceAccessBitfield::Write;
      } else if (accessStr == "rw") {
        accessBits = IREE::Stream::ResourceAccessBitfield::Read |
                     IREE::Stream::ResourceAccessBitfield::Write;
      }
      accessAttrs.push_back(IREE::Stream::ResourceAccessBitfieldAttr::get(
          parser.getBuilder().getContext(), accessBits));
      resourceOperands.emplace_back();
      resourceOffsets.emplace_back();
      resourceLengths.emplace_back();
      if (failed(parser.parseOperand(resourceOperands.back())) ||
          failed(parser.parseLSquare()) ||
          failed(parser.parseOperand(resourceOffsets.back())) ||
          failed(parser.parseKeyword("for")) ||
          failed(parser.parseOperand(resourceLengths.back())) ||
          failed(parser.parseRSquare())) {
        return failure();
      }
    } else {
      // Primitive/custom operand.
      resourceOperands.emplace_back();
      if (failed(parser.parseOperand(resourceOperands.back()))) {
        return failure();
      }
    }
  } while (succeeded(parser.parseOptionalComma()));
  resourceAccesses = parser.getBuilder().getArrayAttr(accessAttrs);
  if (failed(parser.parseRParen())) {
    return failure();
  }
  return success();
}

static void printCmdCallOperands(OpAsmPrinter &p, Operation *op,
                                 ValueRange resourceOperands,
                                 ValueRange resourceOffsets,
                                 ValueRange resourceLengths,
                                 ArrayAttr resourceAccesses) {
  p << "(";
  size_t resourceIndex = 0;
  for (size_t i = 0; i < resourceOperands.size(); ++i) {
    auto operand = resourceOperands[i];
    if (isa<IREE::Stream::ResourceType>(operand.getType())) {
      // Resource type.
      auto resourceOffset = resourceOffsets[resourceIndex];
      auto resourceLength = resourceLengths[resourceIndex];
      auto resourceAccess = cast<IREE::Stream::ResourceAccessBitfieldAttr>(
                                resourceAccesses[resourceIndex])
                                .getValue();
      if (bitEnumContainsAll(resourceAccess,
                             IREE::Stream::ResourceAccessBitfield::Read |
                                 IREE::Stream::ResourceAccessBitfield::Write)) {
        p << "rw";
      } else if (bitEnumContainsAll(
                     resourceAccess,
                     IREE::Stream::ResourceAccessBitfield::Read)) {
        p << "ro";
      } else if (bitEnumContainsAll(
                     resourceAccess,
                     IREE::Stream::ResourceAccessBitfield::Write)) {
        p << "wo";
      }
      p << ' ';
      p.printOperand(operand);
      p << "[";
      p.printOperand(resourceOffset);
      p << " for ";
      p.printOperand(resourceLength);
      p << "]";
      ++resourceIndex;
    } else {
      // Primitive/custom type.
      p.printOperand(operand);
    }
    if (i < resourceOperands.size() - 1) {
      p << ", ";
    }
  }
  p << ")";
}

//===----------------------------------------------------------------------===//
// stream.cmd.execute
//===----------------------------------------------------------------------===//

void CmdExecuteOp::build(OpBuilder &builder, OperationState &state,
                         Value awaitTimepoint, ValueRange operands,
                         ValueRange operandSizes,
                         ArrayRef<NamedAttribute> attributes) {
  state.addTypes(IREE::Stream::TimepointType::get(builder.getContext()));
  state.addOperands(operands);
  state.addOperands(operandSizes);
  if (awaitTimepoint) {
    state.addOperands(awaitTimepoint);
  }
  state.addAttributes(attributes);
  state.attributes.erase(getOperandSegmentSizeAttr());
  state.addAttribute(getOperandSegmentSizeAttr(),
                     builder.getDenseI32ArrayAttr({
                         static_cast<int32_t>(operands.size()),
                         static_cast<int32_t>(operandSizes.size()),
                         awaitTimepoint ? 1 : 0,
                     }));
  state.addRegion();
}

// Returns success if the given op is a known valid stream.cmd.* op for use
// within an execution region.
static LogicalResult verifyCmdOp(Operation *op) {
  if (!op->hasTrait<OpTrait::IREE::Stream::CmdPhaseOp>() &&
      !isa<IREE::Stream::StreamableOpInterface>(op) &&
      !isa<IREE::Stream::YieldOp>(op)) {
    return op->emitOpError()
           << "explicit execution regions must only contain explicit ops";
  }
  return success();
}

LogicalResult CmdExecuteOp::verify() {
  CmdExecuteOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getResourceOperands(),
                                op.getResourceOperandSizes()))) {
    return failure();
  }
  if (failed(verifyAllResourcesCaptured(op.getBody()))) {
    return failure();
  }
  for (auto &nestedOp : op.getBody().front()) {
    if (failed(verifyCmdOp(&nestedOp))) {
      return failure();
    }
  }
  return success();
}

OperandRange
CmdExecuteOp::getEntrySuccessorOperands(RegionSuccessor successor) {
  assert(successor.getSuccessor() == &getBody() && "invalid region index");
  return getResourceOperands();
}

ValueRange CmdExecuteOp::getSuccessorInputs(RegionSuccessor successor) {
  return successor.isOperation() ? ValueRange()
                                 : ValueRange(getBodyRegion().getArguments());
}

void CmdExecuteOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  // Unconditional control flow into the region and back to the parent, so
  // return the correct RegionSuccessor purely based on the index being None or
  // 0.
  if (!point.isParent()) {
    regions.push_back(RegionSuccessor(getOperation()));
  } else {
    regions.push_back(RegionSuccessor(&getBody()));
  }
}

Operation::operand_range CmdExecuteOp::getClosureOperands() {
  return getResourceOperands();
}

Operation::result_range CmdExecuteOp::getClosureResults() {
  return Operation::result_range(nullptr, 0);
}

bool CmdExecuteOp::canClosureContainOp(Operation *op) { return false; }

IREE::Util::ValueAccess CmdExecuteOp::getOperandAccess(unsigned operandIndex) {
  auto arg = getBody().getArgument(operandIndex);
  return computeValueAccess(arg);
}

IREE::Util::ValueAccess CmdExecuteOp::getResultAccess(unsigned resultIndex) {
  return IREE::Util::ValueAccess::None();
}

IREE::Util::ClosureOpInterface
CmdExecuteOp::cloneReplacementExcludingOperandsAndResults(
    ArrayRef<unsigned> excludedOperandIndices,
    ArrayRef<unsigned> excludedResultIndices, PatternRewriter &rewriter) {
  SmallVector<Type> newResultTypes;
  SmallVector<Value> newResultSizes;
  auto newOperandsValues = llvm::to_vector(getResourceOperands());
  auto newOperandSizes = llvm::to_vector(getResourceOperandSizes());
  IREE::Util::excludeClosureOperandsAndResults(
      newOperandsValues, newOperandSizes, excludedOperandIndices,
      newResultTypes, newResultSizes, excludedResultIndices);

  auto newOp = CmdExecuteOp::create(rewriter, getLoc(), getAwaitTimepoint(),
                                    newOperandsValues, newOperandSizes,
                                    getOperation()->getAttrs());
  newOp.setOnce(getOnce());
  auto &newBody = newOp.getClosureBodyRegion();
  newBody.takeBody(getClosureBodyRegion());
  auto &block = newBody.front();
  BitVector eraseIndices(block.getNumArguments());
  for (auto i : excludedOperandIndices) {
    eraseIndices.set(i);
  }
  block.eraseArguments(eraseIndices);
  return newOp;
}

//===----------------------------------------------------------------------===//
// stream.cmd.serial
//===----------------------------------------------------------------------===//

LogicalResult CmdSerialOp::verify() {
  CmdSerialOp op = *this;
  for (auto &nestedOp : op.getBody().front()) {
    if (failed(verifyCmdOp(&nestedOp))) {
      return failure();
    }
  }
  return success();
}

ValueRange CmdSerialOp::getSuccessorInputs(RegionSuccessor successor) {
  return successor.isOperation() ? getOperation()->getResults() : ValueRange();
}

void CmdSerialOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  // Unconditional control flow into the region and back to the parent, so
  // return the correct RegionSuccessor purely based on the index being None or
  // 0.
  if (!point.isParent()) {
    regions.push_back(RegionSuccessor(getOperation()));
  } else {
    regions.push_back(RegionSuccessor(&getBody()));
  }
}

//===----------------------------------------------------------------------===//
// stream.cmd.concurrent
//===----------------------------------------------------------------------===//

LogicalResult CmdConcurrentOp::verify() {
  CmdConcurrentOp op = *this;
  for (auto &nestedOp : op.getBody().front()) {
    if (failed(verifyCmdOp(&nestedOp))) {
      return failure();
    }
  }
  return success();
}

ValueRange CmdConcurrentOp::getSuccessorInputs(RegionSuccessor successor) {
  return successor.isOperation() ? getOperation()->getResults() : ValueRange();
}

void CmdConcurrentOp::getSuccessorRegions(
    RegionBranchPoint point, SmallVectorImpl<RegionSuccessor> &regions) {
  // Unconditional control flow into the region and back to the parent, so
  // return the correct RegionSuccessor purely based on the index being None or
  // 0.
  if (!point.isParent()) {
    regions.push_back(RegionSuccessor(getOperation()));
  } else {
    regions.push_back(RegionSuccessor(&getBody()));
  }
}

//===----------------------------------------------------------------------===//
// stream.cmd.parameter.load
//===----------------------------------------------------------------------===//

LogicalResult CmdParameterLoadOp::verify() {
  CmdParameterLoadOp op = *this;
  size_t expectedCount = op.getSourceKeys().size();
  if (op.getSourceOffsets().size() != expectedCount ||
      op.getResultSizes().size() != expectedCount) {
    return op.emitOpError() << "requires that the source keys, source offsets, "
                               "and result sizes are all 1:1";
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.cmd.parameter.read
//===----------------------------------------------------------------------===//

LogicalResult CmdParameterReadOp::verify() {
  CmdParameterReadOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getTarget(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.cmd.parameter.write
//===----------------------------------------------------------------------===//

LogicalResult CmdParameterWriteOp::verify() {
  CmdParameterWriteOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.cmd.parameter.gather
//===----------------------------------------------------------------------===//

LogicalResult CmdParameterGatherOp::verify() {
  CmdParameterGatherOp op = *this;
  size_t expectedCount = op.getSourceKeys().size();
  if (op.getSourceOffsets().size() != expectedCount ||
      op.getTargetOffsets().size() != expectedCount ||
      op.getTargetLengths().size() != expectedCount) {
    return op.emitOpError()
           << "requires that the source keys, source offsets, target offsets, "
              "and target lengths are all 1:1";
  }
  if (failed(verifyOpValueSizes(op, op.getTarget(), op.getTargetSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.cmd.parameter.scatter
//===----------------------------------------------------------------------===//

LogicalResult CmdParameterScatterOp::verify() {
  CmdParameterScatterOp op = *this;
  size_t expectedCount = op.getTargetKeys().size();
  if (op.getSourceOffsets().size() != expectedCount ||
      op.getSourceLengths().size() != expectedCount ||
      op.getTargetOffsets().size() != expectedCount) {
    return op.emitOpError()
           << "requires that the source offsets, source lengths, target keys, "
              "and target offsets are all 1:1";
  }
  if (failed(verifyOpValueSizes(op, op.getSource(), op.getSourceSize()))) {
    return failure();
  }
  return success();
}

//===----------------------------------------------------------------------===//
// stream.timepoint.join
//===----------------------------------------------------------------------===//

LogicalResult TimepointJoinOp::verify() {
  // We could test if timepoints all come from the same place - this is not
  // strictly required but if we could avoid it things will be easier to
  // implement at runtime (won't have to do a cuda<->vulkan sync, etc).
  return success();
}

// static
Value TimepointJoinOp::join(Location loc, ValueRange timepoints,
                            OpBuilder &builder) {
  assert(!timepoints.empty() && "must have at least one timepoint");
  if (timepoints.size() == 1) {
    return timepoints.front();
  }
  return IREE::Stream::TimepointJoinOp::create(
      builder, loc, builder.getType<IREE::Stream::TimepointType>(), timepoints);
}

// static
Value TimepointJoinOp::join(ValueRange timepoints, OpBuilder &builder) {
  return join(builder.getFusedLoc(llvm::map_to_vector(
                  timepoints, [](Value value) { return value.getLoc(); })),
              timepoints, builder);
}

void TimepointJoinOp::setAwaitTimepoint(Value timepoint) {
  if (timepoint) {
    getAwaitTimepointsMutable().assign(timepoint);
  } else {
    getAwaitTimepointsMutable().clear();
  }
}

void TimepointJoinOp::setAwaitTimepoints(ValueRange timepoints,
                                         OpBuilder &builder) {
  getAwaitTimepointsMutable().assign(timepoints);
}

//===----------------------------------------------------------------------===//
// stream.timepoint.barrier
//===----------------------------------------------------------------------===//

LogicalResult TimepointBarrierOp::verify() {
  TimepointBarrierOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getResource(), op.getResourceSize()))) {
    return failure();
  }
  return success();
}

Value TimepointBarrierOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(getResource());
}

::std::optional<unsigned>
TimepointBarrierOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {0};
}

SmallVector<int64_t> TimepointBarrierOp::getTiedResultOperandIndices() {
  return {0};
}

std::pair<unsigned, unsigned>
TimepointBarrierOp::getTiedResultsIndexAndLength() {
  return {0, 1};
}

//===----------------------------------------------------------------------===//
// stream.timepoint.await
//===----------------------------------------------------------------------===//

void TimepointAwaitOp::build(OpBuilder &builder, OperationState &state,
                             ValueRange operands, ValueRange operandSizes,
                             Value timepoint,
                             ArrayRef<NamedAttribute> attributes) {
  state.addTypes(llvm::map_range(
      operands, [&](Value operand) { return operand.getType(); }));
  state.addOperands(operands);
  state.addOperands(operandSizes);
  state.addOperands(timepoint);
  state.addAttributes(attributes);
  state.attributes.erase(getOperandSegmentSizeAttr());
  state.addAttribute(getOperandSegmentSizeAttr(),
                     builder.getDenseI32ArrayAttr({
                         static_cast<int32_t>(operands.size()),
                         static_cast<int32_t>(operandSizes.size()),
                         static_cast<int32_t>(1), // timepoint
                     }));
}

LogicalResult TimepointAwaitOp::verify() {
  TimepointAwaitOp op = *this;
  if (failed(verifyOpValueSizes(op, op.getResourceOperands(),
                                op.getResourceOperandSizes())) ||
      failed(verifyOpValueSizes(op, op.getResults(),
                                op.getResourceOperandSizes()))) {
    return failure();
  }
  return success();
}

::std::optional<unsigned>
TimepointAwaitOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {resultIndex};
}

SmallVector<int64_t> TimepointAwaitOp::getTiedResultOperandIndices() {
  return llvm::to_vector(llvm::seq<int64_t>(0, getResourceOperands().size()));
}

SmallVector<Value> TimepointAwaitOp::getAwaitTimepoints() {
  if (getAwaitTimepoint()) {
    return {getAwaitTimepoint()};
  } else {
    return {};
  }
}

void TimepointAwaitOp::setAwaitTimepoint(Value timepoint) {
  if (timepoint) {
    getAwaitTimepointMutable().set(timepoint);
  } else {
    getAwaitTimepointMutable().drop();
  }
}

Value TimepointAwaitOp::getResultTimepoint() { return {}; }

//===----------------------------------------------------------------------===//
// stream.channel.create
//===----------------------------------------------------------------------===//

void ChannelCreateOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "channel");
}

//===----------------------------------------------------------------------===//
// stream.channel.split
//===----------------------------------------------------------------------===//

void ChannelSplitOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "channel");
}

//===----------------------------------------------------------------------===//
// stream.channel.rank
//===----------------------------------------------------------------------===//

void ChannelRankOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "ccl_rank");
}

//===----------------------------------------------------------------------===//
// stream.channel.count
//===----------------------------------------------------------------------===//

void ChannelCountOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getResult(), "ccl_count");
}

//===----------------------------------------------------------------------===//
// stream.executable
//===----------------------------------------------------------------------===//

void ExecutableOp::build(OpBuilder &builder, OperationState &state,
                         StringRef sym_name) {
  ensureTerminator(*state.addRegion(), builder, state.location);
  state.addAttribute(mlir::SymbolTable::getSymbolAttrName(),
                     builder.getStringAttr(sym_name));
}

LogicalResult ExecutableOp::verify() {
  // TODO(benvanik): check export name conflicts.
  return success();
}

//===----------------------------------------------------------------------===//
// stream.executable.export
//===----------------------------------------------------------------------===//

void ExecutableExportOp::build(OpBuilder &builder, OperationState &state,
                               StringRef sym_name,
                               FlatSymbolRefAttr function_ref) {
  build(builder, state, /*sym_visibility=*/nullptr,
        builder.getStringAttr(sym_name), function_ref);
}

LogicalResult ExecutableExportOp::verify() {
  // Workgroup count region is optional.
  if (!getWorkgroupCount().empty()) {
    // Verify the return ops all provide XYZ values.
    for (auto returnOp : getWorkgroupCount().getOps<IREE::Stream::ReturnOp>()) {
      if (returnOp.getNumOperands() != 3 ||
          !llvm::all_of(returnOp.getOperandTypes(),
                        [](Type type) { return type.isIndex(); })) {
        return returnOp.emitOpError()
               << "workgroup count region must return the XYZ dimension counts";
      }
    }
  }
  return success();
}

mlir::FunctionOpInterface ExecutableExportOp::lookupFunctionRef() {
  auto executableOp =
      this->getOperation()->getParentOfType<IREE::Stream::ExecutableOp>();
  if (!executableOp) {
    return {};
  }
  auto innerModuleOp = executableOp.getInnerModule();
  if (!innerModuleOp) {
    return {};
  }
  return innerModuleOp.lookupSymbol<mlir::FunctionOpInterface>(
      getFunctionRef());
}

//===----------------------------------------------------------------------===//
// stream.binding.subspan
//===----------------------------------------------------------------------===//

LogicalResult BindingSubspanOp::verify() {
  BindingSubspanOp op = *this;
  if (auto shapedType = dyn_cast<ShapedType>(op.getType())) {
    if (failed(verifyOpDynamicDims(op, shapedType, op.getDynamicDims()))) {
      return failure();
    }
  }

  return success();
}

//===----------------------------------------------------------------------===//
// stream.dispatch.workgroup.*
//===----------------------------------------------------------------------===//

static void getAsmResultNamesForDispatchWorkgroupInfoOp(
    StringRef prefix, const APInt &dimension, Value result,
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(result, (prefix + std::to_string(dimension.getZExtValue())).str());
}

static LogicalResult verifyDispatchWorkgroupInfoOp(Operation *op,
                                                   uint64_t dimension) {
  if (dimension < 0 || dimension >= 3) {
    return op->emitOpError()
           << "dimension " << dimension
           << " out of bounds of dispatch dimensions; expected [0, 3)";
  }
  return success();
}

void DispatchWorkgroupIDOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  getAsmResultNamesForDispatchWorkgroupInfoOp("workgroup_id_", getDimension(),
                                              getResult(), setNameFn);
}

LogicalResult DispatchWorkgroupIDOp::verify() {
  return verifyDispatchWorkgroupInfoOp(getOperation(),
                                       getDimension().getZExtValue());
}

void DispatchWorkgroupCountOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  getAsmResultNamesForDispatchWorkgroupInfoOp(
      "workgroup_count_", getDimension(), getResult(), setNameFn);
}

LogicalResult DispatchWorkgroupCountOp::verify() {
  return verifyDispatchWorkgroupInfoOp(getOperation(),
                                       getDimension().getZExtValue());
}

void DispatchWorkgroupSizeOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  getAsmResultNamesForDispatchWorkgroupInfoOp("workgroup_size_", getDimension(),
                                              getResult(), setNameFn);
}

LogicalResult DispatchWorkgroupSizeOp::verify() {
  return verifyDispatchWorkgroupInfoOp(getOperation(),
                                       getDimension().getZExtValue());
}

//===----------------------------------------------------------------------===//
// stream.yield
//===----------------------------------------------------------------------===//

MutableOperandRange
YieldOp::getMutableSuccessorOperands(RegionSuccessor successor) {
  return getResourceOperandsMutable();
}

MutableOperandRange YieldOp::getClosureResultsMutable() {
  return getResourceOperandsMutable();
}

//===----------------------------------------------------------------------===//
// stream.test.timeline_aware
//===----------------------------------------------------------------------===//

SmallVector<Value>
TestTimelineAwareOp::buildAwaitTimepoints(OpBuilder &builder) {
  // Build timepoint.import ops for each wait fence.
  SmallVector<Value> timepoints;
  for (auto fence : getWaitFenceLikes()) {
    auto importOp = IREE::Stream::TimepointImportOp::create(
        builder, getLoc(), builder.getType<IREE::Stream::TimepointType>(),
        ValueRange{fence}, /*affinity=*/nullptr);
    timepoints.push_back(importOp.getResultTimepoint());
  }
  return timepoints;
}

Value TestTimelineAwareOp::buildResultTimepoint(OpBuilder &builder) {
  // Build timepoint.import for signal fence.
  auto importOp = IREE::Stream::TimepointImportOp::create(
      builder, getLoc(), builder.getType<IREE::Stream::TimepointType>(),
      ValueRange{getSignalFenceLike()}, /*affinity=*/nullptr);
  return importOp.getResultTimepoint();
}

} // namespace mlir::iree_compiler::IREE::Stream

//===----------------------------------------------------------------------===//
// TableGen definitions (intentionally last)
//===----------------------------------------------------------------------===//

#define GET_OP_CLASSES
#include "iree/compiler/Dialect/Stream/IR/StreamOps.cpp.inc" // IWYU pragma: keep
