// 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/Util/IR/UtilOps.h"
#include "iree/compiler/Dialect/Util/IR/UtilTypes.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Support/LogicalResult.h"

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

//===----------------------------------------------------------------------===//
// Utilities
//===----------------------------------------------------------------------===//

namespace {

// Erases an op if it has no uses.
// This is to support ops that are "pure" but can't be marked as such because
// the MLIR CSE pass would deduplicate them.
template <typename Op>
struct ElideUnusedOp : OpRewritePattern<Op> {
  explicit ElideUnusedOp(MLIRContext *context)
      : OpRewritePattern<Op>(context, /*benefit=*/1000) {}
  LogicalResult matchAndRewrite(Op op,
                                PatternRewriter &rewriter) const override {
    if (!op.use_empty()) {
      return failure();
    }
    rewriter.eraseOp(op);
    return success();
  }
};

} // namespace

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

OpFoldResult TensorImportOp::fold(FoldAdaptor operands) {
  // Wait fences and consumes carry availability/ownership effects for the
  // imported tensor value. Folding through them would drop those effects.
  if (getWaitFence() || getConsume()) {
    return {};
  }
  // Cannot fold if there is a byte offset — the import is a subview.
  if (getByteOffset()) {
    return {};
  }
  if (auto exportOp = getSource().getDefiningOp<TensorExportOp>()) {
    if (exportOp.getSource().getType() == getTarget().getType() &&
        exportOp.getSourceEncoding() == getTargetEncoding()) {
      return exportOp.getSource();
    }
  }
  return {};
}

namespace {

static bool isConstantValue(Value value) {
  Attribute constantAttr;
  return matchPattern(value, m_Constant(&constantAttr));
}

static bool tryMatchConstantValue(Value value, Attribute &constantAttr) {
  return matchPattern(value, m_Constant(&constantAttr));
}

static Value getConstantAvailableAt(PatternRewriter &rewriter,
                                    TensorImportOp importOp, Value value) {
  Operation *constantOp = value.getDefiningOp();
  assert(constantOp && "expected constant to have defining op");
  if (constantOp->getBlock() == importOp->getBlock() &&
      constantOp->isBeforeInBlock(importOp)) {
    return value;
  }

  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPoint(importOp);
  Operation *clonedOp = rewriter.clone(*constantOp);
  return clonedOp->getResult(0);
}

struct DimEvidence {
  // First constant dimension value observed from shape-aware consumers.
  Value constantValue;
  // Constant attribute used to compare equivalent dimension values.
  Attribute constantAttr;
  // Whether any shape-aware consumer provided a constant dimension.
  bool hasConstant = false;
  // Whether any shape-aware consumer kept the dimension dynamic.
  bool hasNonConstant = false;
  // Whether shape-aware consumers provided different constant dimensions.
  bool hasConflict = false;
};

struct PropagateImportDimsFromConsumers
    : public OpRewritePattern<TensorImportOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(TensorImportOp importOp,
                                PatternRewriter &rewriter) const override {
    ValueRange targetDims = importOp.getTargetDims();
    if (targetDims.empty()) {
      return failure();
    }
    if (llvm::all_of(targetDims, isConstantValue)) {
      return failure();
    }

    SmallVector<DimEvidence> evidence(targetDims.size());
    for (OpOperand &use : importOp.getTarget().getUses()) {
      Operation *userOp = use.getOwner();
      if (userOp->getBlock() != importOp->getBlock()) {
        // Region-local users may only refine the imported shape on one
        // control-flow path. Keep the import dims unchanged unless all evidence
        // comes from the import's block.
        return failure();
      }

      auto shapeAwareUser = dyn_cast<IREE::Util::ShapeAwareOpInterface>(userOp);
      if (!shapeAwareUser) {
        continue;
      }

      ValueRange userDims =
          shapeAwareUser.getOperandDynamicDims(use.getOperandNumber());
      if (userDims.size() != targetDims.size()) {
        return failure();
      }

      for (auto [index, userDim] : llvm::enumerate(userDims)) {
        if (isConstantValue(targetDims[index])) {
          continue;
        }

        Attribute constantAttr;
        if (!tryMatchConstantValue(userDim, constantAttr)) {
          evidence[index].hasNonConstant = true;
          continue;
        }

        DimEvidence &dimEvidence = evidence[index];
        if (!dimEvidence.hasConstant) {
          dimEvidence.constantValue = userDim;
          dimEvidence.constantAttr = constantAttr;
          dimEvidence.hasConstant = true;
        } else if (dimEvidence.constantAttr != constantAttr) {
          dimEvidence.hasConflict = true;
        }
      }
    }

    bool anyResolved = false;
    SmallVector<Value> resolvedDims(targetDims.begin(), targetDims.end());
    for (auto [index, dimEvidence] : llvm::enumerate(evidence)) {
      if (isConstantValue(targetDims[index]) || !dimEvidence.hasConstant ||
          dimEvidence.hasNonConstant || dimEvidence.hasConflict) {
        continue;
      }
      resolvedDims[index] =
          getConstantAvailableAt(rewriter, importOp, dimEvidence.constantValue);
      anyResolved = true;
    }
    if (!anyResolved) {
      return failure();
    }

    rewriter.modifyOpInPlace(importOp, [&]() {
      importOp.getTargetDimsMutable().assign(resolvedDims);
    });
    return success();
  }
};

} // namespace

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

OpFoldResult TensorExportOp::fold(FoldAdaptor operands) {
  if (auto importOp = getSource().getDefiningOp<TensorImportOp>()) {
    // Wait fences and consumes carry availability/ownership effects for the
    // imported tensor value. Folding through them would drop those effects.
    if (importOp.getWaitFence() || importOp.getConsume()) {
      return {};
    }
    // Cannot fold through an import with byte_offset — it's a subview.
    if (importOp.getByteOffset()) {
      return {};
    }
    if (importOp.getSource().getType() == getTarget().getType() &&
        importOp.getTargetEncoding() == getSourceEncoding()) {
      return importOp.getSource();
    }
  }
  return {};
}

