// Copyright 2019 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

#include "iree/compiler/Dialect/HAL/IR/HALOps.h"

#include "iree/compiler/Dialect/HAL/IR/HALTypes.h"
#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
#include "llvm/ADT/STLExtras.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/IRMapping.h"
#include "mlir/IR/OpImplementation.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/IR/TypeUtilities.h"
#include "mlir/Interfaces/FunctionImplementation.h"
#include "mlir/Interfaces/InferIntRangeInterface.h"
#include "mlir/Interfaces/InferTypeOpInterface.h"

namespace mlir::iree_compiler::IREE::HAL {

namespace {

// We aribtrarily say that unbounded dimensions in a torch program cannot
// exceed 53bits, making the maximum safe dimension 9007199254740991. The
// astute reader will note that this is also the maximum safe value in
// JavaScript, which also "happens" to be the largest mantissa value in a
// 64bit double. We need a maximum and in the absence of a better choice,
// with this one we are at least in good company. This limit is also used
// in the frontends.
static constexpr uint64_t MAX_DIM_VALUE = (static_cast<uint64_t>(1) << 53) - 1;

// Similarly we use a very conservative maximum rank value for specifying
// ranges of runtime rank resolution functions. Various frameworks have hard
// and practical limits ranging from 32 (numpy) to hundreds. At the time of
// writing, PyTorch throws weird errors if trying to print a tensor with a rank
// greater than 992. We really just want a smallish integer value to bound
// arithmetic, so we use an arbitrary maximum.
static constexpr uint64_t MAX_RANK_VALUE = 4096;

} // namespace

//===----------------------------------------------------------------------===//
// custom<DescriptorType>($descriptor_type)
//===----------------------------------------------------------------------===//

// Custom parser/printer to omit the wrapping `<` and `>` unlike autogenerated
// attribute parser/printer.
static ParseResult parseDescriptorType(OpAsmParser &parser,
                                       DescriptorTypeAttr &dtAttr) {
  StringRef enumKeyword;
  if (failed(parser.parseKeyword(&enumKeyword)))
    return failure();
  std::optional<DescriptorType> maybeEnum =
      symbolizeDescriptorType(enumKeyword);
  if (!maybeEnum)
    return failure();
  dtAttr = DescriptorTypeAttr::get(parser.getContext(), *maybeEnum);
  return success();
}

static void printDescriptorType(OpAsmPrinter &p, Operation *,
                                DescriptorTypeAttr dtAttr) {
  p << stringifyDescriptorType(dtAttr.getValue());
}

//===----------------------------------------------------------------------===//
// custom<PipelineBindings>($binding_ordinals,
//                          $binding_buffers,
//                          type($binding_buffers),
//                          $binding_offsets,
//                          $binding_lengths)
//===----------------------------------------------------------------------===//

static ParseResult parsePipelineBindings(
    OpAsmParser &parser,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &ordinals,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &buffers,
    SmallVectorImpl<Type> &bufferTypes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &bufferOffsets,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &bufferLengths) {
  do {
    OpAsmParser::UnresolvedOperand ordinal;
    OpAsmParser::UnresolvedOperand buffer;
    Type bufferType;
    OpAsmParser::UnresolvedOperand bufferOffset;
    OpAsmParser::UnresolvedOperand bufferLength;
    if (failed(parser.parseOperand(ordinal)) || failed(parser.parseEqual()) ||
        failed(parser.parseLParen()) || failed(parser.parseOperand(buffer)) ||
        failed(parser.parseColonType(bufferType)) ||
        failed(parser.parseRParen()) || failed(parser.parseLSquare()) ||
        failed(parser.parseOperand(bufferOffset)) ||
        failed(parser.parseComma()) ||
        failed(parser.parseOperand(bufferLength)) ||
        failed(parser.parseRSquare())) {
      return failure();
    }
    ordinals.push_back(ordinal);
    buffers.push_back(buffer);
    bufferTypes.push_back(bufferType);
    bufferOffsets.push_back(bufferOffset);
    bufferLengths.push_back(bufferLength);
  } while (succeeded(parser.parseOptionalComma()));
  return success();
}

static void printPipelineBindings(OpAsmPrinter &p, Operation *op,
                                  ValueRange ordinals, ValueRange buffers,
                                  TypeRange bufferTypes,
                                  ValueRange bufferOffsets,
                                  ValueRange bufferLengths) {
  llvm::interleaveComma(llvm::zip_equal(ordinals, buffers, bufferTypes,
                                        bufferOffsets, bufferLengths),
                        p,
                        [&](std::tuple<Value, Value, Type, Value, Value> it) {
                          p.printNewline();
                          p << "  ";
                          p.printOperand(std::get<0>(it));
                          p << " = (";
                          p.printOperand(std::get<1>(it));
                          p << " : ";
                          p.printType(std::get<2>(it));
                          p << ")[";
                          p.printOperand(std::get<3>(it));
                          p << ", ";
                          p.printOperand(std::get<4>(it));
                          p << "]";
                        });
  p.printNewline();
}

//===----------------------------------------------------------------------===//
// custom<Bindings>($binding_buffers,
//                  type($binding_buffers),
//                  $binding_offsets,
//                  $binding_lengths)
//===----------------------------------------------------------------------===//

static ParseResult
parseBindings(OpAsmParser &parser,
              SmallVectorImpl<OpAsmParser::UnresolvedOperand> &buffers,
              SmallVectorImpl<Type> &bufferTypes,
              SmallVectorImpl<OpAsmParser::UnresolvedOperand> &bufferOffsets,
              SmallVectorImpl<OpAsmParser::UnresolvedOperand> &bufferLengths) {
  do {
    OpAsmParser::UnresolvedOperand buffer;
    Type bufferType;
    OpAsmParser::UnresolvedOperand bufferOffset;
    OpAsmParser::UnresolvedOperand bufferLength;
    if (failed(parser.parseLParen()) || failed(parser.parseOperand(buffer)) ||
        failed(parser.parseColonType(bufferType)) ||
        failed(parser.parseRParen()) || failed(parser.parseLSquare()) ||
        failed(parser.parseOperand(bufferOffset)) ||
        failed(parser.parseComma()) ||
        failed(parser.parseOperand(bufferLength)) ||
        failed(parser.parseRSquare())) {
      return failure();
    }
    buffers.push_back(buffer);
    bufferTypes.push_back(bufferType);
    bufferOffsets.push_back(bufferOffset);
    bufferLengths.push_back(bufferLength);
  } while (succeeded(parser.parseOptionalComma()));
  return success();
}

static void printBindings(OpAsmPrinter &p, Operation *op, ValueRange buffers,
                          TypeRange bufferTypes, ValueRange bufferOffsets,
                          ValueRange bufferLengths) {
  llvm::interleaveComma(
      llvm::zip_equal(buffers, bufferTypes, bufferOffsets, bufferLengths), p,
      [&](std::tuple<Value, Type, Value, Value> it) {
        p.printNewline();
        p << "  (";
        p.printOperand(std::get<0>(it));
        p << " : ";
        p.printType(std::get<1>(it));
        p << ")[";
        p.printOperand(std::get<2>(it));
        p << ", ";
        p.printOperand(std::get<3>(it));
        p << "]";
      });
  p.printNewline();
}

//===----------------------------------------------------------------------===//
// custom<BindingTable>($binding_buffers,
//                      type($binding_buffers),
//                      $binding_offsets,
//                      $binding_lengths)
//===----------------------------------------------------------------------===//

static ParseResult parseBindingTable(
    OpAsmParser &parser,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &buffers,
    SmallVectorImpl<Type> &bufferTypes,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &bufferOffsets,
    SmallVectorImpl<OpAsmParser::UnresolvedOperand> &bufferLengths) {
  do {
    OpAsmParser::UnresolvedOperand buffer;
    Type bufferType;
    OpAsmParser::UnresolvedOperand bufferOffset;
    OpAsmParser::UnresolvedOperand bufferLength;
    if (failed(parser.parseLParen()) || failed(parser.parseOperand(buffer)) ||
        failed(parser.parseColonType(bufferType)) ||
        failed(parser.parseRParen()) || failed(parser.parseLSquare()) ||
        failed(parser.parseOperand(bufferOffset)) ||
        failed(parser.parseComma()) ||
        failed(parser.parseOperand(bufferLength)) ||
        failed(parser.parseRSquare())) {
      return failure();
    }
    buffers.push_back(buffer);
    bufferTypes.push_back(bufferType);
    bufferOffsets.push_back(bufferOffset);
    bufferLengths.push_back(bufferLength);
  } while (succeeded(parser.parseOptionalComma()));
  return success();
}

static void printBindingTable(OpAsmPrinter &p, Operation *op,
                              ValueRange buffers, TypeRange bufferTypes,
                              ValueRange bufferOffsets,
                              ValueRange bufferLengths) {
  llvm::interleaveComma(
      llvm::zip_equal(buffers, bufferTypes, bufferOffsets, bufferLengths), p,
      [&](std::tuple<Value, Type, Value, Value> it) {
        p.printNewline();
        p << "  ";
        p << "(";
        p.printOperand(std::get<0>(it));
        p << " : ";
        p.printType(std::get<1>(it));
        p << ")[";
        p.printOperand(std::get<2>(it));
        p << ", ";
        p.printOperand(std::get<3>(it));
        p << "]";
      });
  p.printNewline();
}

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

static FunctionType getTargetConditionRegionType(MLIRContext *context) {
  return FunctionType::get(context,
                           {
                               IREE::HAL::DeviceType::get(context),
                           },
                           {
                               IntegerType::get(context, 1),
                           });
}

static LogicalResult verifyTargetConditionRegion(Operation *op,
                                                 Region &region) {
  // Ignore if empty.
  if (region.empty()) {
    return success();
  }

  // Verify region takes a !hal.device.
  if (region.getNumArguments() != 1 ||
      !isa<IREE::HAL::DeviceType>(region.getArgumentTypes().front())) {
    return op->emitOpError()
           << "target condition region must take a !hal.device";
  }

  // Verify i1 return.
  for (auto returnOp : region.getOps<IREE::HAL::ReturnOp>()) {
    if (returnOp.getNumOperands() != 1) {
      return returnOp.emitOpError()
             << "target condition region must return a single i1 result";
    }
    for (auto returnType : returnOp.getOperandTypes()) {
      if (!returnType.isInteger(1)) {
        return returnOp.emitOpError()
               << "target condition region must return a single i1 result";
      }
    }
  }

  return success();
}

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

