// Copyright 2020 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//===- ConvertToGPUPass.cpp -----------------------------------------------===//
//
// Partition computation within dispatch function to workgroups/workitems.
//
//===----------------------------------------------------------------------===//

#include <array>
#include <numeric>

#include "iree/compiler/Conversion/CodegenUtils/FunctionUtils.h"
#include "iree/compiler/Conversion/CodegenUtils/MarkerUtils.h"
#include "iree/compiler/Conversion/Common/Attributes.h"
#include "iree/compiler/Conversion/Common/Transforms.h"
#include "iree/compiler/Conversion/LinalgToSPIRV/KernelDispatchUtils.h"
#include "iree/compiler/Conversion/LinalgToSPIRV/MemorySpace.h"
#include "iree/compiler/Conversion/LinalgToSPIRV/Passes.h"
#include "iree/compiler/Conversion/LinalgToSPIRV/Utils.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
#include "iree/compiler/Dialect/Shape/IR/ShapeDialect.h"
#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/GPU/GPUDialect.h"
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/SCF/SCF.h"
#include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
#include "mlir/Dialect/SPIRV/IR/TargetAndABI.h"
#include "mlir/Dialect/StandardOps/IR/Ops.h"
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/FunctionSupport.h"
#include "mlir/Support/LLVM.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/LoopUtils.h"