//===----------------------------------------------------------------------===//
// hal.tensor.transients
//===----------------------------------------------------------------------===//

namespace {

// Folds consecutive transients ops into one.
// %1 = hal.tensor.transients %0 : tensor<?xf32>{%dim} from %storage1
// %2 = hal.tensor.transients %1 : tensor<?xf32>{%dim} from %storage2
// =>
// %2 = hal.tensor.transients %0 : tensor<?xf32>{%dim} from %storage2
struct FoldConsecutiveTransientsOps : OpRewritePattern<TensorTransientsOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(TensorTransientsOp op,
                                PatternRewriter &rewriter) const override {
    // Check if the source is another transients op.
    auto sourceOp = op.getSource().getDefiningOp<TensorTransientsOp>();
    if (!sourceOp) {
      return failure();
    }

    // Fold by using the original source from the inner transients op,
    // but keeping the outer storage (outer annotation wins).
    rewriter.modifyOpInPlace(op, [&]() {
      op.getSourceMutable().assign(sourceOp.getSource());
      op.getSourceDimsMutable().assign(sourceOp.getSourceDims());
    });
    return success();
  }
};

} // namespace

void TensorTransientsOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                     MLIRContext *context) {
  results.add<FoldConsecutiveTransientsOps>(context);
}

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

namespace {

/// Deduplicates hal.tensor.barrier operands.
struct DeduplicateTensorBarrierSources : OpRewritePattern<TensorBarrierOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(TensorBarrierOp op,
                                PatternRewriter &rewriter) const override {
    DenseMap<Value, unsigned> uniqueSources; // source -> unique index
    SmallVector<Value> orderedSources;
    SmallVector<unsigned> resultMapping; // old -> new result index
    for (auto source : op.getSources()) {
      auto it =
          uniqueSources.insert(std::make_pair(source, orderedSources.size()));
      if (it.second) {
        orderedSources.push_back(source);
      }
      resultMapping.push_back(it.first->second);
    }
    if (orderedSources.size() == op.getSources().size()) {
      return failure();
    }
    auto newOp = TensorBarrierOp::create(rewriter, op.getLoc(), orderedSources,
                                         op.getSignalFence());
    SmallVector<Value> newResults;
    newResults.reserve(newOp.getNumResults());
    for (unsigned newIndex : resultMapping) {
      newResults.push_back(newOp.getResult(newIndex));
    }
    rewriter.replaceOp(op, newResults);
    return success();
  }
};

} // namespace

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

//===----------------------------------------------------------------------===//
// hal.allocator.*
//===----------------------------------------------------------------------===//

namespace {

/// Tries to fold either the device or queue affinity of a select when all
/// potential values of either match.
struct FoldAllocatorSelect : OpRewritePattern<AllocatorSelectOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(AllocatorSelectOp op,
                                PatternRewriter &rewriter) const override {
    // Calculate the unique set of devices and unique set of queue affinities.
    DenseSet<Value> devices;
    DenseSet<Value> queueAffinities;
    for (auto [device, queueAffinity] :
         llvm::zip_equal(op.getDevices(), op.getQueueAffinities())) {
      devices.insert(device);
      queueAffinities.insert(queueAffinity);
    }
    if (devices.size() == 1 && queueAffinities.size() == 1) {
      // Only a single selected combination.
      Value commonDevice = *devices.begin();
      Value commonQueueAffinityMask = *queueAffinities.begin();
      rewriter.replaceOp(op, {commonDevice, commonQueueAffinityMask});
      return success();
    } else if (devices.size() == 1 && !op.getSelectedDevice().use_empty()) {
      // Device is common but queue masks are not.
      // This is unlikely to arise naturally.
      Value commonDevice = *devices.begin();
      rewriter.replaceAllUsesWith(op.getSelectedDevice(), commonDevice);
      return success();
    } else if (queueAffinities.size() == 1 &&
               !op.getSelectedQueueAffinity().use_empty()) {
      // Queue affinities are common (likely "any") but devices are not.
      // We can fold the queue affinity mask to a constant now to allow it to
      // propagate.
      Value commonQueueAffinityMask = *queueAffinities.begin();
      rewriter.replaceAllUsesWith(op.getSelectedQueueAffinity(),
                                  commonQueueAffinityMask);
      return success();
    }
    return failure();
  }
};

} // namespace

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

//===----------------------------------------------------------------------===//
// hal.buffer_view.*
//===----------------------------------------------------------------------===//

namespace {

/// Folds hal.buffer.subspans into buffer view creation subspans.
struct FoldBufferViewCreateSubspan : OpRewritePattern<BufferViewCreateOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(BufferViewCreateOp op,
                                PatternRewriter &rewriter) const override {
    auto ip = rewriter.saveInsertionPoint();
    rewriter.setInsertionPoint(op);
    bool needsUpdate = false;
    auto newSourceBuffer = op.getSourceBuffer();
    auto newSourceOffset = cast<Value>(op.getSourceOffset());
    if (auto subspanOp = dyn_cast_if_present<IREE::HAL::BufferSubspanOp>(
            op.getSourceBuffer().getDefiningOp())) {
      newSourceBuffer = subspanOp.getSourceBuffer();
      newSourceOffset = rewriter.createOrFold<arith::AddIOp>(
          subspanOp.getLoc(), subspanOp.getSourceOffset(),
          op.getSourceOffset());
      needsUpdate = true;
    }
    rewriter.restoreInsertionPoint(ip);
    if (!needsUpdate) {
      return failure();
    }
    rewriter.modifyOpInPlace(op, [&]() {
      op.getSourceBufferMutable().assign(newSourceBuffer);
      op.getSourceOffsetMutable().assign(newSourceOffset);
    });
    return success();
  }
};

} // namespace

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

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

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

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

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