  SmallVector<Type> returnTypes;
  if (failed(parser.parseArrowTypeList(returnTypes))) {
    return failure();
  }
  if (returnTypes.size() != 1 ||
      !llvm::all_of(returnTypes, [](Type type) { return type.isInteger(1); })) {
    return parser.emitError(parser.getCurrentLocation())
           << "target condition region must return one i1";
  }

  return parser.parseRegion(body, args, /*enableNameShadowing=*/false);
}

static void printTargetConditionRegion(OpAsmPrinter &p, Operation *op,
                                       Region &body) {
  if (body.empty())
    return;
  p << "(";
  llvm::interleaveComma(body.getArguments(), p,
                        [&](BlockArgument arg) { p.printRegionArgument(arg); });
  p << ")";
  p.printArrowTypeList(TypeRange{IntegerType::get(body.getContext(), 1)});
  p << " ";
  p.printRegion(body, /*printEntryBlockArgs=*/false,
                /*printBlockTerminators=*/true);
}

static ParseResult parseTargetConditionObjects(
    OpAsmParser &parser, ArrayAttr &targetsAttr, ArrayAttr &targetOrdinalsAttr,
    ArrayAttr &targetObjectsAttr,
    SmallVector<std::unique_ptr<Region>, 2> &targetRegions) {
  SmallVector<Attribute> targetsAttrs;
  SmallVector<Attribute> targetOrdinalsAttrs;
  SmallVector<Attribute> targetObjectsAttrs;
  do {
    // #hal.executable.target<...>
    Attribute targetAttr;
    if (failed(parser.parseAttribute(targetAttr)))
      return failure();
    targetsAttrs.push_back(targetAttr);

    // if(...) -> i1 { ... }
    auto region = std::make_unique<Region>();
    if (succeeded(parser.parseOptionalKeyword("if"))) {
      if (failed(parseTargetConditionRegion(parser, *region)))
        return failure();
    }
    targetRegions.push_back(std::move(region));

    // ordinal(#)
    Attribute targetOrdinalAttr;
    if (failed(parser.parseKeyword("ordinal")) ||
        failed(parser.parseLParen()) ||
        failed(parser.parseAttribute(targetOrdinalAttr,
                                     IndexType::get(parser.getContext()))) ||
        failed(parser.parseRParen()))
      return failure();
    targetOrdinalsAttrs.push_back(targetOrdinalAttr);

    // = [#hal.executable.object<...>, ...]
    ArrayAttr targetObjectsAttr;
    if (failed(parser.parseEqual()) ||
        failed(parser.parseAttribute(targetObjectsAttr)))
      return failure();
    targetObjectsAttrs.push_back(targetObjectsAttr);
  } while (succeeded(parser.parseOptionalComma()));
  targetsAttr = ArrayAttr::get(parser.getContext(), targetsAttrs);
  targetOrdinalsAttr = ArrayAttr::get(parser.getContext(), targetOrdinalsAttrs);
  targetObjectsAttr = ArrayAttr::get(parser.getContext(), targetObjectsAttrs);
  return success();
}

static void printTargetConditionObjects(OpAsmPrinter &p, Operation *op,
                                        ArrayAttr targetsAttr,
                                        ArrayAttr targetOrdinalsAttr,
                                        ArrayAttr targetObjectsAttr,
                                        MutableArrayRef<Region> targetRegions) {
  p.increaseIndent();
  p.printNewline();

  llvm::interleave(
      llvm::zip_equal(targetsAttr, targetOrdinalsAttr, targetObjectsAttr,
                      targetRegions),
      [&](auto it) {
        auto &[targetAttr, targetOrdinalAttr, targetObjectsAttr, targetRegion] =
            it;
        p.printAttribute(targetAttr);
        if (!targetRegion.empty()) {
          p << " if";
          printTargetConditionRegion(p, op, targetRegion);
        }
        p << " ordinal(";
        p.printAttributeWithoutType(targetOrdinalAttr);
        p << ")";
        p << " = ";
        p.printAttribute(targetObjectsAttr);
      },
      [&]() {
        p << ",";
        p.printNewline();
      });

  p.decreaseIndent();
  p.printNewline();
}

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

