// Copyright 2021 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

//===- IREEComprehensiveBufferize.cpp - -----------------------------------===//
//
// Wrapper pass to use MLIR's One-Shot Bufferize pass.
//
//===----------------------------------------------------------------------===//

#include "iree/compiler/Codegen/Common/Passes.h"
#include "iree/compiler/Codegen/Common/Transforms.h"
#include "iree/compiler/Codegen/Dialect/Codegen/IR/IREECodegenOps.h"
#include "iree/compiler/Codegen/Interfaces/BufferizationInterfaces.h"
#include "iree/compiler/Codegen/Utils/Utils.h"
#include "iree/compiler/Dialect/HAL/IR/HALOps.h"
#include "iree/compiler/Dialect/LinalgExt/IR/LinalgExtDialect.h"
#include "iree/compiler/Dialect/Util/IR/UtilDialect.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/Bufferization/IR/Bufferization.h"
#include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
#include "mlir/Dialect/Bufferization/Transforms/Passes.h"
#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/Linalg/Transforms/Transforms.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/MemRef/Transforms/Passes.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/Passes.h"

#define DEBUG_TYPE "iree-codegen-linalg-bufferize"

using mlir::bufferization::BufferizationOptions;
using mlir::bufferization::OneShotAnalysisState;
using mlir::bufferization::OneShotBufferizationOptions;

namespace mlir::iree_compiler {

#define GEN_PASS_DEF_ELIMINATEEMPTYTENSORSPASS
#define GEN_PASS_DEF_IREEBUFFERIZECONSTANTSPASS
#define GEN_PASS_DEF_IREECOMPREHENSIVEBUFFERIZEPASS
#include "iree/compiler/Codegen/Common/Passes.h.inc"

// Default allocation functions.
static FailureOr<Value> defaultAllocationFn(OpBuilder &builder, Location loc,
                                            MemRefType allocationType,
                                            ValueRange dynamicSizes,
                                            unsigned int alignment) {
  MemRefType type = allocationType;
  if (auto storage = type.getMemorySpace()) {
    // We cannot allocate to generate a resultant MemRef type with descriptor
    // type memory space; that's runtime allocations. So erase and fallback to
    // the default 0 memory space. It is fine given this is just the default
    // allocator; backends are expected to control by themselves.
    if (isa<IREE::HAL::DescriptorTypeAttr>(storage)) {
      type = MemRefType::get(type.getShape(), type.getElementType(),
                             type.getLayout());
    }
  }
  return memref::AllocOp::create(builder, loc, type, dynamicSizes).getResult();
}
static LogicalResult defaultMemCpyFn(OpBuilder &builder, Location loc,
                                     Value from, Value to) {
  Operation *copyOp = createLinalgCopyOp(builder, loc, from, to);
  return success(static_cast<bool>(copyOp));
}

namespace {
class EliminateEmptyTensorsPass final
    : public impl::EliminateEmptyTensorsPassBase<EliminateEmptyTensorsPass> {
public:
  void getDependentDialects(DialectRegistry &registry) const override {
    // BufferizationDialect is needed for using type interfaces, like
    // TensorLikeType. Because the builtin types, e.g., RankedTensorType, etc.,
    // implement the type interface in
    // bufferization::BufferizationDialect::initialize().
    registry
        .insert<bufferization::BufferizationDialect, tensor::TensorDialect>();
  }

  void runOnOperation() override;
};

/// Pass to convert from tensor based ops to memref based ops.
class IREEComprehensiveBufferizePass final
    : public impl::IREEComprehensiveBufferizePassBase<
          IREEComprehensiveBufferizePass> {
public:
  using Base::Base;
  explicit IREEComprehensiveBufferizePass(
      BufferizationOptions::AllocationFn allocationFn,
      BufferizationOptions::MemCpyFn memCpyFn)
      : allocationFn(allocationFn), memCpyFn(memCpyFn) {}

  void getDependentDialects(DialectRegistry &registry) const override {
    // clang-format off
    registry
        .insert<affine::AffineDialect,
                amdgpu::AMDGPUDialect,
                arith::ArithDialect,
                bufferization::BufferizationDialect,
                func::FuncDialect,
                gpu::GPUDialect,
                IREE::LinalgExt::IREELinalgExtDialect,
                IREE::Util::UtilDialect,
                linalg::LinalgDialect,
                memref::MemRefDialect,
                scf::SCFDialect,
                tensor::TensorDialect,
                vector::VectorDialect>();
    // clang-format on
  }

  void runOnOperation() override;

private:
  const BufferizationOptions::AllocationFn allocationFn = defaultAllocationFn;
  const BufferizationOptions::MemCpyFn memCpyFn = defaultMemCpyFn;
};

/// Pass to convert from tensor based constants to memref.
class IREEBufferizeConstantsPass final
    : public impl::IREEBufferizeConstantsPassBase<IREEBufferizeConstantsPass> {
public:
  using Base::Base;
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<arith::ArithDialect, bufferization::BufferizationDialect,
                    gpu::GPUDialect, memref::MemRefDialect>();
  }