//===----------------------------------------------------------------------===//
// hal.command_buffer.*
//===----------------------------------------------------------------------===//

namespace {

/// Skips a hal.command_buffer.device accessor when the device was created in
/// the same scope.
struct SkipCommandBufferDeviceOp : OpRewritePattern<CommandBufferDeviceOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(CommandBufferDeviceOp op,
                                PatternRewriter &rewriter) const override {
    if (auto createOp = dyn_cast_if_present<CommandBufferCreateOp>(
            op.getCommandBuffer().getDefiningOp())) {
      rewriter.replaceOp(op, createOp.getDevice());
      return success();
    }
    return failure();
  }
};

} // namespace

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

namespace {

/// Folds hal.buffer.subspans into buffer fill offsets.
struct FoldCommandBufferFillBufferSubspans
    : OpRewritePattern<CommandBufferFillBufferOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(CommandBufferFillBufferOp op,
                                PatternRewriter &rewriter) const override {
    auto ip = rewriter.saveInsertionPoint();
    rewriter.setInsertionPoint(op);
    bool needsUpdate = false;
    auto newTargetBuffer = op.getTargetBuffer();
    auto newTargetOffset = cast<Value>(op.getTargetOffset());
    if (auto subspanOp = dyn_cast_if_present<IREE::HAL::BufferSubspanOp>(
            op.getTargetBuffer().getDefiningOp())) {
      newTargetBuffer = subspanOp.getSourceBuffer();
      newTargetOffset = rewriter.createOrFold<arith::AddIOp>(
          subspanOp.getLoc(), subspanOp.getSourceOffset(),
          op.getTargetOffset());
      needsUpdate = true;
    }
    rewriter.restoreInsertionPoint(ip);
    if (!needsUpdate) {
      return failure();
    }
    rewriter.modifyOpInPlace(op, [&]() {
      op.getTargetBufferMutable().assign(newTargetBuffer);
      op.getTargetOffsetMutable().assign(newTargetOffset);
    });
    return success();
  }
};

} // namespace

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

namespace {

/// Folds hal.buffer.subspans into buffer update offsets.
struct FoldCommandBufferUpdateBufferSubspans
    : OpRewritePattern<CommandBufferUpdateBufferOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(CommandBufferUpdateBufferOp op,
                                PatternRewriter &rewriter) const override {
    auto ip = rewriter.saveInsertionPoint();
    rewriter.setInsertionPoint(op);
    bool needsUpdate = false;
    auto newTargetBuffer = op.getTargetBuffer();
    auto newTargetOffset = cast<Value>(op.getTargetOffset());
    if (auto subspanOp = dyn_cast_if_present<IREE::HAL::BufferSubspanOp>(
            op.getTargetBuffer().getDefiningOp())) {
      newTargetBuffer = subspanOp.getSourceBuffer();
      newTargetOffset = rewriter.createOrFold<arith::AddIOp>(
          subspanOp.getLoc(), subspanOp.getSourceOffset(),
          op.getTargetOffset());
      needsUpdate = true;
    }
    rewriter.restoreInsertionPoint(ip);
    if (!needsUpdate) {
      return failure();
    }
    rewriter.modifyOpInPlace(op, [&]() {
      op.getTargetBufferMutable().assign(newTargetBuffer);
      op.getTargetOffsetMutable().assign(newTargetOffset);
    });
    return success();
  }
};

} // namespace

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

namespace {

/// Folds hal.buffer.subspans into buffer copy offsets.
struct FoldCommandBufferCopyBufferSubspans
    : OpRewritePattern<CommandBufferCopyBufferOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(CommandBufferCopyBufferOp op,
                                PatternRewriter &rewriter) const override {
    auto ip = rewriter.saveInsertionPoint();
    rewriter.setInsertionPoint(op);
    bool needsUpdate = false;
    auto newSourceBuffer = op.getSourceBuffer();
    auto newSourceOffset = cast<Value>(op.getSourceOffset());
    if (auto subspanOp = dyn_cast_if_present<IREE::HAL::BufferSubspanOp>(
            op.getSourceBuffer().getDefiningOp())) {
      newSourceBuffer = subspanOp.getSourceBuffer();
      newSourceOffset = rewriter.createOrFold<arith::AddIOp>(
          subspanOp.getLoc(), subspanOp.getSourceOffset(),
          op.getSourceOffset());
      needsUpdate = true;
    }
    auto newTargetBuffer = op.getTargetBuffer();
    auto newTargetOffset = cast<Value>(op.getTargetOffset());
    if (auto subspanOp = dyn_cast_if_present<IREE::HAL::BufferSubspanOp>(
            op.getTargetBuffer().getDefiningOp())) {
      newTargetBuffer = subspanOp.getSourceBuffer();
      newTargetOffset = rewriter.createOrFold<arith::AddIOp>(
          subspanOp.getLoc(), subspanOp.getSourceOffset(),
          op.getTargetOffset());
      needsUpdate = true;
    }
    rewriter.restoreInsertionPoint(ip);
    if (!needsUpdate) {
      return failure();
    }
    rewriter.modifyOpInPlace(op, [&]() {
      op.getSourceBufferMutable().assign(newSourceBuffer);
      op.getSourceOffsetMutable().assign(newSourceOffset);
      op.getTargetBufferMutable().assign(newTargetBuffer);
      op.getTargetOffsetMutable().assign(newTargetOffset);
    });
    return success();
  }
};

} // namespace

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

namespace {

/// Folds hal.buffer.subspans into dispatch bindings.
/// The binding range is always equal to or a subset of the subspan.
template <typename OpT>
struct FoldCommandBufferDispatchBufferSubspan : OpRewritePattern<OpT> {
  using OpRewritePattern<OpT>::OpRewritePattern;