static ParseResult parseWorkgroupCountRegion(OpAsmParser &parser,
                                             Region &body) {
  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::HAL::ReturnOp>()) {
    for (auto [resultType, returnType] :
         llvm::zip_equal(returnTypes, returnOp.getOperandTypes())) {
      if (resultType != returnType) {
        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 << "(";
  llvm::interleaveComma(body.getArguments(), p,
                        [&](BlockArgument arg) { p.printRegionArgument(arg); });
  p << ")";
  Type indexType = IndexType::get(body.getContext());
  p.printArrowTypeList(TypeRange{indexType, indexType, indexType});
  p << " ";
  p.printRegion(body, /*printEntryBlockArgs=*/false,
                /*printBlockTerminators=*/true);
}

//===----------------------------------------------------------------------===//
// hal.ex.*
//===----------------------------------------------------------------------===//

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

//===----------------------------------------------------------------------===//
// hal.return
//===----------------------------------------------------------------------===//

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

  auto parentFuncOp =
      dyn_cast_or_null<mlir::FunctionOpInterface>(op->getParentOp());
  if (parentFuncOp) {
    auto expectedTypes = parentFuncOp.getResultTypes();
    if (op.getNumOperands() != expectedTypes.size()) {
      return op.emitOpError() << "return must have the same number of operands "
                                 "as the parent result signature (have "
                              << op.getNumOperands() << ", expected "
                              << expectedTypes.size() << ")";
    }
    for (auto &&[index, values] :
         llvm::enumerate(llvm::zip_equal(op.getOperands(), expectedTypes))) {
      auto [operand, expectedType] = values;
      if (operand.getType() != expectedType) {
        return op.emitOpError()
               << "parent expected result " << index << " to be "
               << expectedType << " but returning " << operand.getType();
      }
    }
  }

  return success();
}

//===----------------------------------------------------------------------===//
// hal.tensor.import/export
//===----------------------------------------------------------------------===//

void TensorImportOp::build(OpBuilder &builder, OperationState &result,
                           Type resultType, Value source,
                           TypeAttr targetEncoding, StringAttr name,
                           Attribute affinity) {
  build(builder, result, resultType, source, targetEncoding,
        /*waitFence=*/Value{}, name, affinity);
}

void TensorImportOp::build(OpBuilder &builder, OperationState &result,
                           Type resultType, Value source,
                           TypeAttr targetEncoding, Value waitFence,
                           StringAttr name, Attribute affinity) {
  auto shapedType = llvm::cast<ShapedType>(resultType);
  assert((isa<IREE::HAL::BufferViewType>(source.getType()) ||
          shapedType.hasStaticShape()) &&
         "can only use this constructor for buffer views when shape "
         "information is required");
  SmallVector<Value> dynamicDims;
  for (int64_t i = 0; i < shapedType.getRank(); ++i) {
    if (!shapedType.isDynamicDim(i))
      continue;
    dynamicDims.push_back(builder.createOrFold<IREE::HAL::BufferViewDimOp>(
        result.location, builder.getIndexType(), source,
        builder.getIndexAttr(i)));
  }
  build(builder, result, resultType, source, targetEncoding, dynamicDims,
        waitFence, name, affinity);
}

static LogicalResult verifyTypeStorageCompatibility(Operation *op,
                                                    Type encodingType,
                                                    Type storageType) {
  if (encodingType == storageType)
    return success();
  auto encodingShapedType = llvm::dyn_cast<ShapedType>(encodingType);
  auto storageShapedType = llvm::dyn_cast<ShapedType>(storageType);
  if (!encodingShapedType || !storageShapedType)
    return success();

  if (IREE::Util::getRoundedElementByteWidth(
          encodingShapedType.getElementType()) !=
      IREE::Util::getRoundedElementByteWidth(
          storageShapedType.getElementType())) {
    // TODO(benvanik): more sophisticated logic here. There are a lot of valid
    // cases that are difficult to account for here statically; for example,
    // packing 8xi1 into 1xi8 or complex<f32> into 2xf32. We could try to guess
    // the element count (at least the static part of it) and ensure the scaling
    // matches but that wouldn't account for user variance. Really with this op
    // we are letting the _user_ control the bitcasting and type reflection and
    // purposefully don't want to mess with it (users should be able to put
    // custom types here, etc).
    //
    // NOTE: we round to bytes first as the base type (such as i1) may not be
    // representable in an external form.
    // return op->emitOpError() << "encoding and storage types must be "
    //                             "bitcastable; adjusted encoding bit width "
    //                             "of "
    //                          << encodingShapedType.getElementTypeBitWidth()
    //                          << " != adjusted storage bit width of "
    //                          << storageShapedType.getElementTypeBitWidth();
  }

  if (encodingShapedType.getNumDynamicDims() !=
      storageShapedType.getNumDynamicDims()) {
    // NOTE: we implicitly require that the dimensions are equivalent but
    // dont actually care about their order. For example, tensor<?x1xf32> is
    // compatible with tensor<?xf32>.
    return op->emitOpError()
           << "encoding and storage types must have the same "
              "dynamic dimension values; encoding shape "
           << encodingShapedType << " incompatible with storage shape "
           << storageShapedType;
  }

  return success();
}

LogicalResult TensorImportOp::verify() {
  TensorImportOp op = *this;
  auto targetType = llvm::cast<TensorType>(op.getTarget().getType());
  if (targetType.getNumDynamicDims() != op.getTargetDims().size()) {
    return op->emitOpError() << "number of target_dims must match number of "
                                "dynamic dims in target type";
  }
  return verifyTypeStorageCompatibility(op, op.getTargetEncoding(), targetType);
}

void TensorExportOp::build(OpBuilder &builder, OperationState &result,
                           Type resultType, Value source,
                           TypeAttr sourceEncoding, StringAttr name,
                           Attribute affinity) {
  auto dynamicDims =
      IREE::Util::buildDynamicDimsForValue(result.location, source, builder);
  build(builder, result, resultType, source, sourceEncoding, dynamicDims, name,
        affinity);
}

LogicalResult TensorExportOp::verify() {
  TensorExportOp op = *this;
  auto sourceType = llvm::cast<TensorType>(op.getSource().getType());
  if (sourceType.getNumDynamicDims() != op.getSourceDims().size()) {
    return op->emitOpError() << "number of source_dims must match number of "
                                "dynamic dims in source type";
  }
  return verifyTypeStorageCompatibility(op, op.getSourceEncoding(),
                                        op.getSource().getType());
}

//===----------------------------------------------------------------------===//
// hal.tensor.alias
//===----------------------------------------------------------------------===//

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

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

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

LogicalResult TensorAliasOp::verify() {
  TensorAliasOp op = *this;
  auto type = llvm::cast<TensorType>(op.getSource().getType());
  if (type.getNumDynamicDims() != op.getSourceDims().size()) {
    return op->emitOpError()
           << "number of dynamic dims must match the operand type";
  }
  return success();
}

//===----------------------------------------------------------------------===//
// hal.tensor.barrier
//===----------------------------------------------------------------------===//

void TensorBarrierOp::build(OpBuilder &builder, OperationState &result,
                            ValueRange sources, Value signalFence) {
  auto resultTypes = llvm::map_to_vector(
      sources, [](Value source) { return source.getType(); });
  build(builder, result, resultTypes, sources, signalFence);
}

Value TensorBarrierOp::getTiedResult(unsigned resultIndex) {
  return IREE::Util::TiedOpInterface::findTiedBaseValue(
      getSources()[resultIndex]);
}

::std::optional<unsigned>
TensorBarrierOp::getTiedResultOperandIndex(unsigned resultIndex) {
  return {resultIndex}; // sources[i]
}

SmallVector<int64_t> TensorBarrierOp::getTiedResultOperandIndices() {
  size_t numSources = getSources().size();
  return llvm::to_vector(llvm::seq<int64_t>(0, numSources));
}

//===----------------------------------------------------------------------===//
// hal.dispatch.extern
//===----------------------------------------------------------------------===//

void DispatchExternOp::build(OpBuilder &builder, OperationState &state,
                             ValueRange workload, TypeRange resultTypes,
                             ValueRange resultDims, ValueRange arguments,
                             ValueRange argumentDims,
                             ArrayRef<int64_t> tiedOperands,
                             IREE::HAL::ExecutableObjectsAttr targetObjects,
                             ArrayRef<NamedAttribute> attributes) {
  state.addTypes(resultTypes);
  state.addOperands(workload);
  state.addOperands(arguments);
  state.addOperands(argumentDims);
  state.addOperands(resultDims);
  state.addAttributes(attributes);
  state.attributes.erase(IREE::Util::TiedOpInterface::getStorageAttrName());
  state.addAttribute(IREE::Util::TiedOpInterface::getStorageAttrName(),
                     builder.getIndexArrayAttr(tiedOperands));
  state.addAttribute("targets", targetObjects.getTargets());
  state.addAttribute("target_objects", targetObjects.getTargetObjects());
  state.attributes.erase(getOperandSegmentSizeAttr());
  state.addAttribute(getOperandSegmentSizeAttr(),
                     builder.getDenseI32ArrayAttr({
                         static_cast<int32_t>(workload.size()),
                         static_cast<int32_t>(arguments.size()),
                         static_cast<int32_t>(argumentDims.size()),
                         static_cast<int32_t>(resultDims.size()),
                     }));

  // NOTE: workgroup count region is empty; callers are expected to populate it.
  state.addRegion();

  // Add one empty region per target.
  for (size_t i = 0; i < targetObjects.getTargets().size(); ++i)
    state.addRegion();
}

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

static LogicalResult
verifyWorkgroupCountRegion(Operation *op, ValueRange workload, Region &region) {
  // Verify the workload operands match the expected capture args.
  auto regionArguments =
      llvm::make_filter_range(region.getArgumentTypes(), [](Type type) {
        return !isa<IREE::HAL::DeviceType>(type);
      });
  if (workload.size() != llvm::range_size(regionArguments)) {
    return op->emitOpError()
           << "workload operands and workgroup count args mismatch ("
           << workload.size() << " vs " << llvm::range_size(regionArguments)
           << ")";
  }
  for (auto [index, values] :
       llvm::enumerate(llvm::zip_equal(workload, regionArguments))) {
    auto [workloadValue, capturedType] = values;
    if (workloadValue.getType() != capturedType) {
      return op->emitOpError()
             << "workload value " << index << " type mismatch; operand is "
             << workloadValue.getType() << " but region captures "
             << capturedType;
    }
  }
  return success();
}

LogicalResult DispatchExternOp::verify() {
  Operation *op = getOperation();

  if (failed(verifyOpDynamicDims(getOperation(), getArguments(),
                                 getArgumentDims())) ||
      failed(
          verifyOpDynamicDims(getOperation(), getResults(), getResultDims()))) {
    return failure();
  }

  auto verifyIOType = [&](Type type) -> LogicalResult {
    if (auto shapedType = llvm::dyn_cast<ShapedType>(type)) {
      if (shapedType.getElementType().isIndex()) {
        return op->emitOpError() << "I/O type " << type
                                 << " is invalid: index types must not cross "
                                    "the dispatch boundary";
      }
    }
    return success();
  };
  for (auto type : getOperandTypes()) {
    if (failed(verifyIOType(type)))
      return failure();
  }
  for (auto type : getResultTypes()) {
    if (failed(verifyIOType(type)))
      return failure();
  }

  if (failed(
          verifyWorkgroupCountRegion(op, getWorkload(), getWorkgroupCount()))) {
    return failure();
  }

  if (getTargets().size() != getTargetObjects().size()) {
    return op->emitOpError() << "target and objects arrays must match";
  }
  if (getTargets().size() != getTargetRegions().size()) {
    return op->emitOpError()
           << "target and condition regions must match (but they may be empty)";
  }
  for (auto &targetRegion : getTargetRegions()) {
    if (failed(verifyTargetConditionRegion(op, targetRegion))) {
      return failure();
    }
  }

  return success();
}

std::pair<unsigned, unsigned>
DispatchExternOp::getTiedOperandsIndexAndLength() {
  return getODSOperandIndexAndLength(1);
}

//===----------------------------------------------------------------------===//
// hal.device.memoize
//===----------------------------------------------------------------------===//

void DeviceMemoizeOp::build(OpBuilder &builder, OperationState &state,
                            TypeRange resultTypes, Value device,
                            Value queueAffinity,
                            ArrayRef<NamedAttribute> attributes) {
  state.addTypes(resultTypes);
  state.addOperands(device);
  state.addOperands(queueAffinity);
  state.addAttributes(attributes);
  state.addRegion();
}

void DeviceMemoizeOp::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({}));
  } else {
    regions.push_back(RegionSuccessor(&getBody(), getBody().getArguments()));
  }
}