namespace mlir {
namespace iree_compiler {

//===----------------------------------------------------------------------===//
// Loop utilities
//===----------------------------------------------------------------------===//

/// Builds an empty scf.for operation. The default builder adds an entry basic
/// block which needs to be avoided here.
static scf::ForOp buildEmptyForOp(Location loc, OpBuilder &builder, Value lb,
                                  Value ub, Value step) {
  OperationState state(loc, scf::ForOp::getOperationName());
  state.addOperands({lb, ub, step});
  state.addRegion();
  return cast<scf::ForOp>(builder.createOperation(state));
}

/// Builds an empty scf.if operation without the then and else blocks.
static scf::IfOp buildEmptyIfOp(Location loc, OpBuilder &builder, Value cond) {
  OperationState state(loc, scf::IfOp::getOperationName());
  state.addOperands(cond);
  state.addRegion();
  state.addRegion();
  return cast<scf::IfOp>(builder.createOperation(state));
}

namespace {
struct LoopBounds {
  Value lb;
  Value ub;
  Value step;
};
}  // namespace

/// Replaces a scf.parallelOp with an optional scf.parallel op and nested
/// scf.for operations. To create the scf.parallel op as the outermost loop,
/// pass the lower bound, upper bound and steps in `newPLoopLbs`, `newPLoopUbs`,
/// and `newPLoopStep` respectively. The bounds of the inner scf.for operations
/// to be created are passed in `forLbs`, `forUbs`, and `forStep`. The
/// `permutation` vector contains a mapping from the original loop order, to the
/// loop order to be generated.
static Operation *replacePLoopOp(ConversionPatternRewriter &rewriter,
                                 scf::ParallelOp pLoopOp,
                                 ArrayRef<LoopBounds> newPLoopBounds,
                                 ArrayRef<LoopBounds> forBounds,
                                 ArrayRef<unsigned> permutation) {
  assert(!forBounds.empty() && "unhandled case of no scf.for created");
  unsigned numLoops = pLoopOp.getNumLoops();
  Location loc = pLoopOp.getLoc();
  assert(forBounds.size() + newPLoopBounds.size() == numLoops &&
         "cannot drop loops when splitting scf.parallel operation");
  assert(permutation.size() == numLoops);
  OpBuilder::InsertionGuard guard(rewriter);

  // Need a signature conversion for the body of the scf.parallel operation,
  // before can it can be used as the body of the innermost loop created here.
  TypeConverter::SignatureConversion signatureConverter(numLoops);
  Operation *outermostLoop = nullptr;
  auto permuteIt = permutation.begin();

  // Create the scf.parallel operation as the outermost loop, if specified.
  if (!newPLoopBounds.empty()) {
    auto lbs = llvm::to_vector<2>(llvm::map_range(
        newPLoopBounds, [](LoopBounds bounds) -> Value { return bounds.lb; }));
    auto ubs = llvm::to_vector<2>(llvm::map_range(
        newPLoopBounds, [](LoopBounds bounds) { return bounds.ub; }));
    auto steps = llvm::to_vector<2>(llvm::map_range(
        newPLoopBounds, [](LoopBounds bounds) { return bounds.step; }));
    auto newPLoop = rewriter.create<scf::ParallelOp>(loc, lbs, ubs, steps);
    for (auto iv : newPLoop.getInductionVars()) {
      signatureConverter.remapInput(*permuteIt, iv);
      permuteIt++;
    }
    rewriter.setInsertionPointToStart(newPLoop.getBody());
    outermostLoop = newPLoop.getOperation();
  }

  // Generate the nested scf.for operations with the bounds passed.
  for (auto it : enumerate(forBounds)) {
    Value lb = it.value().lb, ub = it.value().ub, step = it.value().step;
    if (it.index() != forBounds.size() - 1) {
      auto forOp = rewriter.create<scf::ForOp>(loc, lb, ub, step);
      if (!outermostLoop) outermostLoop = forOp.getOperation();
      signatureConverter.remapInput(*permuteIt, forOp.getInductionVar());
      rewriter.setInsertionPointToStart(forOp.getBody());
    } else {
      // For the last loop, move the body of the scf.parallel op as the body of
      // the loop after signature conversion.
      auto forOp = buildEmptyForOp(loc, rewriter, lb, ub, step);
      if (!outermostLoop) outermostLoop = forOp.getOperation();
      signatureConverter.addInputs(*permuteIt, rewriter.getIndexType());
      Region &pLoopOpRegion = pLoopOp.getLoopBody();
      rewriter.applySignatureConversion(&pLoopOpRegion, signatureConverter);
      Region &forOpRegion = forOp.getLoopBody();
      rewriter.inlineRegionBefore(pLoopOpRegion, forOpRegion,
                                  forOpRegion.begin());
    }
    permuteIt++;
  }
  rewriter.eraseOp(pLoopOp);
  return outermostLoop;
}

/// Serializes the dimensions of the scf.parallel specified in
/// `serializedDimensions`, by creating an nested scf.for operation for each
/// dimension.
// TODO(ravishankarm): Move this into LoopUtils.h in MLIR.
static Operation *serializeDimensions(ConversionPatternRewriter &rewriter,
                                      scf::ParallelOp pLoopOp,
                                      ArrayRef<unsigned> serializedDimensions) {
  assert(!serializedDimensions.empty() &&
         "unhandled corner case of no serializing dims");
  OpBuilder::InsertionGuard guard(rewriter);
  DenseSet<unsigned> serializedDimSet;
  serializedDimSet.insert(serializedDimensions.begin(),
                          serializedDimensions.end());
  assert(serializedDimSet.size() == serializedDimensions.size() &&
         "cannot repeat dimensions during serialization of scf.parallel");
  SmallVector<LoopBounds, 2> newPLoopBounds, forBounds;
  SmallVector<unsigned, 2> permutation;
  auto lbs = pLoopOp.lowerBound();
  auto ubs = pLoopOp.upperBound();
  auto steps = pLoopOp.step();
  for (unsigned i : llvm::seq<unsigned>(0, pLoopOp.getNumLoops())) {
    if (serializedDimSet.count(i)) {
      forBounds.push_back({lbs[i], ubs[i], steps[i]});
    } else {
      newPLoopBounds.push_back({lbs[i], ubs[i], steps[i]});
      permutation.push_back(i);
    }
  }
  permutation.append(serializedDimensions.begin(), serializedDimensions.end());
  return replacePLoopOp(rewriter, pLoopOp, newPLoopBounds, forBounds,
                        permutation);
}

/// Serialize all inner dimensions of a `pLoopOp` starting from `serializeFrom`.
static Operation *serializeDimensionsFrom(ConversionPatternRewriter &rewriter,
                                          scf::ParallelOp pLoopOp,
                                          unsigned serializeFrom) {
  unsigned numLoops = pLoopOp.getNumLoops();
  assert(serializeFrom < numLoops &&
         "unhandled corner case of no serialization");
  SmallVector<unsigned, 2> serializedDimensions;
  for (unsigned dim : llvm::seq(serializeFrom, numLoops))
    serializedDimensions.push_back(dim);
  return serializeDimensions(rewriter, pLoopOp, serializedDimensions);
}

/// Collapses all loops in a scf.parallel into one scf.parallel operation. This
/// is done by
/// 1) Normalize the loop bounds to be [0, (ub - lb) / step)
/// 2) Compute the total number of iterations.
/// 3) From the induction variable of the modified loop, compute the values of
///    the original induction variables by de-linearization.
scf::ParallelOp collapseParallelLoops(ConversionPatternRewriter &rewriter,
                                      scf::ParallelOp pLoopOp) {
  if (pLoopOp.getNumReductions()) return nullptr;

  unsigned numLoops = pLoopOp.getNumLoops();
  if (numLoops == 1) return pLoopOp;

  // Compute the number of iterations of each loops starting from the innermost.
  Location loc = pLoopOp.getLoc();
  Value totalNumIterations = rewriter.create<ConstantIndexOp>(loc, 1);

  // Track the "stride" of each loop, i.e. product of the total number of
  // iterations of the inner loops.
  SmallVector<Value, 2> iterationStride;
  iterationStride.resize(pLoopOp.getNumLoops());
  auto lbs = pLoopOp.lowerBound();
  auto ubs = pLoopOp.upperBound();
  auto steps = pLoopOp.step();
  for (int i = numLoops - 1; i >= 0; --i) {
    Value lb = lbs[i], ub = ubs[i], step = steps[i];
    Value iterCount = rewriter.create<SignedDivIOp>(
        loc, rewriter.create<SubIOp>(loc, ub, lb), step);
    iterationStride[i] = totalNumIterations;
    totalNumIterations =
        rewriter.create<MulIOp>(loc, totalNumIterations, iterCount);
  }

  // Create the collapsed parallel loop op with lowerbound 0, step 1 and upper
  // bound being the totalNumIterations.
  Value newLb = rewriter.create<ConstantIndexOp>(loc, 0);
  Value newStep = rewriter.create<ConstantIndexOp>(loc, 1);
  scf::ParallelOp newPLoopOp =
      rewriter.create<scf::ParallelOp>(loc, newLb, totalNumIterations, newStep);

  // Build the body of the collapsed loop by cloning the original loop body. The
  // replacement value of the induction variables of the original loop body,
  // from the induction variable of the new loop, using
  //   origLoopIv[i] = loopIv / iterationStride[i]
  //   loopIv = loopIv % iterationStride[i]
  OpBuilder::InsertionGuard guard(rewriter);
  Block &pLoopBody = pLoopOp.getLoopBody().front();
  rewriter.setInsertionPointToStart(&newPLoopOp.getLoopBody().front());
  Value loopIv = *newPLoopOp.getInductionVars().begin();
  BlockAndValueMapping map;
  edsc::ScopedContext scope(rewriter, loc);
  using namespace edsc::op;
  for (int i : llvm::seq<int>(0, numLoops)) {
    Value iterNum =
        rewriter.create<SignedDivIOp>(loc, loopIv, iterationStride[i]);
    Value newIv = lbs[i] + (iterNum * steps[i]);
    map.map(pLoopBody.getArgument(i), newIv);
    loopIv = rewriter.create<SignedRemIOp>(loc, loopIv, iterationStride[i]);
  }
  for (Operation &op : pLoopBody.without_terminator()) {
    rewriter.clone(op, map);
  }
  rewriter.eraseOp(pLoopOp);
  return newPLoopOp;
}

//===----------------------------------------------------------------------===//
// GPU processor ID mapping utilities
//===----------------------------------------------------------------------===//

/// Distributes scf.parallel to processors with the processors logically
/// arranged with same dimensionality as the number of loops, i.e. a
/// scf.parallel with 2 loops to a 2D grid of processors. `processorIDs` and
/// `numProcessors` must be of same size as the number of loops and are the
/// values to use for process ID and number of processors along each dimension
/// in the distributed code.
/// This method accounts for the case where the number of processors is not
/// enough to execute the entire iteration space with one iteration mapped to
/// each processor. So implements a cyclic distribution of iterations to
/// processors.
static LogicalResult distributeCyclicallyToProcessors(
    ConversionPatternRewriter &rewriter, scf::ParallelOp pLoopOp,
    ArrayRef<linalg::ProcInfo> procInfo) {
  unsigned numLoops = pLoopOp.getNumLoops();
  assert(numLoops == procInfo.size() &&
         "expected as many ids as number of loops");
  SmallVector<LoopBounds, 2> forBounds;
  SmallVector<unsigned, 2> permutation;
  forBounds.reserve(numLoops);
  permutation.reserve(numLoops);
  Location loc = pLoopOp.getLoc();
  auto lbs = pLoopOp.lowerBound(), ubs = pLoopOp.upperBound(),
       steps = pLoopOp.step();
  for (unsigned i : llvm::seq<unsigned>(0, procInfo.size())) {
    Value mappedLb = rewriter.create<AddIOp>(
        loc, lbs[i],
        rewriter.create<MulIOp>(loc, steps[i], procInfo[i].procId));
    Value mappedStep =
        rewriter.create<MulIOp>(loc, steps[i], procInfo[i].nprocs);
    forBounds.push_back({mappedLb, ubs[i], mappedStep});
    permutation.push_back(i);
  }
  replacePLoopOp(rewriter, pLoopOp, /*newPLoopBounds=*/{}, forBounds,
                 permutation);
  return success();
}

/// Distributes scf.parallel to processors with the processors logically
/// arranged with same dimensionality as the number of loops, i.e. a
/// scf.parallel with 2 loops to a 2D grid of processors. `processorIDs` must be
/// of same size as the number of loops and are the values to use for process ID
/// and number of processors along each dimension in the distributed code.  This
/// method assumes that the number of processors is greater than or equal to the
/// number of iterations. So just generates an if statement to mask of
/// processors with no work. When the number of processors is known to be
/// exactly equal to the number of iterations, the if statement is not needed as
/// well. In such cases, `generateGuard` can be set to `false` to avoid
/// generating the if statement.
static LogicalResult distributeSingleIterationPerProcessor(
    ConversionPatternRewriter &rewriter, scf::ParallelOp pLoopOp,
    ArrayRef<linalg::ProcInfo> procInfo, bool generateGuard = false) {
  unsigned numLoops = pLoopOp.getNumLoops();
  Location loc = pLoopOp.getLoc();
  assert(numLoops == procInfo.size() &&
         "expected as many ids as number of loops");

  auto lbs = pLoopOp.lowerBound();
  auto step = pLoopOp.step();
  SmallVector<Value, 2> ivReplacements;
  for (unsigned i : llvm::seq<unsigned>(0, numLoops)) {
    Value iterValue = rewriter.create<AddIOp>(
        loc, lbs[i], rewriter.create<MulIOp>(loc, procInfo[i].procId, step[i]));
    ivReplacements.push_back(iterValue);
  }
  Region &pLoopOpRegion = pLoopOp.getLoopBody();

  if (generateGuard) {
    TypeConverter::SignatureConversion signatureConverter(numLoops);
    Value cond = nullptr;
    auto ubs = pLoopOp.upperBound();
    for (unsigned i : llvm::seq<unsigned>(0, numLoops)) {
      Value cmp = rewriter.create<CmpIOp>(loc, CmpIPredicate::slt,
                                          ivReplacements[i], ubs[i]);
      cond = (cond ? rewriter.create<AndOp>(loc, cond, cmp) : cmp);
      signatureConverter.remapInput(i, ivReplacements[i]);
    }
    rewriter.applySignatureConversion(&pLoopOpRegion, signatureConverter);
    scf::IfOp ifOp = buildEmptyIfOp(loc, rewriter, cond);
    Region &ifOpRegion = ifOp.getRegion(0);
    rewriter.inlineRegionBefore(pLoopOpRegion, ifOpRegion, ifOpRegion.begin());
  } else {
    // The body of the scf.parallel needs to be moved into its parent
    // operation.
    // - Split the block just before the scf.parallel operation.
    // - Move the only block of scf.parallel before the newly created block
    //   (after signature conversion).
    // - Add branch from the original block to the moved block of the
    //   scf.parallel's region, and from the latter to the block created by the
    //   split operation.
    // - Canonicalization will fold these branches away.
    Block *destBlock = pLoopOp.getOperation()->getBlock();
    Block *remainingInst =
        rewriter.splitBlock(destBlock, Block::iterator(pLoopOp));
    Block *sourceBlock = &pLoopOpRegion.front();
    rewriter.eraseOp(sourceBlock->getTerminator());
    rewriter.mergeBlocks(&pLoopOpRegion.front(), destBlock, ivReplacements);
    rewriter.mergeBlocks(remainingInst, destBlock, {});
  }
  rewriter.eraseOp(pLoopOp);
  return success();
}

template <typename GPUIdOp, typename GPUCountOp>
static linalg::ProcInfo getLinearizedGPUProcessorIdAndCount(
    Location loc, ConversionPatternRewriter &rewriter) {
  SmallVector<linalg::ProcInfo, 3> procInfo =
      getGPUProcessorIdsAndCounts<GPUIdOp, GPUCountOp>(rewriter, loc,
                                                       kNumGPUDims);
  linalg::ProcInfo linearized;
  linearized.procId = procInfo[0].procId;
  linearized.nprocs = procInfo[0].nprocs;
  for (unsigned i = 0; i < kNumGPUDims - 1; ++i) {
    linearized.procId =
        rewriter.create<MulIOp>(loc, linearized.procId, procInfo[i + 1].nprocs);
    linearized.procId =
        rewriter.create<AddIOp>(loc, linearized.procId, procInfo[i + 1].procId);
    linearized.nprocs =
        rewriter.create<MulIOp>(loc, linearized.nprocs, procInfo[i + 1].nprocs);
  }
  return linearized;
}

/// Distributes scf.parallel to processors where `IdOp` is used to get the
/// processor ID and `DimOp` is used to get the number of processors along a
/// dimension.
template <typename GPUIdOp, typename GPUCountOp>
static LogicalResult distributeCyclicallyToProcessors(
    ConversionPatternRewriter &rewriter, scf::ParallelOp pLoopOp) {
  unsigned numLoops = pLoopOp.getNumLoops();
  if (numLoops > 3) {
    pLoopOp =
        cast<scf::ParallelOp>(serializeDimensionsFrom(rewriter, pLoopOp, 3));
    numLoops = 3;
  }
  SmallVector<linalg::ProcInfo, 2> procInfo =
      getGPUProcessorIdsAndCounts<GPUIdOp, GPUCountOp>(
          rewriter, pLoopOp.getLoc(), numLoops);
  return distributeCyclicallyToProcessors(rewriter, pLoopOp, procInfo);
}

/// Distributes scf.parallel to processors where `IdOp` is used to get the
/// processor ID and `DimOp` is used to get the number of processors along a
/// dimension. Assumes that the number of processors will be less than equal to
/// the number of iterations of the pLoopOp along all dimensions.
template <typename GPUIdOp, typename GPUCountOp>
static LogicalResult distributeSingleIterationPerProcessor(
    ConversionPatternRewriter &rewriter, scf::ParallelOp pLoopOp,
    bool generateGuard = true) {
  unsigned numLoops = pLoopOp.getNumLoops();
  if (numLoops > 3) {
    pLoopOp =
        cast<scf::ParallelOp>(serializeDimensionsFrom(rewriter, pLoopOp, 3));
    numLoops = 3;
  }
  auto procInfo = getGPUProcessorIdsAndCounts<GPUIdOp, GPUCountOp>(
      rewriter, pLoopOp.getLoc(), numLoops);
  return distributeSingleIterationPerProcessor(rewriter, pLoopOp, procInfo,
                                               generateGuard);
}

/// Distribute the scf.parallel to workgroups.
static LogicalResult mapToWorkgroups(ConversionPatternRewriter &rewriter,
                                     scf::ParallelOp pLoopOp,
                                     bool useCyclicDistribution = false) {
  if (useCyclicDistribution) {
    return distributeCyclicallyToProcessors<gpu::BlockIdOp, gpu::GridDimOp>(
        rewriter, pLoopOp);
  }
  return distributeSingleIterationPerProcessor<gpu::BlockIdOp, gpu::GridDimOp>(
      rewriter, pLoopOp, false);
}

/// Distributes scf.parallel to workitems using local invocation ID.
static LogicalResult mapToLocalInvocationId(
    ConversionPatternRewriter &rewriter, scf::ParallelOp pLoopOp,
    bool useCyclicDistribution = false) {
  if (useCyclicDistribution) {
    return distributeCyclicallyToProcessors<gpu::ThreadIdOp, gpu::BlockDimOp>(
        rewriter, pLoopOp);
  }
  return distributeSingleIterationPerProcessor<gpu::ThreadIdOp,
                                               gpu::BlockDimOp>(rewriter,
                                                                pLoopOp);
}

/// Distributes scf.parallel to workitems using global invocation ID. The GPU
/// dialect doesn't have a direct operation to do this. This could be done using
/// id = blockIdx * blockDim + gridIdx. count = blockDim * gridDim.
static LogicalResult mapToGlobalInvocationId(
    ConversionPatternRewriter &rewriter, scf::ParallelOp pLoopOp) {
  return distributeSingleIterationPerProcessor<GPUGlobalId, GPUGlobalCount>(
      rewriter, pLoopOp);
}

/// Returns the number of bytes copied when loading to/storing from workgorup
/// memory. It is approximated to be the size of the underlying allocation being
/// copied into/from.
static Optional<int64_t> getLinearizedCopySize(linalg::CopyOp copyOp) {
  Value src = copyOp.input();
  Value dst = copyOp.output();
  MemRefType srcType = src.getType().cast<MemRefType>();
  MemRefType dstType = dst.getType().cast<MemRefType>();

  Value workgroupMemoryView;
  MemRefType workgroupMemoryType;
  if (srcType.getMemorySpace() == getWorkgroupMemorySpace()) {
    workgroupMemoryView = src;
    workgroupMemoryType = srcType;
  } else if (dstType.getMemorySpace() == getWorkgroupMemorySpace()) {
    workgroupMemoryView = dst;
    workgroupMemoryType = dstType;
  } else {
    return {};
  }

  SubViewOp workgroupMemorySubviewOp =
      dyn_cast_or_null<SubViewOp>(workgroupMemoryView.getDefiningOp());
  if (!workgroupMemorySubviewOp) return {};
  AllocOp allocOp = dyn_cast_or_null<AllocOp>(
      workgroupMemorySubviewOp.source().getDefiningOp());
  if (!allocOp) return {};

  MemRefType allocOpType = allocOp.getType();
  if (!allocOpType.hasStaticShape()) return {};
  return allocOpType.getNumElements();
}

//===----------------------------------------------------------------------===//
// Pass and patterns.
//===----------------------------------------------------------------------===//

namespace {
/// Pass to convert from tiled and fused linalg ops into gpu.func.
class ConvertToGPUPass
    : public PassWrapper<ConvertToGPUPass,
                         OperationPass<IREE::HAL::ExecutableTargetOp>> {
 public:
  ConvertToGPUPass(const SPIRVCodegenOptions &passOptions)
      : options(passOptions) {}
  ConvertToGPUPass(const ConvertToGPUPass &pass) : options(pass.options) {}

  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<AffineDialect, gpu::GPUDialect, scf::SCFDialect,
                    ShapeDialect>();
  }
  void runOnOperation() override;

