// Copyright 2022 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/Codegen/Common/GPU/Passes.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenAttrs.h"
#include "iree/compiler/Codegen/Utils/GPUUtils.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "mlir/Dialect/AMDGPU/IR/AMDGPUDialect.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
#include "mlir/Dialect/Vector/Transforms/VectorDistribution.h"
#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

#define DEBUG_TYPE "iree-codegen-vector-reduction-to-gpu"

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_VECTORREDUCTIONTOGPUPASS
#include "iree/compiler/Codegen/Common/GPU/Passes.h.inc"

namespace {
static void debugPrint(Operation *op, const char *message) {
  LLVM_DEBUG({
    llvm::dbgs() << "//--- " << message << " ---//\n";
    op->print(llvm::dbgs(), OpPrintingFlags().useLocalScope());
    llvm::dbgs() << "\n\n";
  });
}

/// Emit shared local memory allocation in case it is needed when lowering the
/// warp operations.
static Value allocateGlobalSharedMemory(Location loc, OpBuilder &builder,
                                        gpu::WarpExecuteOnLane0Op warpOp,
                                        Type type) {
  MemRefType memrefType;
  auto addressSpaceAttr = gpu::AddressSpaceAttr::get(
      builder.getContext(), gpu::GPUDialect::getWorkgroupAddressSpace());
  if (auto vectorType = dyn_cast<VectorType>(type)) {
    memrefType =
        MemRefType::get(vectorType.getShape(), vectorType.getElementType(),
                        MemRefLayoutAttrInterface{}, addressSpaceAttr);
  } else {
    memrefType = MemRefType::get({1}, type, MemRefLayoutAttrInterface{},
                                 addressSpaceAttr);
  }
  return memref::AllocOp::create(builder, loc, memrefType);
}

/// Returns true if the given op is a memref.load from a uniform buffer or
/// read-only storage buffer.
static bool isUniformLoad(Operation *op) {
  using namespace IREE::HAL;

  auto loadOp = dyn_cast<memref::LoadOp>(op);
  if (!loadOp) {
    return false;
  }
  if (!hasGlobalMemoryAddressSpace(loadOp.getMemRefType())) {
    return false;
  }
  auto space = loadOp.getMemRefType().getMemorySpace();
  auto descTypeAttr = dyn_cast_if_present<DescriptorTypeAttr>(space);
  if (descTypeAttr &&
      descTypeAttr.getValue() == DescriptorType::UniformBuffer) {
    return true;
  }

  auto subspan = loadOp.getMemRef().getDefiningOp<InterfaceBindingSubspanOp>();
  if (auto fatBufferCast =
          loadOp.getMemRef().getDefiningOp<amdgpu::FatRawBufferCastOp>()) {
    subspan =
        fatBufferCast.getSource().getDefiningOp<InterfaceBindingSubspanOp>();
  }
  if (!subspan) {
    return false;
  }

  descTypeAttr = dyn_cast_if_present<DescriptorTypeAttr>(
      cast<MemRefType>(subspan.getResult().getType()).getMemorySpace());
  if (descTypeAttr &&
      descTypeAttr.getValue() == DescriptorType::UniformBuffer) {
    return true;
  }
  if (auto flags = subspan.getDescriptorFlags()) {
    if (bitEnumContainsAll(*flags, IREE::HAL::DescriptorFlags::ReadOnly)) {
      return true;
    }
  }
  return false;
}

/// Hoist uniform operations as well as special hal operations that have side
/// effect but are safe to move out of the warp single lane region.
static void moveScalarAndBindingUniformCode(gpu::WarpExecuteOnLane0Op warpOp) {
  /// Hoist ops without side effect as well as special binding ops.
  auto canBeHoisted = [](Operation *op,
                         function_ref<bool(Value)> definedOutside) {
    if (op->getNumRegions() != 0) {
      return false;
    }
    if (!llvm::all_of(op->getOperands(), definedOutside)) {
      return false;
    }
    if (isMemoryEffectFree(op)) {
      return true;
    }

    if (isa<IREE::HAL::InterfaceBindingSubspanOp,
            IREE::HAL::InterfaceConstantLoadOp, memref::AssumeAlignmentOp>(
            op)) {
      return true;
    }
    if (isUniformLoad(op)) {
      return true;
    }
    // Shared memory is already scoped to the workgroup and can safely be
    // hoisted out of the the warp op.
    if (auto allocOp = dyn_cast<memref::AllocOp>(op)) {
      if (hasSharedMemoryAddressSpace(allocOp.getType())) {
        return true;
      }
    }

    return false;
  };
  Block *body = warpOp.getBody();

  // Keep track of the ops we want to hoist.
  llvm::SmallSetVector<Operation *, 8> opsToMove;

  // Helper to check if a value is or will be defined outside of the region.
  auto isDefinedOutsideOfBody = [&](Value value) {
    auto *definingOp = value.getDefiningOp();
    return (definingOp && opsToMove.count(definingOp)) ||
           warpOp.isDefinedOutsideOfRegion(value);
  };

  // Do not use walk here, as we do not want to go into nested regions and hoist
  // operations from there.
  for (auto &op : body->without_terminator()) {
    bool hasVectorResult = llvm::any_of(op.getResults(), [](Value result) {
      return isa<VectorType>(result.getType());
    });
    if ((!hasVectorResult || isUniformLoad(&op)) &&
        canBeHoisted(&op, isDefinedOutsideOfBody)) {
      opsToMove.insert(&op);
    }
  }

  // Move all the ops marked as uniform outside of the region.
  for (Operation *op : opsToMove) {
    op->moveBefore(warpOp);
  }
}

/// Pattern to convert single element vector.insert to broadcast, this is a
/// workaround until MultiDimReduction distribution is supported.
struct InsertToBroadcast final : OpRewritePattern<vector::InsertOp> {
  using Base::Base;