//===----------------------------------------------------------------------===//
// hal.allocator.allocate
//===----------------------------------------------------------------------===//

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

Value AllocatorAllocateOp::getOperandSize(unsigned idx) { return {}; }

Value AllocatorAllocateOp::getResultSize(unsigned idx) {
  return getResultSize();
}

//===----------------------------------------------------------------------===//
// hal.allocator.import
//===----------------------------------------------------------------------===//

void AllocatorImportOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getDidImport(), "did_import");
  setNameFn(getResult(), "mapped");
}

Value AllocatorImportOp::getOperandSize(unsigned idx) { return {}; }

Value AllocatorImportOp::getResultSize(unsigned idx) { return getLength(); }

//===----------------------------------------------------------------------===//
// hal.buffer.subspan
//===----------------------------------------------------------------------===//

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

Value BufferSubspanOp::getOperandSize(unsigned idx) { return getLength(); }

Value BufferSubspanOp::getResultSize(unsigned idx) { return getLength(); }

//===----------------------------------------------------------------------===//
// hal.buffer.byte_length
//===----------------------------------------------------------------------===//

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

//===----------------------------------------------------------------------===//
// hal.element_type
//===----------------------------------------------------------------------===//

// Keep these in sync with iree/hal/buffer_view.h
enum class NumericalType : uint32_t {
  kUnknown = 0x00,
  kInteger = 0x10,
  kIntegerSigned = kInteger | 0x01,
  kIntegerUnsigned = kInteger | 0x02,
  kBoolean = kInteger | 0x03,
  kFloat = 0x20,
  kFloatIEEE = kFloat | 0x01,
  kFloatBrain = kFloat | 0x02,
  kFloatComplex = kFloat | 0x03,
  kFloat8E5M2 = kFloat | 0x04,
  kFloat8E4M3 = kFloat | 0x05,
  kFloat8E5M2FNUZ = kFloat | 0x06,
  kFloat8E4M3FNUZ = kFloat | 0x07,
};

constexpr inline int32_t makeElementTypeValue(NumericalType numericalType,
                                              int32_t bitCount) {
  return (static_cast<uint32_t>(numericalType) << 24) | bitCount;
}

// static
std::optional<int32_t> ElementTypeOp::getTypeValue(Type type) {
  if (auto intType = llvm::dyn_cast_if_present<IntegerType>(type)) {
    NumericalType numericalType;
    if (intType.isInteger(1)) {
      return makeElementTypeValue(NumericalType::kBoolean, 8);
    } else if (intType.isSigned()) {
      numericalType = NumericalType::kIntegerSigned;
    } else if (intType.isUnsigned()) {
      numericalType = NumericalType::kIntegerUnsigned;
    } else {
      // There's no such thing as a signless integer in machine types but we
      // need to be able to round-trip the format through the ABI. Exact
      // numerical type equality comparisons may fail if the frontend assumes
      // signed/unsigned but the compiler is propagating signless.
      numericalType = NumericalType::kInteger;
    }
    return makeElementTypeValue(numericalType, intType.getWidth());
  } else if (auto floatType = llvm::dyn_cast_if_present<FloatType>(type)) {
    switch (APFloat::SemanticsToEnum(floatType.getFloatSemantics())) {
    case APFloat::S_Float8E5M2:
      return makeElementTypeValue(NumericalType::kFloat8E5M2, 8);
    case APFloat::S_Float8E4M3:
      return makeElementTypeValue(NumericalType::kFloat8E4M3, 8);
    case APFloat::S_Float8E5M2FNUZ:
      return makeElementTypeValue(NumericalType::kFloat8E5M2FNUZ, 8);
    case APFloat::S_Float8E4M3FNUZ:
      return makeElementTypeValue(NumericalType::kFloat8E4M3FNUZ, 8);
    case APFloat::S_IEEEhalf:
    case APFloat::S_IEEEsingle:
    case APFloat::S_IEEEdouble:
    case APFloat::S_IEEEquad:
      return makeElementTypeValue(NumericalType::kFloatIEEE,
                                  floatType.getWidth());
    case APFloat::S_BFloat:
      return makeElementTypeValue(NumericalType::kFloatBrain,
                                  floatType.getWidth());
    default:
      return std::nullopt;
    }
  } else if (auto complexType = llvm::dyn_cast_if_present<ComplexType>(type)) {
    return makeElementTypeValue(
        NumericalType::kFloatComplex,
        complexType.getElementType().getIntOrFloatBitWidth() * 2);
  }
  return std::nullopt;
}

void ElementTypeOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  // We could make this match the C names.
  std::string name;
  llvm::raw_string_ostream os(name);
  os << "element_type_";
  os << getTypeAttr();
  setNameFn(getResult(), name);
}

LogicalResult ElementTypeOp::verify() {
  ElementTypeOp op = *this;
  auto value = getTypeValue(getTypeAttr().getValue());
  if (!value.has_value()) {
    return op.emitOpError("unsupported element type");
  }
  return success();
}

//===----------------------------------------------------------------------===//
// hal.encoding_type
//===----------------------------------------------------------------------===//

// static
std::optional<int32_t> EncodingTypeOp::getTypeValue(Attribute attr) {
  // TODO(#6762): encoding attribute handling/mapping to enums.
  if (attr)
    return std::nullopt;
  // Default to IREE_HAL_ENCODING_TYPE_DENSE_ROW_MAJOR for now.
  return 1;
}

void EncodingTypeOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  if (!getEncodingAttr())
    setNameFn(getResult(), "dense_row_major");
}

LogicalResult EncodingTypeOp::verify() {
  EncodingTypeOp op = *this;
  auto value = getTypeValue(getEncodingAttr());
  if (!value.has_value()) {
    return op.emitOpError("unsupported encoding type");
  }
  return success();
}

//===----------------------------------------------------------------------===//
// hal.buffer_view.create
//===----------------------------------------------------------------------===//

void BufferViewCreateOp::build(OpBuilder &builder, OperationState &state,
                               Value sourceBuffer, Value sourceOffset,
                               Value sourceLength, int32_t elementType,
                               int32_t encodingType, ValueRange shape) {
  build(builder, state, sourceBuffer, sourceOffset, sourceLength,
        builder.createOrFold<arith::ConstantIntOp>(state.location, elementType,
                                                   32),
        builder.createOrFold<arith::ConstantIntOp>(state.location, encodingType,
                                                   32),
        shape);
}

void BufferViewCreateOp::build(OpBuilder &builder, OperationState &state,
                               Value sourceBuffer, Value sourceOffset,
                               Value sourceLength, Value elementType,
                               Value encodingType, ValueRange shape) {
  state.addOperands(
      {sourceBuffer, sourceOffset, sourceLength, elementType, encodingType});
  state.addOperands(shape);
  state.addTypes({BufferViewType::get(builder.getContext())});
}

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

//===----------------------------------------------------------------------===//
// hal.buffer_view.buffer
//===----------------------------------------------------------------------===//

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

//===----------------------------------------------------------------------===//
// hal.buffer_view.dim
//===----------------------------------------------------------------------===//

void BufferViewDimOp::inferResultRangesFromOptional(
    ArrayRef<IntegerValueRange> argRanges, SetIntLatticeFn setResultRange) {
  const unsigned indexTypeNumBits = 64;
  setResultRange(getResult(), IntegerValueRange(ConstantIntRanges::fromUnsigned(
                                  APInt::getZero(indexTypeNumBits),
                                  APInt(indexTypeNumBits, MAX_DIM_VALUE))));
}

//===----------------------------------------------------------------------===//
// hal.buffer_view.dim
//===----------------------------------------------------------------------===//

void BufferViewRankOp::inferResultRangesFromOptional(
    ArrayRef<IntegerValueRange> argRanges, SetIntLatticeFn setResultRange) {
  const unsigned indexTypeNumBits = 64;
  setResultRange(getResult(), IntegerValueRange(ConstantIntRanges::fromUnsigned(
                                  APInt::getZero(indexTypeNumBits),
                                  APInt(indexTypeNumBits, MAX_RANK_VALUE))));
}

//===----------------------------------------------------------------------===//
// hal.channel.create
//===----------------------------------------------------------------------===//

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

//===----------------------------------------------------------------------===//
// hal.channel.split
//===----------------------------------------------------------------------===//

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

//===----------------------------------------------------------------------===//
// hal.channel.rank_and_count
//===----------------------------------------------------------------------===//

void ChannelRankAndCountOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getRank(), "ccl_rank");
  setNameFn(getCount(), "ccl_count");
}

//===----------------------------------------------------------------------===//
// hal.command_buffer.create
//===----------------------------------------------------------------------===//

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

//===----------------------------------------------------------------------===//
// hal.command_buffer.update_buffer
//===----------------------------------------------------------------------===//

IREE::Util::SubrangeOperand
CommandBufferUpdateBufferOp::getSubrangeOperand(unsigned operandIndex) {
  if (operandIndex == 1) {
    return IREE::Util::SubrangeOperand{getSourceBuffer(), getSourceSize(),
                                       getSourceOffset(), getLength()};
  } else {
    assert(false && "only source is a subrange");
    return {};
  }
}

