MemorySlot.cpp

Go to the documentation of this file.

//===- MemorySlot.cpp - Memory Slot interface implementations for SCF -----===//
//
// Part of the LLVM Project, 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 "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Interfaces/Utils/MemorySlotUtils.h"
 
using namespace mlir;
using namespace mlir::scf;
 
//===----------------------------------------------------------------------===//
// ExecuteRegionOp
//===----------------------------------------------------------------------===//
 
bool ExecuteRegionOp::isRegionPromotable(const MemorySlot &slot, Region *region,
                                         bool hasValueStores) {
  return true;
}
 
void ExecuteRegionOp::setupPromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    llvm::SmallMapVector<Region *, Value, 2> &regionsToProcess) {
  regionsToProcess.insert({&getRegion(), reachingDef});
}
 
Value ExecuteRegionOp::finalizePromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    const llvm::DenseMap<Block *, Value> &reachingAtBlockEnd,
    OpBuilder &builder) {
  if (!hasValueStores)
    return reachingDef;
 
  // Update the yield terminators to return the newly defined reaching
  // definition.
  for (Block &block : getRegion().getBlocks())
    if (isa<YieldOp>(block.getTerminator()))
      memoryslot::updateTerminator(&block, reachingDef, reachingAtBlockEnd);
 
  SmallVector<Type> resultTypes(getResultTypes());
  resultTypes.push_back(slot.elemType);
 
  IRRewriter rewriter(builder);
  Operation *newOp =
      memoryslot::replaceWithNewResults(rewriter, getOperation(), resultTypes);
  return newOp->getResults().back();
}
 
//===----------------------------------------------------------------------===//
// ForOp
//===----------------------------------------------------------------------===//
 
bool ForOp::isRegionPromotable(const MemorySlot &slot, Region *region,
                               bool hasValueStores) {
  return true;
}
 
void ForOp::setupPromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    llvm::SmallMapVector<Region *, Value, 2> &regionsToProcess) {
  Region &bodyRegion = getBodyRegion();
  if (!hasValueStores) {
    regionsToProcess.insert({&bodyRegion, reachingDef});
    return;
  }
 
  getInitArgsMutable().append(reachingDef);
  bodyRegion.addArgument(slot.elemType, slot.ptr.getLoc());
  regionsToProcess.insert({&bodyRegion, bodyRegion.getArguments().back()});
}
 
Value ForOp::finalizePromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    const llvm::DenseMap<Block *, Value> &reachingAtBlockEnd,
    OpBuilder &builder) {
  if (!hasValueStores)
    return reachingDef;
 
  // Update the yield terminator to return the newly defined reaching
  // definition.
  memoryslot::updateTerminator(getBody(), reachingDef, reachingAtBlockEnd);
 
  SmallVector<Type> resultTypes(getResultTypes());
  resultTypes.push_back(slot.elemType);
 
  IRRewriter rewriter(builder);
  Operation *newOp =
      memoryslot::replaceWithNewResults(rewriter, getOperation(), resultTypes);
  return newOp->getResults().back();
}
 
//===----------------------------------------------------------------------===//
// ForallOp
//===----------------------------------------------------------------------===//
 
bool ForallOp::isRegionPromotable(const MemorySlot &slot, Region *region,
                                  bool hasValueStores) {
  // The ForallOp body can be ran in parallel, thus does not support sequenced
  // value passing. Therefore only loads can be handled.
  return !hasValueStores;
}
 
void ForallOp::setupPromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    llvm::SmallMapVector<Region *, Value, 2> &regionsToProcess) {
  assert(!hasValueStores && "ForallOp does not support stores");
  regionsToProcess.insert({&getBodyRegion(), reachingDef});
}
 
Value ForallOp::finalizePromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    const llvm::DenseMap<Block *, Value> &reachingAtBlockEnd,
    OpBuilder &builder) {
  assert(!hasValueStores && "ForallOp does not support stores");
  return reachingDef;
}
 
//===----------------------------------------------------------------------===//
// IfOp
//===----------------------------------------------------------------------===//
 
bool IfOp::isRegionPromotable(const MemorySlot &slot, Region *region,
                              bool hasValueStores) {
  return true;
}
 
void IfOp::setupPromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    llvm::SmallMapVector<Region *, Value, 2> &regionsToProcess) {
  regionsToProcess.insert({&getThenRegion(), reachingDef});
  regionsToProcess.insert({&getElseRegion(), reachingDef});
}
 
Value IfOp::finalizePromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    const llvm::DenseMap<Block *, Value> &reachingAtBlockEnd,
    OpBuilder &builder) {
  if (!hasValueStores)
    return reachingDef;
 
  IRRewriter rewriter(builder);
 
  // Update the yield terminators to return the newly defined reaching
  // definition.
  memoryslot::updateTerminator(&getThenRegion().back(), reachingDef,
                               reachingAtBlockEnd);
  if (getElseRegion().hasOneBlock()) {
    memoryslot::updateTerminator(&getElseRegion().back(), reachingDef,
                                 reachingAtBlockEnd);
  } else {
    OpBuilder::InsertionGuard guard(rewriter);
    rewriter.createBlock(&getElseRegion());
    YieldOp::create(rewriter, getOperation()->getLoc(), reachingDef);
  }
 
  SmallVector<Type> resultTypes(getResultTypes());
  resultTypes.push_back(slot.elemType);
 
  Operation *newOp =
      memoryslot::replaceWithNewResults(rewriter, getOperation(), resultTypes);
  return newOp->getResults().back();
}
 