  void runOnOperation() override;
};
} // namespace

static IREEOneShotBufferizationOptions getBufferizationOptions() {
  IREEOneShotBufferizationOptions options;

  // bufferization.to_buffer is used to bufferize constants in IREE. IREE has
  // it's own logic to handle constants. We'd like to leave the arith.constant
  // as is and insert bufferization.to_buffer to convert the tensor to memref.
  options.opFilter.denyOperation<arith::ConstantOp>();

  // This type converter converts tensor types to memref types when no exact
  // memref type can be inferred from the context.
  options.unknownTypeConverterFn =
      [](bufferization::TensorLikeType tensorLikeType, Attribute memorySpace,
         const BufferizationOptions &options) -> bufferization::BufferLikeType {
    auto tensorType = cast<TensorType>(tensorLikeType);
    if (tensorType.hasStaticShape()) {
      return cast<bufferization::BufferLikeType>(
          bufferization::getMemRefTypeWithStaticIdentityLayout(tensorType,
                                                               memorySpace));
    }
    // Default case: Fully dynamic layout map for best compatibility.
    return cast<bufferization::BufferLikeType>(
        bufferization::getMemRefTypeWithFullyDynamicLayout(tensorType,
                                                           memorySpace));
  };

  return options;
}

LogicalResult
eliminateEmptyTensors(RewriterBase &rewriter, Operation *op,
                      const OneShotBufferizationOptions &options) {
  // Analyze IR.
  OneShotAnalysisState state(op, options);
  if (failed(analyzeOp(op, state))) {
    return failure();
  }

  // Rewrite tensor.empty ops that are anchored on specific ops.
  if (failed(bufferization::eliminateEmptyTensors(rewriter, op, state))) {
    return failure();
  }

  return success();
}

void EliminateEmptyTensorsPass::runOnOperation() {
  mlir::FunctionOpInterface funcOp = getOperation();
  MLIRContext *context = &getContext();

  OpBuilder b(context);
  SmallVector<tensor::EmptyOp> emptyOps;
  funcOp.walk([&](tensor::EmptyOp emptyOp) { emptyOps.push_back(emptyOp); });
  if (llvm::any_of(emptyOps, [&](tensor::EmptyOp emptyOp) {
        return failed(duplicateTensorEmptyOps(b, emptyOp));
      })) {
    return signalPassFailure();
  }

  // Run the convert to destination style patterns.
  {
    RewritePatternSet patterns(context);
    linalg::populateConvertToDestinationStylePatterns(patterns);
    if (failed(applyPatternsGreedily(funcOp, std::move(patterns)))) {
      funcOp->emitOpError("Failed in conversion to destination style patterns");
      return signalPassFailure();
    }
  }

  // iree_codegen.store_to_buffer may consumer reshape operands that got stuck
  // at the end of the funcOp, which will prevent the SubsetInsertionOpInterface
  // implementation from eliminating empty tensors. Push up these reshape ops
  // as far as possible to prevent the reshapes from blocking empty tensor
  // elimination.
  IRRewriter rewriter(funcOp->getContext());
  moveUpMemrefReshapeOps(rewriter, funcOp);
  auto bufferizationOptions = getBufferizationOptions();
  OneShotAnalysisState state(funcOp, bufferizationOptions);
  // Analyze IR.
  if (failed(analyzeOp(funcOp, state))) {
    return signalPassFailure();
  }
  // Eliminate empty tensors.
  if (failed(bufferization::eliminateEmptyTensors(rewriter, funcOp, state))) {
    return signalPassFailure();
  }
}

// The following is copied from bufferization::runOneShotBufferize with
// modifications.
LogicalResult
runIREEOneShotBufferize(Operation *op,
                        const IREEOneShotBufferizationOptions &options,
                        bufferization::BufferizationState &state) {
  OneShotAnalysisState analyzeState(op, options);
  if (failed(analyzeOp(op, analyzeState))) {
    return failure();
  }
  if (options.testAnalysisOnly) {
    return success();
  }
  if (failed(bufferization::runOneShotBufferize(op, options, state))) {
    return failure();
  }

  // All to_buffer ops on constants can be safely marked as read only since
  // constants are immutable and any write conflicts would have already been
  // resolved by the bufferization analysis (which inserts copies as needed).
  op->walk([](bufferization::ToBufferOp toBuffer) {
    if (toBuffer.getTensor().getDefiningOp<arith::ConstantOp>()) {
      toBuffer.setReadOnly(true);
    }
  });

  return success();
}