void CommandBufferUpdateBufferOp::setSubrangeOperand(
    unsigned operandIndex, IREE::Util::SubrangeOperand operand) {
  if (operandIndex == 1) {
    getSourceBufferMutable().assign(operand.resource);
    getSourceSizeMutable().assign(operand.resourceSize);
    getSourceOffsetMutable().assign(operand.offset);
  } else {
    assert(false && "only source is a subrange");
  }
}

//===----------------------------------------------------------------------===//
// hal.command_buffer.dispatch + .indirect
//===----------------------------------------------------------------------===//

void CommandBufferDispatchOp::build(OpBuilder &builder, OperationState &state,
                                    Value commandBuffer, Value executable,
                                    Value entryPoint, ValueRange workgroups,
                                    ValueRange constants,
                                    ArrayRef<BindingValue> bindings,
                                    IREE::HAL::DispatchFlags flags) {
  state.addOperands({commandBuffer, executable, entryPoint});
  state.addOperands(workgroups);
  state.addOperands(constants);
  SmallVector<Value> bindingBuffers;
  SmallVector<Value> bindingOffsets;
  SmallVector<Value> bindingLengths;
  for (auto binding : bindings) {
    bindingBuffers.push_back(binding.buffer);
    bindingOffsets.push_back(binding.byteOffset);
    bindingLengths.push_back(binding.byteLength);
  }
  state.addOperands(bindingBuffers);
  state.addOperands(bindingOffsets);
  state.addOperands(bindingLengths);
  state.addAttribute("flags",
                     builder.getAttr<IREE::HAL::DispatchFlagsAttr>(flags));
  state.addAttribute(getOperandSegmentSizeAttr(),
                     builder.getDenseI32ArrayAttr({
                         1,
                         1,
                         1,
                         1,
                         1,
                         1,
                         static_cast<int32_t>(constants.size()),
                         static_cast<int32_t>(bindingBuffers.size()),
                         static_cast<int32_t>(bindingOffsets.size()),
                         static_cast<int32_t>(bindingLengths.size()),
                     }));
}

void CommandBufferDispatchIndirectOp::build(
    OpBuilder &builder, OperationState &state, Value commandBuffer,
    Value executable, Value entryPoint, Value workgroupsBuffer,
    Value workgroupsOffset, ValueRange constants,
    ArrayRef<BindingValue> bindings, IREE::HAL::DispatchFlags flags) {
  state.addOperands({commandBuffer, executable, entryPoint, workgroupsBuffer,
                     workgroupsOffset});
  state.addOperands(constants);
  SmallVector<Value> bindingBuffers;
  SmallVector<Value> bindingOffsets;
  SmallVector<Value> bindingLengths;
  for (auto binding : bindings) {
    bindingBuffers.push_back(binding.buffer);
    bindingOffsets.push_back(binding.byteOffset);
    bindingLengths.push_back(binding.byteLength);
  }
  state.addOperands(bindingBuffers);
  state.addOperands(bindingOffsets);
  state.addOperands(bindingLengths);
  state.addAttribute("flags",
                     builder.getAttr<IREE::HAL::DispatchFlagsAttr>(flags));
  state.addAttribute(getOperandSegmentSizeAttr(),
                     builder.getDenseI32ArrayAttr({
                         1,
                         1,
                         1,
                         1,
                         1,
                         static_cast<int32_t>(constants.size()),
                         static_cast<int32_t>(bindingBuffers.size()),
                         static_cast<int32_t>(bindingOffsets.size()),
                         static_cast<int32_t>(bindingLengths.size()),
                     }));
}

static LogicalResult verifyDispatchBindings(Operation *op,
                                            ValueRange bindingBuffers,
                                            ValueRange bindingOffsets,
                                            ValueRange bindingLengths) {
  if (bindingBuffers.size() != bindingOffsets.size() ||
      bindingBuffers.size() != bindingLengths.size()) {
    return op->emitOpError() << "requires that binding fields all have the "
                                "same number of elements";
  }
  return success();
}

LogicalResult CommandBufferDispatchOp::verify() {
  CommandBufferDispatchOp op = *this;
  return verifyDispatchBindings(op, op.getBindingBuffers(),
                                op.getBindingOffsets(), op.getBindingLengths());
}

LogicalResult CommandBufferDispatchIndirectOp::verify() {
  CommandBufferDispatchIndirectOp op = *this;
  return verifyDispatchBindings(op, op.getBindingBuffers(),
                                op.getBindingOffsets(), op.getBindingLengths());
}

//===----------------------------------------------------------------------===//
// hal.device.resolve
//===----------------------------------------------------------------------===//

void DeviceResolveOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  for (auto result : getResults()) {
    if (isa<IREE::HAL::DeviceType>(result.getType())) {
      setNameFn(result, "device");
    } else if (isa<IREE::HAL::AllocatorType>(result.getType())) {
      setNameFn(result, "allocator");
    } else if (isa<IntegerType>(result.getType())) {
      setNameFn(result, "queue_affinity");
    }
  }
}

//===----------------------------------------------------------------------===//
// hal.device.allocator
//===----------------------------------------------------------------------===//

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

//===----------------------------------------------------------------------===//
// hal.device.query
//===----------------------------------------------------------------------===//

LogicalResult DeviceQueryOp::verify() {
  DeviceQueryOp op = *this;
  if (op.getDefaultValue().has_value()) {
    if (op.getDefaultValue()->getType() != op.getValue().getType()) {
      return op.emitOpError()
             << "type mismatch between result and default value";
    }
  }
  return success();
}

// static
Value DeviceQueryOp::createI1(Location loc, Value device, StringRef category,
                              StringRef key, OpBuilder &builder) {
  auto i1Type = builder.getI1Type();
  return builder
      .create<IREE::HAL::DeviceQueryOp>(
          loc, i1Type, i1Type, device, builder.getStringAttr(category),
          builder.getStringAttr(key), builder.getIntegerAttr(i1Type, 0))
      .getValue();
}

//===----------------------------------------------------------------------===//
// hal.device.queue.*
//===----------------------------------------------------------------------===//

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

Value DeviceQueueAllocaOp::getOperandSize(unsigned idx) { return {}; }

Value DeviceQueueAllocaOp::getResultSize(unsigned idx) {
  return getResultSize();
}

static LogicalResult verifyDeviceQueueFences(Operation *queueOp,
                                             Value waitFence,
                                             Value signalFence) {
  if (waitFence == signalFence &&
      !isa<IREE::Util::NullOp>(waitFence.getDefiningOp())) {
    return queueOp->emitOpError() << "device queue operations cannot wait and "
                                     "signal on the same fence.";
  }
  return success();
}

LogicalResult DeviceQueueAllocaOp::verify() {
  return verifyDeviceQueueFences(*this, getWaitFence(), getSignalFence());
}

LogicalResult DeviceQueueDeallocaOp::verify() {
  return verifyDeviceQueueFences(*this, getWaitFence(), getSignalFence());
}

LogicalResult DeviceQueueFillOp::verify() {
  return verifyDeviceQueueFences(*this, getWaitFence(), getSignalFence());
}

LogicalResult DeviceQueueUpdateOp::verify() {
  return verifyDeviceQueueFences(*this, getWaitFence(), getSignalFence());
}

LogicalResult DeviceQueueCopyOp::verify() {
  return verifyDeviceQueueFences(*this, getWaitFence(), getSignalFence());
}

LogicalResult DeviceQueueReadOp::verify() {
  return verifyDeviceQueueFences(*this, getWaitFence(), getSignalFence());
}

LogicalResult DeviceQueueWriteOp::verify() {
  return verifyDeviceQueueFences(*this, getWaitFence(), getSignalFence());
}

LogicalResult DeviceQueueExecuteOp::verify() {
  return verifyDeviceQueueFences(*this, getWaitFence(), getSignalFence());
}

void DeviceQueueExecuteIndirectOp::build(OpBuilder &builder,
                                         OperationState &state, Value device,
                                         Value queueAffinity, Value waitFence,
                                         Value signalFence, Value commandBuffer,
                                         ArrayRef<BindingValue> bindings) {
  state.addOperands(
      {device, queueAffinity, waitFence, signalFence, commandBuffer});
  SmallVector<Value> bindingBuffers;
  SmallVector<Value> bindingOffsets;
  SmallVector<Value> bindingLengths;
  for (auto binding : bindings) {
    bindingBuffers.push_back(binding.buffer);
    bindingOffsets.push_back(binding.byteOffset);
    bindingLengths.push_back(binding.byteLength);
  }
  state.addOperands(bindingBuffers);
  state.addOperands(bindingOffsets);
  state.addOperands(bindingLengths);
}

LogicalResult DeviceQueueExecuteIndirectOp::verify() {
  return verifyDeviceQueueFences(*this, getWaitFence(), getSignalFence());
}

//===----------------------------------------------------------------------===//
// hal.devices.*
//===----------------------------------------------------------------------===//

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

void DevicesGetOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  APInt index;
  if (matchPattern(getIndex(), m_ConstantInt(&index))) {
    llvm::SmallString<16> str("device_");
    index.toStringUnsigned(str);
    setNameFn(getResult(), str);
  } else {
    setNameFn(getResult(), "device_n");
  }
}

//===----------------------------------------------------------------------===//
// hal.executable.source
//===----------------------------------------------------------------------===//

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

  auto conditionOps = getOps<IREE::HAL::ExecutableConditionOp>();
  if (llvm::range_size(conditionOps) > 1)
    return op.emitOpError()
           << "only one condition op is allowed in an executable";

  return success();
}

//===----------------------------------------------------------------------===//
// hal.executable
//===----------------------------------------------------------------------===//

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

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

//===----------------------------------------------------------------------===//
// hal.executable.export
//===----------------------------------------------------------------------===//

ParseResult ExecutableExportOp::parse(OpAsmParser &parser,
                                      OperationState &result) {
  StringAttr visibilityAttr;
  if (failed(parseSymbolVisibility(parser, visibilityAttr))) {
    return failure();
  }

  StringAttr nameAttr;
  IREE::HAL::PipelineLayoutAttr layoutAttr;
  if (failed(parser.parseSymbolName(nameAttr,
                                    mlir::SymbolTable::getSymbolAttrName(),
                                    result.attributes))) {
    return failure();
  }
  if (succeeded(parser.parseOptionalKeyword("ordinal"))) {
    IntegerAttr ordinalAttr;
    if (failed(parser.parseLParen()) ||
        failed(parser.parseAttribute(ordinalAttr,
                                     parser.getBuilder().getIndexType())) ||
        failed(parser.parseRParen())) {
      return failure();
    }
    result.addAttribute("ordinal", ordinalAttr);
  }
  if (failed(parser.parseKeyword("layout")) || failed(parser.parseLParen()) ||
      failed(parser.parseAttribute(layoutAttr)) ||
      failed(parser.parseRParen()) ||
      failed(parser.parseOptionalAttrDictWithKeyword(result.attributes))) {
    return failure();
  }
  result.addAttribute("layout", layoutAttr);

  std::unique_ptr<Region> region;
  SmallVector<OpAsmParser::Argument> regionOperands;
  // A missing optional region is materialized as an empty region.
  (void)parser.parseOptionalRegion(region, regionOperands);
  result.addRegion(std::move(region));

  return success();
}

void ExecutableExportOp::print(OpAsmPrinter &p) {
  Operation *op = getOperation();
  p << ' ';
  printSymbolVisibility(p, op, op->getAttrOfType<StringAttr>("sym_visibility"));
  p << ' ';
  p.printSymbolName(getSymName());
  if (getOrdinalAttr()) {
    p << " ordinal(";
    p.printAttributeWithoutType(getOrdinalAttr());
    p << ")";
  }
  p << " layout(";
  p.printAttribute(getLayout());
  p << ")";
  p.printOptionalAttrDictWithKeyword(
      op->getAttrs(),
      /*elidedAttrs=*/{"sym_name", "layout", "ordinal"});
  if (getWorkgroupCount().empty())
    return;
  p << " ";
  p.printRegion(getWorkgroupCount());
}

LogicalResult ExecutableExportOp::verify() {
  ExecutableExportOp op = *this;
  Block *body = getWorkgroupCountBody();
  // When there is no body, nothing to verify.
  if (!body)
    return success();

  if (!llvm::hasSingleElement(getWorkgroupCount())) {
    return op.emitOpError() << "expected a single region block";
  }
  bool validArguments = true;
  if (body->getNumArguments() == 0) {
    // Need at least a !hal.device.
    validArguments = false;
  } else if (!llvm::isa<IREE::HAL::DeviceType>(
                 body->getArgument(0).getType())) {
    // !hal.device must come first.
    validArguments = false;
  } else {
    // All remaining arguments need to be of type index (today).
    for (BlockArgument &blockArg : body->getArguments().drop_front(1)) {
      if (!llvm::isa<IndexType>(blockArg.getType())) {
        validArguments = false;
        break;
      }
    }
  }
  if (!validArguments) {
    return op.emitOpError(
        "expected workgroup_count to take (%device: !hal.device, "
        "%workload_0: index, %workload_1: index, ...");
  }
  // Check that the last statement in the block is `hal.return` operation.
  // TODO(ravishankarm): The SingleBlockImplicitTerminator<"HAL::ReturnOp">
  // should generate this check, but it doesnt.
  auto returnOp = dyn_cast<ReturnOp>(body->getTerminator());
  if (!returnOp || returnOp.getOperands().size() != getNumWorkgroupDims()) {
    return op.emitOpError("expected operation to yield ")
           << getNumWorkgroupDims() << " values";
  }
  return success();
}

// Calculates the workgroup count (x, y, z) given the total N-dimensional
// |workload| and specific |workgroupSize|.
static std::array<Value, 3>
calculateWorkloadWorkgroupCount(Location loc, ValueRange workload,
                                const std::array<Value, 3> &workgroupSize,
                                OpBuilder &builder) {
  std::array<Value, 3> result;

  auto constantOne = builder.createOrFold<arith::ConstantIndexOp>(loc, 1);
  if (workload.size() <= 3) {
    // 1-D to 3-D are easy (pad 2 to 0 dimensions) and divide by workgroup
    // size.
    for (int i = 0; i < 3; ++i) {
      // Round up: (workload[i] + workgroup_size - 1) / workgroup_size;
      Value workloadI = i < workload.size() ? workload[i] : constantOne;
      workloadI = builder.createOrFold<arith::SubIOp>(
          loc,
          builder.createOrFold<arith::AddIOp>(loc, workloadI, workgroupSize[i]),
          constantOne);
      result[i] = builder.createOrFold<arith::DivUIOp>(loc, workloadI,
                                                       workgroupSize[i]);
    }
  } else {
    // TODO(#4140): remapping of N-D to 3-D: this is not how you do this!
    Value flatWorkload = constantOne;
    for (auto workloadI : workload) {
      flatWorkload =
          builder.createOrFold<arith::MulIOp>(loc, flatWorkload, workloadI);
    }
    for (int i = 0; i < 3; ++i) {
      // Round up: (workload[i] + workgroup_size - 1) / workgroup_size;
      auto rounded = builder.createOrFold<arith::SubIOp>(
          loc,
          builder.createOrFold<arith::AddIOp>(loc, flatWorkload,
                                              workgroupSize[i]),
          constantOne);
      auto workgroupCountI =
          builder.createOrFold<arith::DivUIOp>(loc, rounded, workgroupSize[i]);
      result[i] = workgroupCountI;

      // Multiply back out and subtract from invocations.
      flatWorkload = builder.createOrFold<arith::SubIOp>(
          loc, flatWorkload,
          builder.createOrFold<arith::MulIOp>(loc, workgroupCountI, rounded));
    }
  }

  return result;
}

static std::array<Value, 3>
calculateWorkgroupCountFromRegion(Location loc, Block *body, Value device,
                                  ValueRange workload, OpBuilder &builder) {
  // TODO(benvanik): replace with region inlining util.
  IRMapping bvm;
  bvm.map(body->getArgument(0), device);
  // For now use the number of args to minimum of number of args used by
  // the body, and number of workload entries. When there is a more explicit
  // propagation of number of workload entries to the `hal.executable.variant`
  // this will be the same by construction.
  unsigned numArgs =
      std::min<unsigned>(body->getNumArguments() - 1, workload.size());
  for (unsigned argNum : llvm::seq<unsigned>(0, numArgs)) {
    bvm.map(body->getArgument(/*device*/ 1 + argNum), workload[argNum]);
  }
  for (Operation &op : body->without_terminator()) {
    builder.clone(op, bvm);
  }
  auto returnOp = cast<IREE::HAL::ReturnOp>(body->getTerminator());
  assert(returnOp.getNumOperands() == 3 && "must return xyz");
  return {
      bvm.lookup(returnOp.getOperands()[0]),
      bvm.lookup(returnOp.getOperands()[1]),
      bvm.lookup(returnOp.getOperands()[2]),
  };
}

// Calculates the workgroup count (x, y, z) for dispatching to the entry point.
// The provided N-dimensional |workload| is the total number of invocations
// required as calculated by the generic workload logic (basically, number of
// output elements in tensors).
std::array<Value, 3> ExecutableExportOp::calculateWorkgroupCount(
    Location loc, Value device, ValueRange workload, OpBuilder &builder) {
  Block *body = getWorkgroupCountBody();
  if (body) {
    return calculateWorkgroupCountFromRegion(loc, body, device, workload,
                                             builder);
  }
  auto workgroupSize = calculateWorkgroupSize(loc, device, workload, builder);
  return calculateWorkloadWorkgroupCount(loc, workload, workgroupSize, builder);
}

// Calculates the workgroup size (x, y, z). These are the dimension numbers
// for a single workgroup.
std::array<Value, 3> ExecutableExportOp::calculateWorkgroupSize(
    Location loc, Value device, ValueRange workload, OpBuilder &builder) {
  // When no workgroup size is specified we just assume [1,1,1].
  // This yields a workgroup count that models the extents of the workload.
  return {
      builder.createOrFold<arith::ConstantIndexOp>(loc, 1),
      builder.createOrFold<arith::ConstantIndexOp>(loc, 1),
      builder.createOrFold<arith::ConstantIndexOp>(loc, 1),
  };
}

