//===- SCFToSPIRV.cpp - SCF to SPIR-V Patterns ----------------------------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// This file implements patterns to convert SCF dialect to SPIR-V dialect.
//
//===----------------------------------------------------------------------===//

#include "mlir/Conversion/SCFToSPIRV/SCFToSPIRV.h"
#include "mlir/Dialect/SCF/SCF.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
#include "mlir/Dialect/SPIRV/Transforms/SPIRVConversion.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/Transforms/DialectConversion.h"

using namespace mlir;

//===----------------------------------------------------------------------===//
// Context
//===----------------------------------------------------------------------===//

namespace mlir {
struct ScfToSPIRVContextImpl {
  // Map between the spirv region control flow operation (spv.mlir.loop or
  // spv.mlir.selection) to the VariableOp created to store the region results.
  // The order of the VariableOp matches the order of the results.
  DenseMap<Operation *, SmallVector<spirv::VariableOp, 8>> outputVars;
};
} // namespace mlir

/// We use ScfToSPIRVContext to store information about the lowering of the scf
/// region that need to be used later on. When we lower scf.for/scf.if we create
/// VariableOp to store the results. We need to keep track of the VariableOp
/// created as we need to insert stores into them when lowering Yield. Those
/// StoreOp cannot be created earlier as they may use a different type than
/// yield operands.
ScfToSPIRVContext::ScfToSPIRVContext() {
  impl = std::make_unique<ScfToSPIRVContextImpl>();
}

ScfToSPIRVContext::~ScfToSPIRVContext() = default;

//===----------------------------------------------------------------------===//
// Pattern Declarations
//===----------------------------------------------------------------------===//

namespace {
/// Common class for all vector to GPU patterns.
template <typename OpTy>
class SCFToSPIRVPattern : public OpConversionPattern<OpTy> {
public:
  SCFToSPIRVPattern<OpTy>(MLIRContext *context, SPIRVTypeConverter &converter,
                          ScfToSPIRVContextImpl *scfToSPIRVContext)
      : OpConversionPattern<OpTy>::OpConversionPattern(converter, context),
        scfToSPIRVContext(scfToSPIRVContext), typeConverter(converter) {}

protected:
  ScfToSPIRVContextImpl *scfToSPIRVContext;
  // FIXME: We explicitly keep a reference of the type converter here instead of
  // passing it to OpConversionPattern during construction. This effectively
  // bypasses the conversion framework's automation on type conversion. This is
  // needed right now because the conversion framework will unconditionally
  // legalize all types used by SCF ops upon discovering them, for example, the
  // types of loop carried values. We use SPIR-V variables for those loop
  // carried values. Depending on the available capabilities, the SPIR-V
  // variable can be different, for example, cooperative matrix or normal
  // variable. We'd like to detach the conversion of the loop carried values
  // from the SCF ops (which is mainly a region). So we need to "mark" types
  // used by SCF ops as legal, if to use the conversion framework for type
  // conversion. There isn't a straightforward way to do that yet, as when
  // converting types, ops aren't taken into consideration. Therefore, we just
  // bypass the framework's type conversion for now.
  SPIRVTypeConverter &typeConverter;
};

/// Pattern to convert a scf::ForOp within kernel functions into spirv::LoopOp.
class ForOpConversion final : public SCFToSPIRVPattern<scf::ForOp> {
public:
  using SCFToSPIRVPattern<scf::ForOp>::SCFToSPIRVPattern;

  LogicalResult
  matchAndRewrite(scf::ForOp forOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};

/// Pattern to convert a scf::IfOp within kernel functions into
/// spirv::SelectionOp.
class IfOpConversion final : public SCFToSPIRVPattern<scf::IfOp> {
public:
  using SCFToSPIRVPattern<scf::IfOp>::SCFToSPIRVPattern;

  LogicalResult
  matchAndRewrite(scf::IfOp ifOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};

class TerminatorOpConversion final : public SCFToSPIRVPattern<scf::YieldOp> {
public:
  using SCFToSPIRVPattern<scf::YieldOp>::SCFToSPIRVPattern;

  LogicalResult
  matchAndRewrite(scf::YieldOp terminatorOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};

class WhileOpConversion final : public SCFToSPIRVPattern<scf::WhileOp> {
public:
  using SCFToSPIRVPattern<scf::WhileOp>::SCFToSPIRVPattern;