  LogicalResult matchAndRewrite(vector::InsertOp insertOp,
                                PatternRewriter &rewriter) const override {
    if (insertOp.getDestVectorType().getNumElements() != 1) {
      return failure();
    }
    rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
        insertOp, insertOp.getDestVectorType(), insertOp.getValueToStore());
    return success();
  }
};

/// Pattern to sink `gpu.barrier` ops out of a `warp_execute_on_lane_0` op.
struct WarpOpBarrier final : OpRewritePattern<gpu::WarpExecuteOnLane0Op> {
  using Base::Base;

  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,
                                PatternRewriter &rewriter) const override {
    auto yield = cast<gpu::YieldOp>(
        warpOp.getBodyRegion().getBlocks().begin()->getTerminator());
    Operation *lastNode = yield->getPrevNode();
    auto barrierOp = dyn_cast_if_present<gpu::BarrierOp>(lastNode);
    if (!barrierOp) {
      return failure();
    }

    rewriter.setInsertionPointAfter(warpOp);
    gpu::BarrierOp::create(rewriter, barrierOp.getLoc(),
                           barrierOp.getAddressSpacesAttr(),
                           /*named_barrier=*/Value{});
    rewriter.eraseOp(barrierOp);
    return success();
  }
};

static Value simpleWarpShuffleFunction(Location loc, OpBuilder &builder,
                                       Value val, Value srcIdx,
                                       int64_t warpSz) {
  assert((val.getType().isF32() || val.getType().isInteger(32)) &&
         "unsupported shuffle type");
  Type i32Type = builder.getIntegerType(32);
  Value srcIdxI32 = arith::IndexCastOp::create(builder, loc, i32Type, srcIdx);
  Value warpSzI32 = arith::ConstantOp::create(
      builder, loc, builder.getIntegerAttr(i32Type, warpSz));
  Value result = gpu::ShuffleOp::create(builder, loc, val, srcIdxI32, warpSzI32,
                                        gpu::ShuffleMode::IDX)
                     .getResult(0);
  return result;
}