//===----------------------------------------------------------------------===//
// IndexSwitchOp
//===----------------------------------------------------------------------===//
 
bool IndexSwitchOp::isRegionPromotable(const MemorySlot &slot, Region *region,
                                       bool hasValueStores) {
  return true;
}
 
void IndexSwitchOp::setupPromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    llvm::SmallMapVector<Region *, Value, 2> &regionsToProcess) {
  regionsToProcess.insert({&getDefaultRegion(), reachingDef});
  for (Region &caseRegion : getCaseRegions())
    regionsToProcess.insert({&caseRegion, reachingDef});
}
 
Value IndexSwitchOp::finalizePromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    const llvm::DenseMap<Block *, Value> &reachingAtBlockEnd,
    OpBuilder &builder) {
  if (!hasValueStores)
    return reachingDef;
 
  IRRewriter rewriter(builder);
 
  // Update the yield terminators to return the newly defined reaching
  // definition.
  memoryslot::updateTerminator(&getDefaultRegion().back(), reachingDef,
                               reachingAtBlockEnd);
  for (Region &caseRegion : getCaseRegions())
    memoryslot::updateTerminator(&caseRegion.back(), reachingDef,
                                 reachingAtBlockEnd);
 
  SmallVector<Type> resultTypes(getResultTypes());
  resultTypes.push_back(slot.elemType);
 
  Operation *newOp =
      memoryslot::replaceWithNewResults(rewriter, getOperation(), resultTypes);
  return newOp->getResults().back();
}
 
//===----------------------------------------------------------------------===//
// ParallelOp
//===----------------------------------------------------------------------===//
 
bool ParallelOp::isRegionPromotable(const MemorySlot &slot, Region *region,
                                    bool hasValueStores) {
  // The ParallelOp body can be ran in parallel, thus does not support sequenced
  // value passing. Therefore only loads can be handled.
  return !hasValueStores;
}
 
void ParallelOp::setupPromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    llvm::SmallMapVector<Region *, Value, 2> &regionsToProcess) {
  assert(!hasValueStores && "ParallelOp does not support stores");
  regionsToProcess.insert({&getBodyRegion(), reachingDef});
}
 
Value ParallelOp::finalizePromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    const llvm::DenseMap<Block *, Value> &reachingAtBlockEnd,
    OpBuilder &builder) {
  assert(!hasValueStores && "ParallelOp does not support stores");
  return reachingDef;
}
 
//===----------------------------------------------------------------------===//
// ReduceOp
//===----------------------------------------------------------------------===//
 
bool ReduceOp::isRegionPromotable(const MemorySlot &slot, Region *region,
                                  bool hasValueStores) {
  // The ReduceOp body can be ran in parallel, thus does not support sequenced
  // value passing. Therefore only loads can be handled.
  return !hasValueStores;
}
 
void ReduceOp::setupPromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    llvm::SmallMapVector<Region *, Value, 2> &regionsToProcess) {
  assert(!hasValueStores && "ReduceOp does not support stores");
  for (Region &reduction : getReductions())
    regionsToProcess.insert({&reduction, reachingDef});
}
 
Value ReduceOp::finalizePromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    const llvm::DenseMap<Block *, Value> &reachingAtBlockEnd,
    OpBuilder &builder) {
  assert(!hasValueStores && "ReduceOp does not support stores");
  return reachingDef;
}
 
//===----------------------------------------------------------------------===//
// WhileOp
//===----------------------------------------------------------------------===//
 
bool WhileOp::isRegionPromotable(const MemorySlot &slot, Region *region,
                                 bool hasValueStores) {
  return true;
}
 
void WhileOp::setupPromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    llvm::SmallMapVector<Region *, Value, 2> &regionsToProcess) {
  Region &beforeRegion = getBefore();
  Region &afterRegion = getAfter();
  if (!hasValueStores) {
    regionsToProcess.insert({&beforeRegion, reachingDef});
    regionsToProcess.insert({&afterRegion, reachingDef});
    return;
  }
 
  getInitsMutable().append(reachingDef);
 
  beforeRegion.addArgument(slot.elemType, slot.ptr.getLoc());
  regionsToProcess.insert({&beforeRegion, beforeRegion.getArguments().back()});
 
  afterRegion.addArgument(slot.elemType, slot.ptr.getLoc());
  regionsToProcess.insert({&afterRegion, afterRegion.getArguments().back()});
}
 
Value WhileOp::finalizePromotion(
    const MemorySlot &slot, Value reachingDef, bool hasValueStores,
    const llvm::DenseMap<Block *, Value> &reachingAtBlockEnd,
    OpBuilder &builder) {
  if (!hasValueStores)
    return reachingDef;
 
  // Update the yield terminators to return the newly defined reaching
  // definition.
  memoryslot::updateTerminator(&getBefore().back(),
                               getBefore().getArguments().back(),
                               reachingAtBlockEnd);
  memoryslot::updateTerminator(
      &getAfter().back(), getAfter().getArguments().back(), reachingAtBlockEnd);
 
  SmallVector<Type> resultTypes(getResultTypes());
  resultTypes.push_back(slot.elemType);
 
  IRRewriter rewriter(builder);
  Operation *newOp =
      memoryslot::replaceWithNewResults(rewriter, getOperation(), resultTypes);
  return newOp->getResults().back();
}