  LogicalResult matchAndRewrite(OpT op,
                                PatternRewriter &rewriter) const override {
    auto ip = rewriter.saveInsertionPoint();
    rewriter.setInsertionPoint(op);
    bool needsUpdate = false;
    auto bindingBuffers = llvm::to_vector(op.getBindingBuffers());
    auto bindingOffsets = llvm::to_vector(op.getBindingOffsets());
    for (size_t i = 0; i < bindingBuffers.size(); ++i) {
      auto *definingOp = bindingBuffers[i].getDefiningOp();
      if (!definingOp) {
        continue;
      }
      if (auto subspanOp = dyn_cast<IREE::HAL::BufferSubspanOp>(definingOp)) {
        needsUpdate = true;
        bindingBuffers[i] = subspanOp.getSourceBuffer();
        bindingOffsets[i] = rewriter.createOrFold<arith::AddIOp>(
            subspanOp.getLoc(), subspanOp.getSourceOffset(), bindingOffsets[i]);
      }
    }
    rewriter.restoreInsertionPoint(ip);
    if (!needsUpdate) {
      return failure();
    }
    rewriter.modifyOpInPlace(op, [&]() {
      auto mutableBindingBuffers = op.getBindingBuffersMutable();
      mutableBindingBuffers.clear();
      mutableBindingBuffers.append(bindingBuffers);
      auto mutableBindingOffsets = op.getBindingOffsetsMutable();
      mutableBindingOffsets.clear();
      mutableBindingOffsets.append(bindingOffsets);
    });
    return success();
  }
};

} // namespace

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

namespace {

/// Folds hal.buffer.subspans into the indirect dispatch workgroup count.
/// The binding range is always equal to or a subset of the subspan.
struct FoldCommandBufferDispatchIndirectBufferSubspan
    : OpRewritePattern<CommandBufferDispatchIndirectOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(CommandBufferDispatchIndirectOp op,
                                PatternRewriter &rewriter) const override {
    Value workgroupsBuffer = op.getWorkgroupsBuffer();
    auto *definingOp = workgroupsBuffer.getDefiningOp();
    if (!definingOp) {
      return failure();
    }
    Value workgroupsOffset = op.getWorkgroupsOffset();
    if (auto subspanOp = dyn_cast<IREE::HAL::BufferSubspanOp>(definingOp)) {
      workgroupsBuffer = subspanOp.getSourceBuffer();
      workgroupsOffset = rewriter.createOrFold<arith::AddIOp>(
          subspanOp.getLoc(), subspanOp.getSourceOffset(), workgroupsOffset);
    } else {
      return failure();
    }
    rewriter.modifyOpInPlace(op, [&]() {
      op.getWorkgroupsBufferMutable().set(workgroupsBuffer);
      op.getWorkgroupsOffsetMutable().set(workgroupsOffset);
    });
    return success();
  }
};

} // namespace

void CommandBufferDispatchIndirectOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.insert<FoldCommandBufferDispatchIndirectBufferSubspan>(context);
  results.insert<
      FoldCommandBufferDispatchBufferSubspan<CommandBufferDispatchIndirectOp>>(
      context);
}

//===----------------------------------------------------------------------===//
// hal.device.queue.execute
//===----------------------------------------------------------------------===//

// Returns true if |before| is always executed by the time |after| is reached.
// NOTE: this is currently very conservative and only looks for ops in the
// same basic block. We need an abstract interpreter to do much more as we'd
// need to track conditionals/branching logic.
static bool isOpAlwaysExecutedWith(Operation *before, Operation *after) {
  if (before == after) {
    return true;
  }
  if (before->getBlock() != after->getBlock()) {
    return false;
  }
  return before->isBeforeInBlock(after);
}

// Returns true if |op| was hoisted before |insertBefore| without breaking
// SSA invariants. Returns false if no IR modifications were made.
static bool tryHoistOpBeforeUser(Operation *op, Operation *insertBefore) {
  if (op == insertBefore) {
    return false;
  }

  // Currently conservative - should be doing a domination check.
  if (op->getBlock() != insertBefore->getBlock()) {
    // Today only doing within the same block.
    return false;
  }

  // Ensure all operands are defined above the insertion target.
  // TODO(benvanik): hoist dependent ops too (constants are common).
  if (!llvm::all_of(op->getOperands(), [&](Value operand) {
        auto *definingOp = operand.getDefiningOp();
        if (!definingOp || definingOp->getBlock() != insertBefore->getBlock()) {
          // Function/block args or values defined outside the insertion block
          // are ok since we are limiting to 1 block.
          return true;
        }
        return definingOp->isBeforeInBlock(insertBefore);
      })) {
    return false;
  }

  // Should be safe to hoist the op 🤞.
  op->moveBefore(insertBefore);
  return true;
}

namespace {

/// Swaps a device queue barrier with an immediate host fence signal when the
/// wait fence is immediately resolved (null).
struct ImmediatelyResolveDeviceQueueBarrier
    : OpRewritePattern<DeviceQueueBarrierOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(DeviceQueueBarrierOp barrierOp,
                                PatternRewriter &rewriter) const override {
    // Check for whether we know the wait fence is immediately resolved in the
    // local scope. A more involved data flow analysis would let us handle more
    // cases (function calls, block edges, etc) that commonly arise.
    if (!isa_and_nonnull<IREE::Util::NullOp>(
            barrierOp.getWaitFence().getDefiningOp())) {
      return failure();
    }

    rewriter.replaceOpWithNewOp<IREE::HAL::FenceSignalOp>(
        barrierOp, barrierOp.getSignalFence());
    return success();
  }
};