struct VectorReductionToGPUPass final
    : impl::VectorReductionToGPUPassBase<VectorReductionToGPUPass> {
  using VectorReductionToGPUPassBase::VectorReductionToGPUPassBase;
  void runOnOperation() override {
    FunctionOpInterface funcOp = getOperation();
    MLIRContext *ctx = &getContext();

    debugPrint(funcOp, "after step #0: before vector reduction to gpu");

    // 1. Pre-process multiDimReductions.
    // TODO: Remove once MultiDimReduce is supported by distribute patterns.
    {
      RewritePatternSet patterns(ctx);
      vector::populateVectorMultiReductionReorderPatterns(
          patterns, vector::VectorMultiReductionLowering::InnerReduction);
      vector::populateVectorMultiReductionFlatteningPatterns(
          patterns, vector::VectorMultiReductionLowering::InnerReduction);
      vector::populateVectorMultiReductionUnrollingPatterns(
          patterns, vector::VectorMultiReductionLowering::InnerReduction);
      // Add clean up patterns after lowering of multidimreduce lowering.
      patterns.add<InsertToBroadcast>(ctx);
      vector::ShapeCastOp::getCanonicalizationPatterns(patterns, ctx);
      vector::BroadcastOp::getCanonicalizationPatterns(patterns, ctx);
      vector::ExtractOp::getCanonicalizationPatterns(patterns, ctx);
      (void)applyPatternsGreedily(getOperation(), std::move(patterns));
    }

    debugPrint(funcOp, "after step #1: preprocessing reduction ops");

    std::optional<SmallVector<int64_t>> maybeWorkgroupSize =
        getWorkgroupSize(funcOp);
    if (!maybeWorkgroupSize) {
      funcOp->emitOpError(
          "expected workgroup size to be set as part of `translation_info`");
      return signalPassFailure();
    }
    SmallVector<int64_t> &workgroupSize = maybeWorkgroupSize.value();
    assert(workgroupSize[1] == 1 && workgroupSize[2] == 1);
    // 2. Create the warp op and move the body into it.
    //
    // When the function body contains a workgroup-level scf.forall, wrap the
    // forall body rather than the function body. The forall body is
    // structurally equivalent to the function body on the non-forall path
    // (same vector ops, just with the forall induction var instead of
    // hal.interface.workgroup.id). Wrapping the function body would place the
    // forall inside the warp op, where distribution patterns cannot propagate
    // through the forall region boundary, leaving vectors at full size.
    const int groupSize = workgroupSize[0];
    Location loc = funcOp.getLoc();
    OpBuilder builder(funcOp);

    // Check for a workgroup-level scf.forall.
    scf::ForallOp workgroupForall;
    for (auto &op : funcOp.getFunctionBody().front()) {
      if (auto forall = dyn_cast<scf::ForallOp>(&op)) {
        if (forallOpHasMappingType<IREE::Codegen::WorkgroupMappingAttr>(
                forall)) {
          workgroupForall = forall;
          break;
        }
      }
    }

    gpu::WarpExecuteOnLane0Op warpOp;
    if (workgroupForall) {
      Block *forallBody = workgroupForall.getBody();
      Operation *terminator = forallBody->getTerminator();

      builder.setInsertionPoint(terminator);
      auto threadX = gpu::ThreadIdOp::create(
          builder, loc, builder.getIndexType(), gpu::Dimension::x);
      auto cstGroupSize =
          arith::ConstantIndexOp::create(builder, loc, groupSize);
      warpOp = gpu::WarpExecuteOnLane0Op::create(
          builder, loc, TypeRange(), threadX.getResult(), groupSize);

      // Move all existing forall body ops into the warp region. The create
      // call above produces a region with one empty block; use that block.
      Block &warpBody = warpOp.getWarpRegion().front();
      SmallVector<Operation *> opsToMove;
      for (auto &op : *forallBody) {
        if (&op == threadX.getOperation() ||
            &op == cstGroupSize.getOperation() ||
            &op == warpOp.getOperation() || &op == terminator) {
          continue;
        }
        opsToMove.push_back(&op);
      }
      for (auto *op : opsToMove) {
        op->moveBefore(&warpBody, warpBody.end());
      }
      builder.setInsertionPointToEnd(&warpBody);
      gpu::YieldOp::create(builder, loc);
    } else {
      auto threadX = gpu::ThreadIdOp::create(
          builder, loc, builder.getIndexType(), gpu::Dimension::x);
      auto cstGroupSize =
          arith::ConstantIndexOp::create(builder, loc, groupSize);
      warpOp = gpu::WarpExecuteOnLane0Op::create(
          builder, loc, TypeRange(), threadX.getResult(), groupSize);
      warpOp.getWarpRegion().takeBody(funcOp.getFunctionBody());
      Block &newBlock = funcOp.getFunctionBody().emplaceBlock();
      threadX->moveBefore(&newBlock, newBlock.end());
      cstGroupSize->moveBefore(&newBlock, newBlock.end());
      warpOp->moveBefore(&newBlock, newBlock.end());
      warpOp.getWarpRegion().getBlocks().back().back().moveBefore(
          &newBlock, newBlock.end());
      builder.setInsertionPointToEnd(
          &warpOp.getWarpRegion().getBlocks().back());
      gpu::YieldOp::create(builder, loc);
    }

    debugPrint(funcOp, "after step #2: wrapping code with the warp execute op");

    // 3. Hoist the scalar code outside of the warp region.
    moveScalarAndBindingUniformCode(warpOp);

    debugPrint(funcOp, "after step #3: hosting uniform code");

    // 4. Distribute transfer write operations and propagate vector
    // distribution.
    {
      std::optional<int> subgroupSize = getGPUSubgroupSize(funcOp);
      if (!subgroupSize) {
        funcOp->emitOpError("missing subgroup size");
        return signalPassFailure();
      }
      auto groupReductionFn = [subgroupSize, &expandSubgroupReduction =
                                                 this->expandSubgroupReduction](
                                  Location loc, OpBuilder &builder, Value input,
                                  vector::CombiningKind kind,
                                  uint32_t size) -> Value {
        return emitGPUGroupReduction(loc, builder, input, kind, size,
                                     *subgroupSize, expandSubgroupReduction);
      };
      auto distributionFn = [](Value val) {
        auto vecType = dyn_cast<VectorType>(val.getType());
        if (!vecType) {
          return AffineMap::get(val.getContext());
        }
        // Create an identity dim map of rank |vecRank|. This greedily divides
        // threads along the outermost vector dimensions to the innermost ones.
        int64_t vecRank = vecType.getRank();
        OpBuilder builder(val.getContext());
        return builder.getMultiDimIdentityMap(vecRank);
      };

      RewritePatternSet patterns(ctx);
      vector::populatePropagateWarpVectorDistributionPatterns(
          patterns, distributionFn, simpleWarpShuffleFunction);
      vector::populateDistributeReduction(patterns, groupReductionFn);

      // We don't want to sink large transfer writes to a single lane -- pick a
      // conservative value based on the group size.
      unsigned maxWriteElementsToExtract = std::max(groupSize / 4, 1);
      vector::populateDistributeTransferWriteOpPatterns(
          patterns, distributionFn, maxWriteElementsToExtract);
      patterns.add<WarpOpBarrier>(patterns.getContext(), 3);
      vector::ReductionOp::getCanonicalizationPatterns(patterns, ctx);
      (void)applyPatternsGreedily(getOperation(), std::move(patterns));
    }

    debugPrint(funcOp, "after step #4: propagating distribution");

    // 5. Lower the remaining WarpExecuteOnLane0 ops.
    {
      RewritePatternSet patterns(ctx);
      vector::WarpExecuteOnLane0LoweringOptions options;
      options.warpAllocationFn = allocateGlobalSharedMemory;
      options.warpSynchronizationFn = [](Location loc, OpBuilder &builder,
                                         gpu::WarpExecuteOnLane0Op warpOp) {
        // There's no communication via global memory occurring, so we only need
        // to fence on workgroup memory.
        gpu::BarrierOp::create(builder, loc, gpu::AddressSpace::Workgroup);
      };
      vector::populateWarpExecuteOnLane0OpToScfForPattern(patterns, options);
      (void)applyPatternsGreedily(getOperation(), std::move(patterns));
    }

    debugPrint(funcOp, "after step #5: lowering remaining ops");
  }
};

} // namespace
} // namespace mlir::iree_compiler