//===----------------------------------------------------------------------===//
// hal.executable.variant
//===----------------------------------------------------------------------===//

void ExecutableVariantOp::build(OpBuilder &builder, OperationState &state,
                                StringRef symName,
                                IREE::HAL::ExecutableTargetAttr target) {
  ensureTerminator(*state.addRegion(), builder, state.location);
  state.addAttribute(mlir::SymbolTable::getSymbolAttrName(),
                     builder.getStringAttr(symName));
  state.addAttribute("target", target);
}

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

  auto conditionOps = getOps<IREE::HAL::ExecutableConditionOp>();
  if (llvm::range_size(conditionOps) > 1)
    return op.emitOpError() << "only one condition op is allowed in a variant";

  return success();
}

DenseMap<Attribute, int> ExecutableVariantOp::gatherConstantOrdinals() {
  DenseMap<Attribute, int> map;
  for (auto blockOp : getConstantBlockOps()) {
    int baseCount = map.size();
    for (auto [i, keyAttr] : llvm::enumerate(blockOp.getKeys())) {
      map.try_emplace(keyAttr, baseCount + i);
    }
  }
  return map;
}

Value ExecutableVariantOp::createConditionOp(OpBuilder &builder) {
  assert(!getConditionOp() && "condition op already exists");

  builder.setInsertionPointToStart(&getRegion().front());
  auto conditionOp = builder.create<IREE::HAL::ExecutableConditionOp>(getLoc());
  Block *entryPoint = conditionOp.addEntryBlock();
  Value device = entryPoint->getArgument(0);

  builder.setInsertionPointToStart(entryPoint);
  return device;
}

Value ExecutableVariantOp::buildCondition(Value device, OpBuilder &builder) {
  // Base case dependent on target information.
  Value selected = IREE::HAL::DeviceQueryOp::createI1(
      getLoc(), device, "hal.executable.format",
      getTarget().getFormat().getValue(), builder);

  // Factor in variant condition region, if any.
  auto conditionOp = getConditionOp();
  if (conditionOp) {
    auto regionOp = builder.create<scf::ExecuteRegionOp>(conditionOp.getLoc(),
                                                         builder.getI1Type());

    IRMapping mapper;
    mapper.map(conditionOp.getRegion().getArgument(0), device);
    conditionOp.getRegion().cloneInto(&regionOp.getRegion(), mapper);

    for (auto returnOp :
         llvm::make_early_inc_range(regionOp.getOps<IREE::HAL::ReturnOp>())) {
      OpBuilder(returnOp).create<scf::YieldOp>(returnOp.getLoc(),
                                               returnOp.getOperands());
      returnOp.erase();
    }

    selected = builder.create<arith::AndIOp>(getLoc(), selected,
                                             regionOp.getResult(0));
  }

  return selected;
}

//===----------------------------------------------------------------------===//
// hal.executable.condition
//===----------------------------------------------------------------------===//

LogicalResult ExecutableConditionOp::verify() {
  ExecutableConditionOp op = *this;
  return verifyTargetConditionRegion(op, op.getBody());
}

void ExecutableConditionOp::build(OpBuilder &builder, OperationState &result,
                                  ArrayRef<NamedAttribute> attrs) {
  result.addAttribute(
      "function_type",
      TypeAttr::get(getTargetConditionRegionType(builder.getContext())));
  result.addRegion();
  result.attributes.append(attrs.begin(), attrs.end());
}

ParseResult ExecutableConditionOp::parse(OpAsmParser &parser,
                                         OperationState &result) {
  if (parseTargetConditionRegion(parser, *result.addRegion()))
    return failure();
  result.addAttribute(
      "function_type",
      TypeAttr::get(getTargetConditionRegionType(parser.getContext())));
  if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
    return failure();
  return success();
}

void ExecutableConditionOp::print(OpAsmPrinter &p) {
  Operation *op = getOperation();
  printTargetConditionRegion(p, op, getBody());
  p.printOptionalAttrDictWithKeyword(op->getAttrs(),
                                     /*elidedAttrs=*/{"function_type"});
}

Block *ExecutableConditionOp::addEntryBlock() {
  assert(empty() && "function already has an entry block");
  auto *entry = new Block();
  auto argTypes = getArgumentTypes();
  SmallVector<Location> argLocs(argTypes.size(), getLoc());
  entry->addArguments(argTypes, argLocs);
  push_back(entry);
  return entry;
}

Block *ExecutableConditionOp::addBlock() {
  assert(!empty() && "function should at least have an entry block");
  push_back(new Block());
  return &back();
}

//===----------------------------------------------------------------------===//
// hal.executable.constant.block
//===----------------------------------------------------------------------===//

ParseResult ExecutableConstantBlockOp::parse(OpAsmParser &parser,
                                             OperationState &result) {
  auto &builder = parser.getBuilder();

  // Parse the function signature.
  SmallVector<OpAsmParser::Argument> entryArgs;
  bool isVariadic = false;
  SmallVector<DictionaryAttr> resultAttrs;
  SmallVector<Type> resultTypes;
  if (mlir::function_interface_impl::parseFunctionSignature(
          parser, /*allowVariadic=*/false, entryArgs, isVariadic, resultTypes,
          resultAttrs)) {
    return failure();
  }
  SmallVector<Type> argTypes;
  for (auto &arg : entryArgs)
    argTypes.push_back(arg.type);
  auto fnType = builder.getFunctionType(argTypes, resultTypes);
  result.addAttribute(getFunctionTypeAttrName(result.name),
                      TypeAttr::get(fnType));

  // Parse the keys used for each yielded constant value.
  // There must be one key per result. Note that we support omitted parens when
  // only one result is present.
  SmallVector<Attribute> keyAttrs;
  if (failed(parser.parseKeyword("as")))
    return failure();
  if (resultTypes.size() == 1) {
    std::string key;
    if (failed(parser.parseString(&key)))
      return failure();
    keyAttrs.push_back(builder.getStringAttr(key));
  } else {
    if (failed(parser.parseCommaSeparatedList(
            AsmParser::Delimiter::OptionalParen,
            [&]() {
              std::string key;
              if (failed(parser.parseString(&key)))
                return failure();
              keyAttrs.push_back(builder.getStringAttr(key));
              return success();
            },
            "containing a 1:1 list of keys per yielded value"))) {
      return failure();
    }
  }
  result.addAttribute("keys", builder.getArrayAttr(keyAttrs));

  // If function attributes are present, parse them.
  NamedAttrList parsedAttributes;
  if (parser.parseOptionalAttrDictWithKeyword(parsedAttributes)) {
    return failure();
  }
  result.attributes.append(parsedAttributes);

  // Add the attributes to the function arguments.
  assert(resultAttrs.size() == resultTypes.size());
  mlir::function_interface_impl::addArgAndResultAttrs(
      builder, result, entryArgs, resultAttrs, getArgAttrsAttrName(result.name),
      getResAttrsAttrName(result.name));

  // Parse the optional function body. The printer will not print the body if
  // its empty, so disallow parsing of empty body in the parser.
  auto *body = result.addRegion();
  SMLoc loc = parser.getCurrentLocation();
  if (failed(parser.parseRegion(*body, entryArgs,
                                /*enableNameShadowing=*/false))) {
    return failure();
  }
  // Function body was parsed, make sure its not empty.
  if (body->empty()) {
    return parser.emitError(loc, "expected non-empty function body");
  }

  return success();
}

void ExecutableConstantBlockOp::print(OpAsmPrinter &p) {
  Operation *op = getOperation();
  ArrayRef<Type> argTypes = getArgumentTypes();
  ArrayRef<Type> resultTypes = getResultTypes();
  mlir::function_interface_impl::printFunctionSignature(
      p, cast<mlir::FunctionOpInterface>(op), argTypes, /*isVariadic=*/false,
      resultTypes);
  p << " as ";
  if (resultTypes.size() != 1)
    p << '(';
  llvm::interleaveComma(getKeys().getValue(), p,
                        [&](Attribute attr) { p << attr; });
  if (resultTypes.size() != 1)
    p << ')';
  mlir::function_interface_impl::printFunctionAttributes(
      p, op, {getFunctionTypeAttrName(), getKeysAttrName()});
  p << " ";
  p.printRegion(getBody(), /*printEntryBlockArgs=*/false,
                /*printBlockTerminators=*/true);
}

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

  // Verify the function takes either nothing or a device.
  auto argTypes = op.getArgumentTypes();
  if (!argTypes.empty() &&
      (argTypes.size() > 1 || !llvm::isa<IREE::HAL::DeviceType>(argTypes[0]))) {
    return op->emitOpError()
           << "initializer must take a !hal.device or nothing";
  }

  // Verify the return types are all i32 (today).
  for (auto resultType : llvm::enumerate(op.getResultTypes())) {
    if (!resultType.value().isInteger(32)) {
      return op->emitOpError()
             << "initializer must return only i32 values (result "
             << resultType.index() << " is " << resultType.value() << ")";
    }
  }

  // Verify there's a key for every result.
  if (op.getNumResults() != op.getKeys().size()) {
    return op->emitOpError() << "must have one key for every result";
  }

  return success();
}