/// Aliases a signal fence to a wait fence when there's a direct execution
/// dependency through the barrier. This only checks the local scope but could
/// be extended across CFG boundaries.
///
/// Example:
///  %fence0 = hal.fence.create
///  hal.device.queue.barrier signal(%fence0)
///  hal.device.queue.execute wait(%fence0) signal(%fence1)
/// ->
///  hal.device.queue.barrier signal(%fence1)
struct HoistDeviceQueueBarrierChain : OpRewritePattern<DeviceQueueBarrierOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(DeviceQueueBarrierOp barrierOp,
                                PatternRewriter &rewriter) const override {
    // See if we can observe the original fence creation in the local scope.
    auto waitFence = barrierOp.getWaitFence();
    auto createOp = dyn_cast_if_present<IREE::HAL::FenceCreateOp>(
        waitFence.getDefiningOp());
    if (!createOp) {
      return rewriter.notifyMatchFailure(barrierOp,
                                         "cannot analyze wait fence creation");
    }

    // Today this simple pattern only deals with the local block. We should
    // extend this to support a must-be-executed context such that we can deal
    // with the common case of some basic control flow handling errors/etc.
    if (createOp->getBlock() != barrierOp->getBlock()) {
      return rewriter.notifyMatchFailure(
          barrierOp,
          "create and barrier are in different blocks; analysis TBD");
    }

    // To ensure we don't break SSA invariants we need to only hoist if the
    // signal fence is or can be defined before all users of the waitFence we
    // are replacing. Note that because we are only matching on ops within
    // the same block if we don't have the defining op it means it's a block
    // argument and is always available.
    auto signalFence = barrierOp.getSignalFence();
    auto signalDefiningOp = signalFence.getDefiningOp();
    if (signalDefiningOp) {
      // Try to hoist up to the defining op.
      if (!tryHoistOpBeforeUser(signalDefiningOp, createOp)) {
        return rewriter.notifyMatchFailure(
            barrierOp, "signal defining op cannot be hoisted");
      }
    }

    // Replace the original fence with the new one and drop the create.
    rewriter.replaceAllUsesWith(waitFence, signalFence);
    rewriter.eraseOp(createOp);

    // Drop the barrier now that it is a no-op.
    rewriter.eraseOp(barrierOp);

    return success();
  }
};

/// Elides queue barriers that are used for sequencing fences when the operation
/// could be performed by way of the originating queue operation.
///
/// Example (where %b is only used by the two ops):
///  hal.device.queue.execute wait(%a) signal(%b) commands(...)
///  hal.device.queue.barrier wait(%b) signal(%c)
/// ->
///  hal.device.queue.execute wait(%a) signal(%c) commands(...)
struct ElideDeviceQueueBarrierOp : OpRewritePattern<DeviceQueueBarrierOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(DeviceQueueBarrierOp barrierOp,
                                PatternRewriter &rewriter) const override {

    // We're looking at the wait fence on the barrier back up to the signal
    // operation on that fence.
    Value oldFence = barrierOp.getWaitFence();
    Value newFence = barrierOp.getSignalFence();

    // If the only two users of the fence are the signaling op and this barrier
    // then we know we can drop the fence. We look through the users and expect
    // to find only one other: the op that signals the fence. Just because we
    // find it doesn't mean that it's been executed yet, though.
    Operation *signalingOp = nullptr;
    for (auto &use : oldFence.getUses()) {
      if (use.getOwner() == barrierOp) {
        // This barrier op.
        continue;
      } else if (isSignalingOp(use.getOwner(), oldFence)) {
        // Found an op signaling the fence.
        // This may be one of many, though, so we keep scanning.
        signalingOp = use.getOwner();
        continue;
      }
      // Found some other user that is not the signaling op - we bail for now.
      return rewriter.notifyMatchFailure(barrierOp,
                                         "other fence user preventing elision");
    }
    if (!signalingOp || !isOpAlwaysExecutedWith(signalingOp, barrierOp)) {
      // Unable to find the signaling op or it isn't known to have been executed
      // prior to this op (control flow/etc).
      return rewriter.notifyMatchFailure(barrierOp, "signaling op not found");
    }

    rewriter.startOpModification(signalingOp);

    // Try to move the fence producer before the signaling op. This will fail if
    // the op creating the fence has dependencies with hazards.
    if (!IREE::Util::tryMoveProducerBefore(newFence, signalingOp)) {
      rewriter.cancelOpModification(signalingOp);
      return rewriter.notifyMatchFailure(barrierOp,
                                         "fence is not usable by signaling op");
    }

    // Rewrite the signaling op to signal the barrier fence.
    if (failed(updateOpToSignalFence(signalingOp, newFence))) {
      rewriter.cancelOpModification(signalingOp);
      return rewriter.notifyMatchFailure(barrierOp,
                                         "unrecognized signaling op");
    }

    rewriter.finalizeOpModification(signalingOp);

    // Elide the barrier. The fence should be cleaned up as part of DCE.
    rewriter.eraseOp(barrierOp);

    return success();
  }

  // Returns true if |op| signals |fence|.
  static bool isSignalingOp(Operation *op, Value fence) {
    if (auto queueOp = dyn_cast<IREE::HAL::DeviceQueueOpInterface>(op)) {
      return queueOp.getQueueSignalFence() == fence;
    }
    return false;
  }

  // Updates |op| to signal |fence|.
  static LogicalResult updateOpToSignalFence(Operation *op, Value fence) {
    if (auto queueOp = dyn_cast<IREE::HAL::DeviceQueueOpInterface>(op)) {
      queueOp.setQueueSignalFence(fence);
      return success();
    }
    return failure();
  }
};

} // namespace

void DeviceQueueBarrierOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.insert<ImmediatelyResolveDeviceQueueBarrier>(context);
  results.insert<HoistDeviceQueueBarrierChain>(context);
  results.insert<ElideDeviceQueueBarrierOp>(context);
}

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