 private:
  SPIRVCodegenOptions options;
};

struct SerializeParallelLoopPattern
    : public OpConversionPattern<scf::ParallelOp> {
  using OpConversionPattern<scf::ParallelOp>::OpConversionPattern;
  LogicalResult matchAndRewrite(
      scf::ParallelOp pLoopOp, ArrayRef<Value> operands,
      ConversionPatternRewriter &rewriter) const override {
    return success(serializeDimensionsFrom(rewriter, pLoopOp, 0) != nullptr);
  }
};

/// Implementation of the mapping of tiled linalg op to workitems within a
/// workgroup.
template <typename LinalgOpTy>
static LogicalResult mapLinalgOpToLocalInvocationIdImpl(
    LinalgOpTy linalgOp, ArrayRef<Value> operands,
    ConversionPatternRewriter &rewriter, bool optimizeControlFlow) {
  // Check for marker that specifies that the linalg op is to be partitioned
  // across threads within a workgroup.
  if (!hasMarker(linalgOp)) return failure();
  Optional<linalg::LinalgLoops> loops =
      linalg::linalgLowerOpToLoops<scf::ParallelOp>(rewriter, linalgOp);
  if (!loops) return failure();
  if (loops.getValue().empty()) return success();

  auto pLoopOp = cast<scf::ParallelOp>(loops.getValue()[0]);
  return mapToLocalInvocationId(rewriter, pLoopOp, optimizeControlFlow);
}

static LogicalResult distributeCopyOp(linalg::CopyOp copyOp,
                                      scf::ParallelOp pLoopOp,
                                      ConversionPatternRewriter &rewriter) {
  pLoopOp = collapseParallelLoops(rewriter, pLoopOp);
  if (!pLoopOp) return failure();

  Optional<int64_t> copyLength = getLinearizedCopySize(copyOp);
  linalg::ProcInfo idAndCount =
      getLinearizedGPUProcessorIdAndCount<gpu::ThreadIdOp, gpu::BlockDimOp>(
          copyOp.getLoc(), rewriter);
  auto workgroupSize =
      spirv::lookupLocalWorkGroupSize(copyOp).getValues<APInt>();
  int64_t linearizedWorkgroupSize = std::accumulate(
      workgroupSize.begin(), workgroupSize.end(), 1,
      [](int64_t total, APInt value) { return total * value.getSExtValue(); });

  if (copyLength.hasValue() && !workgroupSize.empty() &&
      copyLength.getValue() <= linearizedWorkgroupSize) {
    return distributeSingleIterationPerProcessor(rewriter, pLoopOp, idAndCount,
                                                 /*generateGuard=*/true);
  }
  return distributeCyclicallyToProcessors(rewriter, pLoopOp, idAndCount);
}

/// CopyOp that are loading to/storing from workgroup memory are special cased
/// to use all workitems to do a copy. This is done by linearizing the copy
/// operation.
// TODO(ravishankarm): This linearization is achieved through collapsing the
// generated parallel loops from a multi-dimensional copy. Such lowering results
// in mods/divs in the collapsed loop body. This can be removed by reshaping the
// copy to be a 1D copy. This seems to be hitting an error in reshape
// canonicalization. Investigate this further.
template <>
LogicalResult mapLinalgOpToLocalInvocationIdImpl<linalg::CopyOp>(
    linalg::CopyOp copyOp, ArrayRef<Value> operands,
    ConversionPatternRewriter &rewriter, bool optimizeControlFlow) {
  if (!hasMarker(copyOp,
                 {getCopyToWorkgroupMemoryMarker(), getWorkgroupMarker()}))
    return failure();
  Optional<linalg::LinalgLoops> loops =
      linalg::linalgLowerOpToLoops<scf::ParallelOp>(rewriter, copyOp);
  if (!loops) return failure();
  if (loops.getValue().empty()) return success();

  auto pLoopOp = cast<scf::ParallelOp>(loops.getValue()[0]);
  if (hasMarker(copyOp, getWorkgroupMarker())) {
    return mapToLocalInvocationId(rewriter, pLoopOp, optimizeControlFlow);
  }
  return distributeCopyOp(copyOp, pLoopOp, rewriter);
}

/// Map tiled linalg op to workitems by lowering it to scf.parallel and
/// partitioning it to workitems.
template <typename LinalgOpTy>
struct MapLinalgOpToLocalInvocationId : public OpConversionPattern<LinalgOpTy> {
  MapLinalgOpToLocalInvocationId(MLIRContext *context,
                                 bool usingLinalgOnTensorsPath,
                                 PatternBenefit benefit = 1)
      : OpConversionPattern<LinalgOpTy>(context, benefit),
        usingLinalgOnTensorsPath(usingLinalgOnTensorsPath) {}
  LogicalResult matchAndRewrite(
      LinalgOpTy linalgOp, ArrayRef<Value> operands,
      ConversionPatternRewriter &rewriter) const override {
    if (failed(mapLinalgOpToLocalInvocationIdImpl(linalgOp, operands, rewriter,
                                                  usingLinalgOnTensorsPath)))
      return failure();

    // If the `linalgOp` writes to workgroup memory insert barrier after the
    // op.
    if (llvm::any_of(linalgOp.getOperands(), [](Value output) {
          MemRefType outputType = output.getType().dyn_cast<MemRefType>();
          return outputType &&
                 outputType.getMemorySpace() == getWorkgroupMemorySpace();
        })) {
      rewriter.create<spirv::ControlBarrierOp>(
          linalgOp.getLoc(), spirv::Scope::Workgroup, spirv::Scope::Workgroup,
          spirv::MemorySemantics::AcquireRelease);
    }
    rewriter.eraseOp(linalgOp);
    return success();
  }