  LogicalResult
  matchAndRewrite(scf::WhileOp forOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override;
};
} // namespace

/// Helper function to replaces SCF op outputs with SPIR-V variable loads.
/// We create VariableOp to handle the results value of the control flow region.
/// spv.mlir.loop/spv.mlir.selection currently don't yield value. Right after
/// the loop we load the value from the allocation and use it as the SCF op
/// result.
template <typename ScfOp, typename OpTy>
static void replaceSCFOutputValue(ScfOp scfOp, OpTy newOp,
                                  ConversionPatternRewriter &rewriter,
                                  ScfToSPIRVContextImpl *scfToSPIRVContext,
                                  ArrayRef<Type> returnTypes) {

  Location loc = scfOp.getLoc();
  auto &allocas = scfToSPIRVContext->outputVars[newOp];
  // Clearing the allocas is necessary in case a dialect conversion path failed
  // previously, and this is the second attempt of this conversion.
  allocas.clear();
  SmallVector<Value, 8> resultValue;
  for (Type convertedType : returnTypes) {
    auto pointerType =
        spirv::PointerType::get(convertedType, spirv::StorageClass::Function);
    rewriter.setInsertionPoint(newOp);
    auto alloc = rewriter.create<spirv::VariableOp>(
        loc, pointerType, spirv::StorageClass::Function,
        /*initializer=*/nullptr);
    allocas.push_back(alloc);
    rewriter.setInsertionPointAfter(newOp);
    Value loadResult = rewriter.create<spirv::LoadOp>(loc, alloc);
    resultValue.push_back(loadResult);
  }
  rewriter.replaceOp(scfOp, resultValue);
}

static Region::iterator getBlockIt(Region &region, unsigned index) {
  return std::next(region.begin(), index);
}

//===----------------------------------------------------------------------===//
// scf::ForOp
//===----------------------------------------------------------------------===//

LogicalResult
ForOpConversion::matchAndRewrite(scf::ForOp forOp, OpAdaptor adaptor,
                                 ConversionPatternRewriter &rewriter) const {
  // scf::ForOp can be lowered to the structured control flow represented by
  // spirv::LoopOp by making the continue block of the spirv::LoopOp the loop
  // latch and the merge block the exit block. The resulting spirv::LoopOp has a
  // single back edge from the continue to header block, and a single exit from
  // header to merge.
  auto loc = forOp.getLoc();
  auto loopOp = rewriter.create<spirv::LoopOp>(loc, spirv::LoopControl::None);
  loopOp.addEntryAndMergeBlock();

  OpBuilder::InsertionGuard guard(rewriter);
  // Create the block for the header.
  auto *header = new Block();
  // Insert the header.
  loopOp.body().getBlocks().insert(getBlockIt(loopOp.body(), 1), header);

  // Create the new induction variable to use.
  BlockArgument newIndVar = header->addArgument(adaptor.lowerBound().getType());
  for (Value arg : adaptor.initArgs())
    header->addArgument(arg.getType());
  Block *body = forOp.getBody();

  // Apply signature conversion to the body of the forOp. It has a single block,
  // with argument which is the induction variable. That has to be replaced with
  // the new induction variable.
  TypeConverter::SignatureConversion signatureConverter(
      body->getNumArguments());
  signatureConverter.remapInput(0, newIndVar);
  for (unsigned i = 1, e = body->getNumArguments(); i < e; i++)
    signatureConverter.remapInput(i, header->getArgument(i));
  body = rewriter.applySignatureConversion(&forOp.getLoopBody(),
                                           signatureConverter);

  // Move the blocks from the forOp into the loopOp. This is the body of the
  // loopOp.
  rewriter.inlineRegionBefore(forOp->getRegion(0), loopOp.body(),
                              getBlockIt(loopOp.body(), 2));

  SmallVector<Value, 8> args(1, adaptor.lowerBound());
  args.append(adaptor.initArgs().begin(), adaptor.initArgs().end());
  // Branch into it from the entry.
  rewriter.setInsertionPointToEnd(&(loopOp.body().front()));
  rewriter.create<spirv::BranchOp>(loc, header, args);

  // Generate the rest of the loop header.
  rewriter.setInsertionPointToEnd(header);
  auto *mergeBlock = loopOp.getMergeBlock();
  auto cmpOp = rewriter.create<spirv::SLessThanOp>(
      loc, rewriter.getI1Type(), newIndVar, adaptor.upperBound());

  rewriter.create<spirv::BranchConditionalOp>(
      loc, cmpOp, body, ArrayRef<Value>(), mergeBlock, ArrayRef<Value>());

  // Generate instructions to increment the step of the induction variable and
  // branch to the header.
  Block *continueBlock = loopOp.getContinueBlock();
  rewriter.setInsertionPointToEnd(continueBlock);

  // Add the step to the induction variable and branch to the header.
  Value updatedIndVar = rewriter.create<spirv::IAddOp>(
      loc, newIndVar.getType(), newIndVar, adaptor.step());
  rewriter.create<spirv::BranchOp>(loc, header, updatedIndVar);

  // Infer the return types from the init operands. Vector type may get
  // converted to CooperativeMatrix or to Vector type, to avoid having complex
  // extra logic to figure out the right type we just infer it from the Init
  // operands.
  SmallVector<Type, 8> initTypes;
  for (auto arg : adaptor.initArgs())
    initTypes.push_back(arg.getType());
  replaceSCFOutputValue(forOp, loopOp, rewriter, scfToSPIRVContext, initTypes);
  return success();
}

//===----------------------------------------------------------------------===//
// scf::IfOp
//===----------------------------------------------------------------------===//

LogicalResult
IfOpConversion::matchAndRewrite(scf::IfOp ifOp, OpAdaptor adaptor,
                                ConversionPatternRewriter &rewriter) const {
  // When lowering `scf::IfOp` we explicitly create a selection header block
  // before the control flow diverges and a merge block where control flow
  // subsequently converges.
  auto loc = ifOp.getLoc();

  // Create `spv.selection` operation, selection header block and merge block.
  auto selectionOp =
      rewriter.create<spirv::SelectionOp>(loc, spirv::SelectionControl::None);
  auto *mergeBlock =
      rewriter.createBlock(&selectionOp.body(), selectionOp.body().end());
  rewriter.create<spirv::MergeOp>(loc);

  OpBuilder::InsertionGuard guard(rewriter);
  auto *selectionHeaderBlock =
      rewriter.createBlock(&selectionOp.body().front());

  // Inline `then` region before the merge block and branch to it.
  auto &thenRegion = ifOp.thenRegion();
  auto *thenBlock = &thenRegion.front();
  rewriter.setInsertionPointToEnd(&thenRegion.back());
  rewriter.create<spirv::BranchOp>(loc, mergeBlock);
  rewriter.inlineRegionBefore(thenRegion, mergeBlock);

  auto *elseBlock = mergeBlock;
  // If `else` region is not empty, inline that region before the merge block
  // and branch to it.
  if (!ifOp.elseRegion().empty()) {
    auto &elseRegion = ifOp.elseRegion();
    elseBlock = &elseRegion.front();
    rewriter.setInsertionPointToEnd(&elseRegion.back());
    rewriter.create<spirv::BranchOp>(loc, mergeBlock);
    rewriter.inlineRegionBefore(elseRegion, mergeBlock);
  }

  // Create a `spv.BranchConditional` operation for selection header block.
  rewriter.setInsertionPointToEnd(selectionHeaderBlock);
  rewriter.create<spirv::BranchConditionalOp>(loc, adaptor.condition(),
                                              thenBlock, ArrayRef<Value>(),
                                              elseBlock, ArrayRef<Value>());

  SmallVector<Type, 8> returnTypes;
  for (auto result : ifOp.results()) {
    auto convertedType = typeConverter.convertType(result.getType());
    returnTypes.push_back(convertedType);
  }
  replaceSCFOutputValue(ifOp, selectionOp, rewriter, scfToSPIRVContext,
                        returnTypes);
  return success();
}

//===----------------------------------------------------------------------===//
// scf::YieldOp
//===----------------------------------------------------------------------===//

/// Yield is lowered to stores to the VariableOp created during lowering of the
/// parent region. For loops we also need to update the branch looping back to
/// the header with the loop carried values.
LogicalResult TerminatorOpConversion::matchAndRewrite(
    scf::YieldOp terminatorOp, OpAdaptor adaptor,
    ConversionPatternRewriter &rewriter) const {
  ValueRange operands = adaptor.getOperands();

  // If the region is return values, store each value into the associated
  // VariableOp created during lowering of the parent region.
  if (!operands.empty()) {
    auto &allocas = scfToSPIRVContext->outputVars[terminatorOp->getParentOp()];
    if (allocas.size() != operands.size())
      return failure();

    auto loc = terminatorOp.getLoc();
    for (unsigned i = 0, e = operands.size(); i < e; i++)
      rewriter.create<spirv::StoreOp>(loc, allocas[i], operands[i]);
    if (isa<spirv::LoopOp>(terminatorOp->getParentOp())) {
      // For loops we also need to update the branch jumping back to the header.
      auto br =
          cast<spirv::BranchOp>(rewriter.getInsertionBlock()->getTerminator());
      SmallVector<Value, 8> args(br.getBlockArguments());
      args.append(operands.begin(), operands.end());
      rewriter.setInsertionPoint(br);
      rewriter.create<spirv::BranchOp>(terminatorOp.getLoc(), br.getTarget(),
                                       args);
      rewriter.eraseOp(br);
    }
  }
  rewriter.eraseOp(terminatorOp);
  return success();
}

//===----------------------------------------------------------------------===//
// scf::WhileOp
//===----------------------------------------------------------------------===//

LogicalResult
WhileOpConversion::matchAndRewrite(scf::WhileOp whileOp, OpAdaptor adaptor,
                                   ConversionPatternRewriter &rewriter) const {
  auto loc = whileOp.getLoc();
  auto loopOp = rewriter.create<spirv::LoopOp>(loc, spirv::LoopControl::None);
  loopOp.addEntryAndMergeBlock();

  OpBuilder::InsertionGuard guard(rewriter);

  Region &beforeRegion = whileOp.before();
  Region &afterRegion = whileOp.after();

  Block &entryBlock = *loopOp.getEntryBlock();
  Block &beforeBlock = beforeRegion.front();
  Block &afterBlock = afterRegion.front();
  Block &mergeBlock = *loopOp.getMergeBlock();

  auto cond = cast<scf::ConditionOp>(beforeBlock.getTerminator());
  SmallVector<Value> condArgs;
  if (failed(rewriter.getRemappedValues(cond.args(), condArgs)))
    return failure();

  Value conditionVal = rewriter.getRemappedValue(cond.condition());
  if (!conditionVal)
    return failure();

  auto yield = cast<scf::YieldOp>(afterBlock.getTerminator());
  SmallVector<Value> yieldArgs;
  if (failed(rewriter.getRemappedValues(yield.results(), yieldArgs)))
    return failure();

  // Move the while before block as the initial loop header block.
  rewriter.inlineRegionBefore(beforeRegion, loopOp.body(),
                              getBlockIt(loopOp.body(), 1));

  // Move the while after block as the initial loop body block.
  rewriter.inlineRegionBefore(afterRegion, loopOp.body(),
                              getBlockIt(loopOp.body(), 2));

  // Jump from the loop entry block to the loop header block.
  rewriter.setInsertionPointToEnd(&entryBlock);
  rewriter.create<spirv::BranchOp>(loc, &beforeBlock, adaptor.inits());

  auto condLoc = cond.getLoc();

  SmallVector<Value> resultValues(condArgs.size());

  // For other SCF ops, the scf.yield op yields the value for the whole SCF op.
  // So we use the scf.yield op as the anchor to create/load/store SPIR-V local
  // variables. But for the scf.while op, the scf.yield op yields a value for
  // the before region, which may not matching the whole op's result. Instead,
  // the scf.condition op returns values matching the whole op's results. So we
  // need to create/load/store variables according to that.
  for (auto it : llvm::enumerate(condArgs)) {
    auto res = it.value();
    auto i = it.index();
    auto pointerType =
        spirv::PointerType::get(res.getType(), spirv::StorageClass::Function);

    // Create local variables before the scf.while op.
    rewriter.setInsertionPoint(loopOp);
    auto alloc = rewriter.create<spirv::VariableOp>(
        condLoc, pointerType, spirv::StorageClass::Function,
        /*initializer=*/nullptr);

    // Load the final result values after the scf.while op.
    rewriter.setInsertionPointAfter(loopOp);
    auto loadResult = rewriter.create<spirv::LoadOp>(condLoc, alloc);
    resultValues[i] = loadResult;

    // Store the current iteration's result value.
    rewriter.setInsertionPointToEnd(&beforeBlock);
    rewriter.create<spirv::StoreOp>(condLoc, alloc, res);
  }

  rewriter.setInsertionPointToEnd(&beforeBlock);
  rewriter.replaceOpWithNewOp<spirv::BranchConditionalOp>(
      cond, conditionVal, &afterBlock, condArgs, &mergeBlock, llvm::None);

  // Convert the scf.yield op to a branch back to the header block.
  rewriter.setInsertionPointToEnd(&afterBlock);
  rewriter.replaceOpWithNewOp<spirv::BranchOp>(yield, &beforeBlock, yieldArgs);

  rewriter.replaceOp(whileOp, resultValues);
  return success();
}

//===----------------------------------------------------------------------===//
// Hooks
//===----------------------------------------------------------------------===//

void mlir::populateSCFToSPIRVPatterns(SPIRVTypeConverter &typeConverter,
                                      ScfToSPIRVContext &scfToSPIRVContext,
                                      RewritePatternSet &patterns) {
  patterns.add<ForOpConversion, IfOpConversion, TerminatorOpConversion,
               WhileOpConversion>(patterns.getContext(), typeConverter,
                                  scfToSPIRVContext.getImpl());
}