//===----------------------------------------------------------------------===//
// hal.executable.binary
//===----------------------------------------------------------------------===//

void ExecutableBinaryOp::build(OpBuilder &builder, OperationState &state,
                               StringRef symName, StringRef format,
                               std::vector<uint8_t> data) {
  state.addAttribute(mlir::SymbolTable::getSymbolAttrName(),
                     builder.getStringAttr(symName));
  state.addAttribute("format", builder.getStringAttr(format));
  state.addAttribute("data",
                     DenseIntElementsAttr::get(
                         VectorType::get({static_cast<int64_t>(data.size())},
                                         builder.getIntegerType(8)),
                         data));
}

void ExecutableBinaryOp::build(OpBuilder &builder, OperationState &state,
                               StringRef symName, StringAttr format,
                               DenseIntElementsAttr data) {
  state.addAttribute(mlir::SymbolTable::getSymbolAttrName(),
                     builder.getStringAttr(symName));
  state.addAttribute("format", format);
  state.addAttribute("data", data);
}

//===----------------------------------------------------------------------===//
// hal.executable.create
//===----------------------------------------------------------------------===//

void ExecutableCreateOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  // TODO(benvanik): name after sanitized symbol.
  setNameFn(getResult(), StringRef("executable"));
}

//===----------------------------------------------------------------------===//
// hal.executable.lookup
//===----------------------------------------------------------------------===//

void ExecutableLookupOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  // TODO(benvanik): name after sanitized symbol.
  setNameFn(getResult(), "exe");
}

//===----------------------------------------------------------------------===//
// hal.executable.export.ordinal
//===----------------------------------------------------------------------===//

void ExecutableExportOrdinalOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  // TODO(benvanik): name after sanitized symbol.
  setNameFn(getResult(), "ordinal");
}

//===----------------------------------------------------------------------===//
// hal.interface.constant.load
//===----------------------------------------------------------------------===//

LogicalResult InterfaceConstantLoadOp::verify() {
  InterfaceConstantLoadOp op = *this;
  auto layoutAttr = op.getLayout();
  if (op.getOrdinal().getZExtValue() >= layoutAttr.getConstants()) {
    return op.emitOpError("push constant ordinal out of bounds");
  }
  return success();
}

//===----------------------------------------------------------------------===//
// hal.interface.binding.subspan
//===----------------------------------------------------------------------===//

void InterfaceBindingSubspanOp::build(OpBuilder &builder,
                                      OperationState &result, Type resultType,
                                      IREE::HAL::PipelineLayoutAttr layout,
                                      APInt binding, Value byte_offset,
                                      ValueRange dynamic_dims,
                                      IntegerAttr alignment,
                                      std::optional<DescriptorFlags> flags) {
  IREE::HAL::DescriptorFlagsAttr descriptorAttr;
  if (flags.has_value()) {
    descriptorAttr = IREE::HAL::DescriptorFlagsAttr::get(builder.getContext(),
                                                         flags.value());
  }
  build(builder, result, resultType, layout, binding, byte_offset, dynamic_dims,
        alignment, descriptorAttr);
}

LogicalResult InterfaceBindingSubspanOp::verify() {
  InterfaceBindingSubspanOp op = *this;
  if (ShapedType shapedType = llvm::dyn_cast<ShapedType>(op.getType())) {
    if (shapedType.getNumDynamicDims() != op.getDynamicDims().size()) {
      return op.emitOpError("result type ")
             << op.getType() << " has " << shapedType.getNumDynamicDims()
             << " dynamic dimensions but " << op.getDynamicDims().size()
             << " associated dimension SSA values";
    }
  }
  uint64_t binding = op.getBinding().getZExtValue();
  if (binding >= op.getLayout().getBindings().size()) {
    return op.emitOpError("binding ordinal ")
           << binding << " out of bounds in layout " << op.getLayout();
  }
  return success();
}

IREE::HAL::PipelineBindingAttr
InterfaceBindingSubspanOp::getPipelineBindingAttr() {
  return getLayout().getBinding(getBinding());
}

IREE::HAL::DescriptorType InterfaceBindingSubspanOp::getDescriptorType() {
  auto bindingAttr = getPipelineBindingAttr();
  return bindingAttr.getType();
}

llvm::MaybeAlign InterfaceBindingSubspanOp::getBaseAlignment() {
  if (auto baseAlignmentInt = getAlignment()) {
    return llvm::MaybeAlign(baseAlignmentInt.value().getZExtValue());
  }
  return std::nullopt;
}

llvm::Align InterfaceBindingSubspanOp::calculateAlignment() {
  // If we can't calculate an alignment we fall back to the natural alignment of
  // the element type (for example, a memref<?xi32> is known to be at least
  // 4-byte aligned).
  llvm::Align naturalAlignment(1);
  auto resultType = getType();
  if (auto shapedType = llvm::dyn_cast<ShapedType>(resultType)) {
    naturalAlignment = llvm::Align(
        IREE::Util::getRoundedElementByteWidth(shapedType.getElementType()));
  }

  // If the binding has no assigned alignment we fall back to natural alignment.
  auto baseAlignment = getBaseAlignment();
  if (!baseAlignment)
    return naturalAlignment;

  // If there's no offset specified then we can use the binding alignment
  // directly.
  if (!getByteOffset())
    return baseAlignment.value();

  // Try to get the alignment of the byte offset. If it's a constant then we can
  // find a common alignment between it and the base and otherwise we need to
  // try to infer the alignment from the IR - otherwise we fall back.
  auto offsetOrAlignment = lookupOffsetOrAlignment(getByteOffset());
  if (!offsetOrAlignment.has_value())
    return naturalAlignment;

  // Compute the common alignment between that of the binding base and that of
  // the byte offset.
  return llvm::commonAlignment(baseAlignment.value(),
                               offsetOrAlignment.value());
}

LogicalResult InterfaceBindingSubspanOp::reifyResultShapes(
    OpBuilder &builder, ReifiedRankedShapedTypeDims &reifiedReturnShapes) {
  auto resultShapedType = dyn_cast<ShapedType>(getResult().getType());
  if (!resultShapedType) {
    return failure();
  }
  SmallVector<OpFoldResult> resultShape = mlir::getMixedValues(
      resultShapedType.getShape(), getDynamicDims(), builder);
  reifiedReturnShapes.emplace_back(std::move(resultShape));
  return success();
}

//===----------------------------------------------------------------------===//
// hal.interface.workgroup.*
//===----------------------------------------------------------------------===//

static void getAsmResultNamesForInterfaceWorkgroupOp(
    StringRef prefix, const APInt &dimension, Value result,
    function_ref<void(Value, StringRef)> setNameFn) {
  switch (dimension.getZExtValue()) {
  case 0:
    setNameFn(result, (prefix + "x").str());
    return;
  case 1:
    setNameFn(result, (prefix + "y").str());
    return;
  case 2:
    setNameFn(result, (prefix + "z").str());
    return;
  }
}

// Minimum is the smallest possible result we could get. It's 0 for ID-like
// operations and 1 for count-like ones.
static void setResultRangesForInterfaceWorkgroupOp(
    Value result, const std::optional<APInt> &upperBound,
    SetIntRangeFn setResultRanges, int64_t minimum) {
  unsigned width = ConstantIntRanges::getStorageBitwidth(result.getType());
  if (!upperBound.has_value()) {
    setResultRanges(
        result, ConstantIntRanges::fromSigned(APInt(width, minimum),
                                              APInt::getSignedMaxValue(width)));
    return;
  }
  setResultRanges(result,
                  ConstantIntRanges::fromUnsigned(APInt(width, minimum),
                                                  *upperBound + minimum - 1));
}

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

void InterfaceWorkgroupIDOp::inferResultRanges(
    ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRanges) {
  setResultRangesForInterfaceWorkgroupOp(getResult(), getUpperBound(),
                                         setResultRanges, /*minimum=*/0);
}

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

void InterfaceWorkgroupCountOp::inferResultRanges(
    ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRanges) {
  setResultRangesForInterfaceWorkgroupOp(getResult(), getUpperBound(),
                                         setResultRanges, /*minimum=*/1);
}

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

void InterfaceWorkgroupSizeOp::inferResultRanges(
    ArrayRef<ConstantIntRanges> argRanges, SetIntRangeFn setResultRanges) {
  setResultRangesForInterfaceWorkgroupOp(getResult(), getUpperBound(),
                                         setResultRanges, /*minimum=*/1);
}

//===----------------------------------------------------------------------===//
// hal.fence.*
//===----------------------------------------------------------------------===//

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

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

void FenceAwaitOp::getAsmResultNames(
    function_ref<void(Value, StringRef)> setNameFn) {
  setNameFn(getStatus(), "status");
}

} // namespace mlir::iree_compiler::IREE::HAL

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

#define GET_OP_CLASSES
#include "iree/compiler/Dialect/HAL/IR/HALOps.cpp.inc"