 private:
  /// Flag to signify if Linalg on tensors path is being used. The control flow
  /// optimizations implemented on legacy path seems to be failing on this
  /// path. Assuming this overhead is not too much, for now just generated the
  /// extra loops.
  bool usingLinalgOnTensorsPath;
};

/// Given the workload return the workgroup count along X obtained by
/// linearizing the workload and dividing by the workgroup size.
static Value getWorkgroupCountX(OpBuilder &builder, Location loc,
                                ArrayRef<Value> values,
                                int64_t workgroupSizeX) {
  AffineExpr expr = builder.getAffineConstantExpr(1);
  for (auto val : enumerate(values)) {
    expr = expr * builder.getAffineSymbolExpr(val.index());
  }
  expr = expr.ceilDiv(workgroupSizeX);
  return linalg::applyMapToValues(
      builder, loc, AffineMap::get(0, values.size(), expr), values)[0];
}

/// Map linalg operation to execute on GPU in parallel by mapping the parallel
/// loops to "GlobalInvocationId".
template <typename LinalgOpTy>
struct MapLinalgOpToGlobalInvocationId
    : public OpConversionPattern<LinalgOpTy> {
  MapLinalgOpToGlobalInvocationId(MLIRContext *context,
                                  bool usingLinalgOnTensorsPath,
                                  PatternBenefit benefit = 1)
      : OpConversionPattern<LinalgOpTy>(context, benefit),
        usingLinalgOnTensorsPath(usingLinalgOnTensorsPath) {}

  LogicalResult matchAndRewrite(
      LinalgOpTy linalgOp, ArrayRef<Value> operands,
      ConversionPatternRewriter &rewriter) const override {
    // If marker exists do nothing.
    if (hasMarker(linalgOp)) return failure();
    FuncOp funcOp = linalgOp->template getParentOfType<FuncOp>();
    if (!funcOp) return failure();
    Optional<linalg::LinalgLoops> loops =
        linalg::linalgLowerOpToLoops<scf::ParallelOp>(rewriter, linalgOp);
    if (!loops) return failure();

    SmallVector<int64_t, 3> workgroupSize(3, 1);
    if (!loops.getValue().empty()) {
      scf::ParallelOp pLoopOp = dyn_cast<scf::ParallelOp>(loops.getValue()[0]);
      // If there are parallel loops partition them to threads using global
      // invocation ID.
      if (pLoopOp) {
        pLoopOp = collapseParallelLoops(rewriter, pLoopOp);
        if (!pLoopOp) return failure();
        if (failed(mapToGlobalInvocationId(rewriter, pLoopOp)))
          return rewriter.notifyMatchFailure(
              linalgOp, "mapping to GlobalInvocationID failed");
        workgroupSize = {32, 1, 1};
      }
    }
    if (usingLinalgOnTensorsPath) {
      WorkgroupCountRegionBuilder regionBuilder =
          [&workgroupSize](
              OpBuilder &b, Location loc,
              std::array<Value, 3> workload) -> std::array<Value, 3> {
        Value one = b.create<ConstantIndexOp>(loc, 1);
        return {getWorkgroupCountX(b, loc, workload, workgroupSize[0]), one,
                one};
      };
      if (failed(defineWorkgroupCountRegion(rewriter, funcOp, regionBuilder))) {
        return failure();
      }
    } else {
      // TODO (GH-4901): Only support static shapes on this path. This should be
      // removed when moved to linalg on tensors.
      Optional<SmallVector<int64_t, 4>> staticLoopRange =
          linalg::getStaticLoopRanges(linalgOp);
      if (!staticLoopRange ||
          llvm::any_of(staticLoopRange.getValue(), [](int64_t d) {
            return d == ShapedType::kDynamicSize;
          })) {
        return linalgOp.emitError("failed to find statlc loop bounds");
      }
      ArrayRef<int64_t> parallelLoopRange(staticLoopRange.getValue());
      unsigned numOuterParallel = getNumOuterParallelLoops(linalgOp);
      parallelLoopRange = parallelLoopRange.take_front(numOuterParallel);
      WorkgroupCountRegionBuilder regionBuilder =
          [&parallelLoopRange, &workgroupSize](
              OpBuilder &b, Location loc,
              std::array<Value, 3> workload) -> std::array<Value, 3> {
        Value one = b.create<ConstantIndexOp>(loc, 1);
        auto values = llvm::to_vector<4>(
            llvm::map_range(parallelLoopRange, [&](int64_t dim) -> Value {
              return b.create<ConstantIndexOp>(loc, dim);
            }));
        return {getWorkgroupCountX(b, loc, values, workgroupSize[0]), one, one};
      };
      if (failed(defineWorkgroupCountRegion(rewriter, funcOp, regionBuilder))) {
        return failure();
      }
    }
    if (failed(updateWorkGroupSize(funcOp, workgroupSize))) {
      return failure();
    }
    rewriter.eraseOp(linalgOp);
    return success();
  }

 private:
  /// Flag to signify if Linalg on tensors path is being used. This changes the
  /// way the number of workgroups is computed. With the linalg on tensors path,
  /// the hal.executable.entry_point will be updated to contain a region that
  /// gives the number of workgroups to use.
  bool usingLinalgOnTensorsPath;
};

/// Remove the linalg.range operation created when lowering to loops.
struct RemoveLinalgRange : public OpConversionPattern<linalg::RangeOp> {
  using OpConversionPattern<linalg::RangeOp>::OpConversionPattern;
  LogicalResult matchAndRewrite(
      linalg::RangeOp rangeOp, ArrayRef<Value> operands,
      ConversionPatternRewriter &rewriter) const override {
    if (!rangeOp.getResult().use_empty()) return failure();
    rewriter.eraseOp(rangeOp);
    return success();
  }
};
}  // namespace

// Applies tiling followed to load/store optimized size then distribute on
// incovations.
static LogicalResult linalgCopyTileAndDistribute(
    linalg::CopyOp copyOp, ArrayRef<Value> operands,
    ConversionPatternRewriter &rewriter) {
  linalg::LinalgTilingOptions options;
  // Tile to memory access of 128bits as those tend to be optimal on most GPUs.
  constexpr unsigned vecLoadBits = 128;
  unsigned elementBits =
      copyOp.getSource().getType().cast<MemRefType>().getElementTypeBitWidth();
  if (elementBits == 0 || vecLoadBits % elementBits != 0) return failure();
  unsigned numElement = vecLoadBits / elementBits;
  options.setTileSizes({1, numElement})
      .setLoopType(linalg::LinalgTilingLoopType::ParallelLoops);
  Optional<linalg::TiledLinalgOp> tiledOp = linalg::tileLinalgOp(
      rewriter, cast<linalg::LinalgOp>(copyOp.getOperation()), options);
  if (!tiledOp) return failure();
  if (tiledOp->loops.empty()) return success();
  setMarker(tiledOp->op, getVectorizeMarker());
  auto pLoopOp = cast<scf::ParallelOp>(tiledOp->loops[0]);
  return distributeCopyOp(copyOp, pLoopOp, rewriter);
}

namespace {
// Pattern to tile and distribute linalg::CopyOp.
struct TileAndDistributeCopyOp : public OpConversionPattern<linalg::CopyOp> {
  using OpConversionPattern<linalg::CopyOp>::OpConversionPattern;
  LogicalResult matchAndRewrite(
      linalg::CopyOp linalgOp, ArrayRef<Value> operands,
      ConversionPatternRewriter &rewriter) const override {
    if (!hasMarker(linalgOp, getCopyToWorkgroupMemoryMarker())) {
      return failure();
    }
    if (failed(linalgCopyTileAndDistribute(linalgOp, operands, rewriter))) {
      return failure();
    }

    // Insert a barrier if read or write shared memory.
    if (llvm::any_of(linalgOp.getOperands(), [](Value output) {
          return output.getType().cast<MemRefType>().getMemorySpace() ==
                 getWorkgroupMemorySpace();
        })) {
      rewriter.create<spirv::ControlBarrierOp>(
          linalgOp.getLoc(), spirv::Scope::Workgroup, spirv::Scope::Workgroup,
          spirv::MemorySemantics::AcquireRelease);
    }
    rewriter.eraseOp(linalgOp);
    return success();
  }
};
}  // namespace

void populateLinalgTileAndDistributePatterns(
    MLIRContext *context, OwningRewritePatternList &patterns) {
  patterns.insert<TileAndDistributeCopyOp>(context);
}

void ConvertToGPUPass::runOnOperation() {
  MLIRContext *context = &getContext();
  ConversionTarget target(*context);
  // After this pass Linalg and scf.parallel ops should be gone.
  target.addIllegalOp<scf::ParallelOp>();
  target.addIllegalDialect<linalg::LinalgDialect>();
  // Reshape ops are treated legal since they just change the way the underlying
  // buffer is viewed. These are legalized downstream. They become no ops when
  // lowering to SPIR-V since the SPIR-V code uses linearized arrays.
  target.addLegalOp<linalg::ReshapeOp>();
  // Let the rest fall through.
  target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });

  OwningRewritePatternList patterns;

  patterns.insert<
      MapLinalgOpToGlobalInvocationId<linalg::CopyOp>,
      MapLinalgOpToGlobalInvocationId<linalg::FillOp>,
      MapLinalgOpToGlobalInvocationId<linalg::GenericOp>,
      MapLinalgOpToGlobalInvocationId<linalg::IndexedGenericOp>,
      MapLinalgOpToLocalInvocationId<linalg::ConvOp>,
      MapLinalgOpToLocalInvocationId<linalg::ConvInputNWCFilterWCFOp>,
      MapLinalgOpToLocalInvocationId<linalg::ConvInputNHWCFilterHWCFOp>,
      MapLinalgOpToLocalInvocationId<linalg::ConvInputNDHWCFilterDHWCFOp>,
      MapLinalgOpToLocalInvocationId<linalg::CopyOp>,
      MapLinalgOpToLocalInvocationId<linalg::FillOp>,
      MapLinalgOpToLocalInvocationId<linalg::GenericOp>,
      MapLinalgOpToLocalInvocationId<linalg::IndexedGenericOp>,
      MapLinalgOpToLocalInvocationId<linalg::MatmulOp>,
      MapLinalgOpToLocalInvocationId<linalg::BatchMatmulOp>,
      MapLinalgOpToLocalInvocationId<linalg::PoolingMaxOp>,
      MapLinalgOpToLocalInvocationId<linalg::PoolingMinOp>,
      MapLinalgOpToLocalInvocationId<linalg::PoolingSumOp>, RemoveLinalgRange,
      SerializeParallelLoopPattern>(context, options.usingLinalgOnTensors);
  FrozenRewritePatternList frozenPatterns(std::move(patterns));

  for (FuncOp funcOp : getOperation().getInnerModule().getOps<FuncOp>()) {
    if (!isEntryPoint(funcOp)) continue;
    Region &body = funcOp.getBody();
    if (!llvm::hasSingleElement(body)) {
      funcOp.emitError("unhandled dispatch function with multiple blocks");
      return signalPassFailure();
    }
    if (failed(applyFullConversion(funcOp, target, frozenPatterns)))
      return signalPassFailure();
  }
}

std::unique_ptr<OperationPass<IREE::HAL::ExecutableTargetOp>>
createConvertToGPUPass(const SPIRVCodegenOptions &options) {
  return std::make_unique<ConvertToGPUPass>(options);
}

static PassRegistration<ConvertToGPUPass> pass(
    "iree-codegen-convert-to-gpu", "Map tiled linalg and loop ops to GPU", [] {
      SPIRVCodegenOptions options = getSPIRVCodegenOptionsFromClOptions();
      return std::make_unique<ConvertToGPUPass>(options);
    });

}  // namespace iree_compiler
}  // namespace mlir