namespace {

/// Removes the condition region and fallback from an export that will always be
/// selected. This happens if the condition region is folded using a specialized
/// target environment that allows for compile-time query evaluation.
struct DropTrueConditionRegion : OpRewritePattern<ExecutableExportOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(ExecutableExportOp exportOp,
                                PatternRewriter &rewriter) const override {
    auto *block = exportOp.getConditionBody();
    if (!block) {
      return rewriter.notifyMatchFailure(exportOp, "no condition region");
    }

    auto returnOp =
        dyn_cast_if_present<IREE::HAL::ReturnOp>(block->getTerminator());
    if (!returnOp) {
      return rewriter.notifyMatchFailure(exportOp,
                                         "invalid condition terminator");
    }

    Value conditionResult = returnOp.getOperand(0);
    if (!matchPattern(conditionResult, m_One())) {
      return rewriter.notifyMatchFailure(exportOp,
                                         "non-constant true condition result");
    }

    rewriter.eraseBlock(block);
    rewriter.modifyOpInPlace(exportOp,
                             [&]() { exportOp.removeConditionFallbackAttr(); });
    return success();
  }
};

} // namespace

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

namespace {

// Returns a set of fused locations for each result from all return sites.
static SmallVector<Location> gatherResultLocations(int numResults,
                                                   Region &region) {
  SmallVector<SmallVector<Location>> allLocs;
  allLocs.resize(numResults);
  for (auto returnOp : region.getOps<IREE::HAL::ReturnOp>()) {
    for (auto [i, result] : llvm::enumerate(returnOp.getOperands())) {
      allLocs[i].push_back(result.getLoc());
    }
  }
  return llvm::map_to_vector(allLocs, [&](auto resultLocs) {
    return FusedLoc::get(region.getContext(), resultLocs);
  });
}

// Rewrites |region| to have a single hal.return with all prior return sites
// branching to it. Upon return the exit block may not be the last!
static void rewriteToOneReturn(int numResults, Region &region,
                               PatternRewriter &rewriter) {
  // Get all of the return ops - if there's only one then the requirement is
  // already satisfied and we can exit early.
  auto returnOps = llvm::to_vector(region.getOps<IREE::HAL::ReturnOp>());
  if (returnOps.size() <= 1) {
    return; // no-op
  }
  SmallVector<Location> returnLocs;
  for (auto returnOp : returnOps) {
    returnLocs.push_back(returnOp.getLoc());
  }

  // Create the new exit block with arguments matching 1:1 with results.
  auto anyReturnOp = returnOps.front();
  auto resultLocs = gatherResultLocations(anyReturnOp.getNumOperands(), region);
  auto &exitBlock = region.emplaceBlock();
  exitBlock.addArguments(anyReturnOp.getOperandTypes(), resultLocs);
  OpBuilder builder = OpBuilder::atBlockBegin(&exitBlock);
  IREE::HAL::ReturnOp::create(builder,
                              FusedLoc::get(region.getContext(), returnLocs),
                              exitBlock.getArguments());

  // Rewrite all return ops to branch to the exit block.
  for (auto returnOp : returnOps) {
    OpBuilder builder(returnOp);
    cf::BranchOp::create(builder, returnOp.getLoc(), &exitBlock,
                         returnOp.getOperands());
    rewriter.eraseOp(returnOp);
  }
}

/// Merges hal.executable.constant.block ops together into one.
/// Duplicate keys are ignored and will be cleaned up by
/// DeduplicateExecutableConstantBlockKeys.
struct MergeExecutableConstantBlocks : OpRewritePattern<ExecutableVariantOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(ExecutableVariantOp variantOp,
                                PatternRewriter &rewriter) const override {
    auto blockOps = llvm::to_vector(variantOp.getConstantBlockOps());
    if (blockOps.size() <= 1) {
      return rewriter.notifyMatchFailure(variantOp,
                                         "not enough blocks to merge");
    }

    rewriter.startOpModification(variantOp);

    // Gather all constants initialized by the blocks.
    SmallVector<Location> blockLocs;
    bool anyRequireDevice = false;
    SmallVector<Type> resultTypes;
    SmallVector<Attribute> resultKeys;
    SmallVector<Location> resultLocs;
    for (auto blockOp : blockOps) {
      blockLocs.push_back(blockOp.getLoc());
      if (blockOp.getNumArguments() > 0) {
        anyRequireDevice = true;
      }
      llvm::append_range(resultTypes, blockOp.getResultTypes());
      llvm::append_range(resultKeys, blockOp.getKeys().getValue());
      llvm::append_range(
          resultLocs,
          gatherResultLocations(blockOp.getNumResults(), blockOp.getRegion()));
    }
    SmallVector<Type> inputTypes;
    if (anyRequireDevice) {
      inputTypes.push_back(IREE::HAL::DeviceType::get(rewriter.getContext()));
    }

    // Create the new combined block op at the location of the first block to
    // keep things in a deterministic order; this makes it look like we are
    // merging all subsequent blocks into the first but without having to worry
    // about making that work.
    rewriter.setInsertionPoint(blockOps.front());
    auto fusedLoc = rewriter.getFusedLoc(blockLocs);
    auto newBlockOp = ExecutableConstantBlockOp::create(
        rewriter, fusedLoc, rewriter.getFunctionType(inputTypes, resultTypes),
        rewriter.getArrayAttr(resultKeys), /*arg_attrs=*/ArrayAttr(),
        /*res_attrs=*/ArrayAttr());

    // Create the entry block that captures the optional device argument and
    // the exit block that returns the final flattened set of keys.
    auto &targetRegion = newBlockOp.getRegion();
    auto *preBlock = newBlockOp.addEntryBlock();
    SmallVector<Block *> targetBlocks;
    for (size_t i = 0; i < blockOps.size(); ++i) {
      targetBlocks.push_back(&targetRegion.emplaceBlock());
    }
    auto *postBlock = &targetRegion.emplaceBlock();
    OpBuilder builder = OpBuilder::atBlockBegin(preBlock);
    cf::BranchOp::create(builder, blockOps.front().getLoc(),
                         targetBlocks.front());

    // Inline all source constant block regions (which may have multiple
    // Blocks).
    SmallVector<Value> resultValues;
    for (unsigned i = 0; i < targetBlocks.size(); ++i) {
      auto *headerBlock = targetBlocks[i];
      auto *nextBlock =
          i < targetBlocks.size() - 1 ? targetBlocks[i + 1] : postBlock;
      auto blockOp = blockOps[i];
      auto &sourceRegion = blockOp.getRegion();

      // Ensure there's only one hal.return in the region.
      // This makes it easier to splice in as we can capture the returned values
      // for use in our combined return.
      rewriteToOneReturn(resultTypes.size(), sourceRegion, rewriter);

      // Inline the entire CFG of the constant block into the target.
      rewriter.cloneRegionBefore(sourceRegion, nextBlock);

      // Branch from the header block into the first block of the region. Note
      // that it may have a %device argument.
      Block *firstBlock = headerBlock->getNextNode();
      SmallVector<Value> firstBranchOperands;
      if (firstBlock->getNumArguments() > 0) {
        firstBranchOperands.push_back(newBlockOp.getArgument(0));
      }
      OpBuilder headerBuilder = OpBuilder::atBlockEnd(headerBlock);
      cf::BranchOp::create(headerBuilder, newBlockOp.getLoc(), firstBlock,
                           firstBranchOperands);

      // Find the single expected return, capture its operands, and rewrite it
      // to branch to the next block.
      for (auto returnOp : llvm::make_early_inc_range(
               targetRegion.getOps<IREE::HAL::ReturnOp>())) {
        OpBuilder builder(returnOp);
        llvm::append_range(resultValues, returnOp.getOperands());
        cf::BranchOp::create(builder, returnOp.getLoc(), nextBlock);
        rewriter.eraseOp(returnOp);
      }
    }

    // Return from the constant block with all operands.
    OpBuilder postBlockbuilder = OpBuilder::atBlockBegin(postBlock);
    IREE::HAL::ReturnOp::create(postBlockbuilder, fusedLoc, resultValues);

    rewriter.finalizeOpModification(variantOp);

    // Erase all the old blocks.
    for (auto blockOp : blockOps) {
      rewriter.eraseOp(blockOp);
    }

    return success();
  }
};

} // namespace

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