/// Run comprehensive bufferize.
void IREEComprehensiveBufferizePass::runOnOperation() {
  mlir::FunctionOpInterface funcOp = getOperation();

  // First drop all fusion barriers. Fusion is done at this point and these can
  // be safely removed.
  IRRewriter rewriter(funcOp.getContext());
  funcOp.walk([&](IREE::Codegen::FusionBarrierOp barrier) {
    rewriter.replaceOp(barrier, barrier.getSource());
  });

  IREEOneShotBufferizationOptions options = getBufferizationOptions();
  options.testAnalysisOnly = testAnalysisOnly;
  options.printConflicts = printConflicts;
  options.allocationFn = allocationFn;
  options.memCpyFn = memCpyFn;
  // Turning off checkParallelRegions assumes that we are not relying too much
  // on bufferization being conservative. If we are, then this could cause race
  // conditions. Turning this option off could be a good step in diagnosing
  // data races on GPU.
  options.checkParallelRegions = false;

  // Place constants in the GPU constant address space to enable
  // backend-specific optimizations (e.g., scalar reads on AMDGPU).
  if (isGPUBackend(IREE::HAL::ExecutableTargetAttr::lookup(funcOp))) {
    options.defaultMemorySpaceFn =
        [](bufferization::TensorLikeType t) -> std::optional<Attribute> {
      return gpu::AddressSpaceAttr::get(t.getContext(),
                                        gpu::AddressSpace::Constant);
    };
  }

  bufferization::BufferizationState bufferizationState;
  if (failed(runIREEOneShotBufferize(funcOp, options, bufferizationState))) {
    return signalPassFailure();
  }

  // Remove redundant args and unused results.
  {
    RewritePatternSet patterns(&getContext());
    linalg::populateEraseUnusedOperandsAndResultsPatterns(patterns);
    if (failed(applyPatternsGreedily(funcOp, std::move(patterns)))) {
      return signalPassFailure();
    }
  }
}

void IREEBufferizeConstantsPass::runOnOperation() {
  mlir::bufferization::OneShotBufferizationOptions opt;
  opt.copyBeforeWrite = true;
  opt.opFilter.allowOperation(arith::ConstantOp::getOperationName());

  // Place constants in the GPU constant address space to enable
  // backend-specific optimizations (e.g., scalar reads on AMDGPU).
  if (isGPUBackend(IREE::HAL::ExecutableTargetAttr::lookup(getOperation()))) {
    opt.defaultMemorySpaceFn =
        [](bufferization::TensorLikeType t) -> std::optional<Attribute> {
      return gpu::AddressSpaceAttr::get(t.getContext(),
                                        gpu::AddressSpace::Constant);
    };
  }

  bufferization::BufferizationState bufferizationState;
  if (failed(mlir::bufferization::runOneShotBufferize(
          getOperation(), opt, bufferizationState,
          /*statistics=*/nullptr))) {
    signalPassFailure();
    return;
  }

  // All to_buffer ops on constants can be safely marked as read only since
  // constants are immutable and any write conflicts would have already been
  // resolved by the bufferization analysis (which inserts copies as needed).
  getOperation()->walk([](bufferization::ToBufferOp toBuffer) {
    if (toBuffer.getTensor().getDefiningOp<arith::ConstantOp>()) {
      toBuffer.setReadOnly(true);
    }
  });
}

std::unique_ptr<InterfacePass<mlir::FunctionOpInterface>>
createIREEComprehensiveBufferizePass(
    std::optional<BufferizationOptions::AllocationFn> allocationFn,
    std::optional<BufferizationOptions::MemCpyFn> memCpyFn) {
  if (!allocationFn) {
    allocationFn = defaultAllocationFn;
  }
  if (!memCpyFn) {
    memCpyFn = defaultMemCpyFn;
  }
  return std::make_unique<IREEComprehensiveBufferizePass>(allocationFn.value(),
                                                          memCpyFn.value());
}

void addIREEPostBufferizationPasses(OpPassManager &funcPassManager) {
  funcPassManager.addPass(createIREEInjectAssumeAlignmentPass());
  funcPassManager.addPass(memref::createResolveShapedTypeResultDimsPass());
  funcPassManager.addPass(createIREECodegenCanonicalizerPass());
  funcPassManager.addPass(createCSEPass());
  // There are redundant memcpy (with linalg.generic form) ops created, which
  // can be deleted by canonicalizer. We have to run it again because the
  // memrefs are unified in CSE pass, so we can truly remove redundant memcpy.
  funcPassManager.addPass(createIREECodegenCanonicalizerPass());
  funcPassManager.addPass(createCleanupBufferAllocViewPass());
}

void addIREEComprehensiveBufferizePasses(
    OpPassManager &funcPassManager,
    std::optional<BufferizationOptions::AllocationFn> allocationFn,
    std::optional<BufferizationOptions::MemCpyFn> memCpyFn) {
  funcPassManager.addPass(createEliminateEmptyTensorsPass());
  funcPassManager.addPass(bufferization::createEmptyTensorToAllocTensorPass());
  funcPassManager.addPass(
      createIREEComprehensiveBufferizePass(allocationFn, memCpyFn));
  addIREEPostBufferizationPasses(funcPassManager);
}

} // namespace mlir::iree_compiler