namespace {

static void filterReturnOperands(ExecutableConstantBlockOp blockOp,
                                 const BitVector &preservedIndices) {
  for (auto returnOp :
       llvm::make_early_inc_range(blockOp.getOps<IREE::HAL::ReturnOp>())) {
    SmallVector<Value> operands;
    for (auto [i, operand] : llvm::enumerate(returnOp.getOperands())) {
      if (preservedIndices.test(i)) {
        operands.push_back(operand);
      }
    }
    returnOp.getOperandsMutable().assign(operands);
  }
}

/// Drops the %device argument of a constant block region if unused.
struct DropUnusedExecutableConstantBlockDeviceArg
    : OpRewritePattern<ExecutableConstantBlockOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(ExecutableConstantBlockOp blockOp,
                                PatternRewriter &rewriter) const override {
    if (blockOp.getNumArguments() == 0) {
      return failure();
    }
    auto deviceArg = blockOp.getArgument(0);
    if (!deviceArg.use_empty()) {
      return failure();
    }
    rewriter.modifyOpInPlace(blockOp, [&]() {
      // Type conversion here shouldn't fail.
      (void)blockOp.eraseArgument(0);
      blockOp.setFunctionTypeAttr(TypeAttr::get(
          rewriter.getFunctionType(/*inputs=*/{}, blockOp.getResultTypes())));
    });
    return success();
  }
};

/// Deduplicates constant values that have matching keys, choosing the first
/// one found. There's no verification that the values produced are the same
/// as users are expected to uniquely name their keys.
struct DeduplicateExecutableConstantBlockKeys
    : OpRewritePattern<ExecutableConstantBlockOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(ExecutableConstantBlockOp blockOp,
                                PatternRewriter &rewriter) const override {
    // Build a set of preserved result indices (those with unique keys).
    BitVector resultIndices(blockOp.getNumResults(), /*t=*/false);
    SmallVector<Type> resultTypes;
    SetVector<Attribute> resultKeys;
    int i = 0;
    for (auto [resultKey, resultType] : llvm::zip_equal(
             blockOp.getKeys().getValue(), blockOp.getResultTypes())) {
      if (resultKeys.insert(resultKey)) {
        resultIndices.set(i);
        resultTypes.push_back(resultType);
      }
      ++i;
    }

    // If all results are preserved this is a no-op.
    if (resultIndices.all()) {
      return rewriter.notifyMatchFailure(blockOp, "no duplicate keys");
    }

    // Update function in-place.
    rewriter.modifyOpInPlace(blockOp, [&]() {
      // Update metadata.
      blockOp.setFunctionTypeAttr(TypeAttr::get(
          rewriter.getFunctionType(blockOp.getArgumentTypes(), resultTypes)));
      blockOp.setKeysAttr(rewriter.getArrayAttr(resultKeys.takeVector()));
      // Drop all unneeded results from each return.
      filterReturnOperands(blockOp, resultIndices);
    });
    return success();
  }
};

} // namespace

void ExecutableConstantBlockOp::getCanonicalizationPatterns(
    RewritePatternSet &results, MLIRContext *context) {
  results.insert<DropUnusedExecutableConstantBlockDeviceArg>(context);
  results.insert<DeduplicateExecutableConstantBlockKeys>(context);
}

//===----------------------------------------------------------------------===//
// hal.fence.create
//===----------------------------------------------------------------------===//

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

//===----------------------------------------------------------------------===//
// hal.fence.join
//===----------------------------------------------------------------------===//

OpFoldResult FenceJoinOp::fold(FoldAdaptor operands) {
  if (getFences().size() == 1) {
    return getFences().front();
  }
  return {};
}

namespace {

/// Replaces a fence join with no operands with a null value.
struct ElideEmptyFenceJoin : OpRewritePattern<FenceJoinOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(FenceJoinOp op,
                                PatternRewriter &rewriter) const override {
    if (op.getNumOperands() != 0) {
      return failure();
    }
    rewriter.replaceOpWithNewOp<IREE::Util::NullOp>(op,
                                                    op.getResult().getType());
    return success();
  }
};

// Produces a deduplicated and null-elided operand list.
// Returns std::nullopt if nothing changed.
static std::optional<SmallVector<Value, 0>>
deduplicateFenceOperands(ValueRange operands) {
  SetVector<Value> newOperands;
  for (auto operand : operands) {
    if (isa_and_nonnull<IREE::Util::NullOp>(operand.getDefiningOp())) {
      // Drop null values as they don't mean anything. Ideally we'd reach back
      // a little further here but that's best done in an IPO pass.
      continue;
    }
    newOperands.insert(operand);
  }

  if (newOperands.size() == operands.size()) {
    return std::nullopt;
  }
  return newOperands.takeVector();
}

/// Deduplicates fence join operands and drops nulls.
struct DeduplicateFenceJoinFences : OpRewritePattern<FenceJoinOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(FenceJoinOp op,
                                PatternRewriter &rewriter) const override {
    auto newOperands = deduplicateFenceOperands(op.getFences());
    if (!newOperands) {
      return failure();
    }
    rewriter.replaceOpWithNewOp<FenceJoinOp>(
        op, op.getResult().getType(), op.getFlagsAttr(), newOperands.value());
    return success();
  }
};

} // namespace

void FenceJoinOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                              MLIRContext *context) {
  results.insert<ElideEmptyFenceJoin>(context);
  results.insert<DeduplicateFenceJoinFences>(context);
}

//===----------------------------------------------------------------------===//
// hal.fence.signal
//===----------------------------------------------------------------------===//

namespace {

/// Replaces a fence that is immediately signaled on the host with a null fence.
/// This is only safe if there are no users of the fence between where it is
/// created and where it is signaled. We keep things in the local block scope
/// but a larger data flow analysis pass would be useful for propagating across
/// block/function boundaries (common in larger loops/call trees where signal
/// fences are passed as arguments).
///
/// Example:
///  %fence = hal.fence.create
///  hal.fence.signal<%fence : !hal.fence>
/// ->
///  %fence = util.null : !hal.fence
struct ElideSignaledFence : OpRewritePattern<FenceSignalOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(FenceSignalOp signalOp,
                                PatternRewriter &rewriter) const override {
    auto fence = signalOp.getFence();
    auto createOp =
        dyn_cast_if_present<IREE::HAL::FenceCreateOp>(fence.getDefiningOp());
    if (!createOp) {
      return failure();
    }

    // TODO(benvanik): broader analysis - likely in a dedicated fence elision
    // pass so we can do IPO. For now block-only.
    if (createOp->getBlock() != signalOp->getBlock()) {
      return rewriter.notifyMatchFailure(
          signalOp,
          "fence create and signal are in different blocks; analysis TBD");
    }

    // Ensure there are no uses between the create and the signal.
    // There are probably some uses we could allow (selects, etc) but we'll
    // reserve that for a larger analysis.
    for (auto userOp : fence.getUsers()) {
      if (userOp->getBlock() == signalOp->getBlock() &&
          userOp->isBeforeInBlock(signalOp)) {
        return rewriter.notifyMatchFailure(
            signalOp, "interleaved fence usage; cannot elide");
      }
    }

    // Safe to elide.
    Value nullFence = IREE::Util::NullOp::create(
        rewriter, rewriter.getFusedLoc({createOp.getLoc(), signalOp.getLoc()}),
        fence.getType());
    rewriter.replaceAllUsesWith(fence, nullFence);
    rewriter.eraseOp(signalOp);

    return success();
  }
};

} // namespace

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

//===----------------------------------------------------------------------===//
// hal.fence.await
//===----------------------------------------------------------------------===//

namespace {

/// Elides a fence await with no fences.
struct ElideEmptyFenceAwait : OpRewritePattern<FenceAwaitOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(FenceAwaitOp op,
                                PatternRewriter &rewriter) const override {
    if (!op.getFences().empty()) {
      return failure();
    }
    rewriter.replaceOpWithNewOp<arith::ConstantIntOp>(op, /*ok=*/0, 32);
    return success();
  }
};

/// Deduplicates fence await operands and drops nulls.
struct DeduplicateFenceAwaitFences : OpRewritePattern<FenceAwaitOp> {
  using Base::Base;
  LogicalResult matchAndRewrite(FenceAwaitOp op,
                                PatternRewriter &rewriter) const override {
    auto newOperands = deduplicateFenceOperands(op.getFences());
    if (newOperands == std::nullopt) {
      return failure();
    }
    // TODO(benvanik): resolve flag sets.
    rewriter.replaceOpWithNewOp<FenceAwaitOp>(
        op, op.getStatus().getType(), op.getTimeoutMillis(), op.getFlagsAttr(),
        newOperands.value());
    return success();
  }
};

} // namespace

void FenceAwaitOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                               MLIRContext *context) {
  results.insert<ElideEmptyFenceAwait>(context);
  results.insert<DeduplicateFenceAwaitFences>(context);
}

//===----------------------------------------------------------------------===//
// hal.buffer_usage
//===----------------------------------------------------------------------===//

OpFoldResult BufferUsageOp::fold(FoldAdaptor operands) {
  return IntegerAttr::get(IntegerType::get(getContext(), 32),
                          getUsageAttr().getInt());
}

//===----------------------------------------------------------------------===//
// hal.memory_type
//===----------------------------------------------------------------------===//

OpFoldResult MemoryTypeOp::fold(FoldAdaptor operands) {
  return IntegerAttr::get(IntegerType::get(getContext(), 32),
                          getTypeAttr().getInt());
}

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