//===- ComprehensiveBufferize.cpp - Single pass bufferization -------------===//
//
// 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 "PassDetail.h"

#include "mlir/Dialect/Arithmetic/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.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/Func/IR/FuncOps.h"
#include "mlir/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.h"
#include "mlir/Dialect/Linalg/Passes.h"
#include "mlir/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/SCF/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "mlir/Transforms/Passes.h"

using namespace mlir;
using namespace mlir::bufferization;
using namespace mlir::linalg;
using namespace mlir::linalg::comprehensive_bufferize;

namespace {
struct LinalgComprehensiveModuleBufferize
    : public LinalgComprehensiveModuleBufferizeBase<
          LinalgComprehensiveModuleBufferize> {
  LinalgComprehensiveModuleBufferize() = default;

  LinalgComprehensiveModuleBufferize(
      const LinalgComprehensiveModuleBufferize &p) = default;

  explicit LinalgComprehensiveModuleBufferize(
      const OneShotBufferizationOptions &options)
      : options(options) {}

  void runOnOperation() override;

  void getDependentDialects(DialectRegistry &registry) const override {
    registry
        .insert<bufferization::BufferizationDialect, linalg::LinalgDialect,
                memref::MemRefDialect, tensor::TensorDialect,
                vector::VectorDialect, scf::SCFDialect,
                arith::ArithmeticDialect, func::FuncDialect, AffineDialect>();
    arith::registerBufferizableOpInterfaceExternalModels(registry);
    bufferization::registerAllocationOpInterfaceExternalModels(registry);
    linalg::registerBufferizableOpInterfaceExternalModels(registry);
    scf::registerBufferizableOpInterfaceExternalModels(registry);
    std_ext::registerModuleBufferizationExternalModels(registry);
    tensor::registerBufferizableOpInterfaceExternalModels(registry);
    vector::registerBufferizableOpInterfaceExternalModels(registry);
  }

private:
  llvm::Optional<OneShotBufferizationOptions> options;
};
} // namespace

static void applyEnablingTransformations(ModuleOp moduleOp) {
  RewritePatternSet patterns(moduleOp.getContext());
  patterns.add<GeneralizePadOpPattern>(moduleOp.getContext());
  (void)applyPatternsAndFoldGreedily(moduleOp, std::move(patterns));
}

static FailureOr<Value> allocationFnUsingAlloca(OpBuilder &b, Location loc,
                                                MemRefType type,
                                                ValueRange dynShape,
                                                unsigned int bufferAlignment) {
  Value allocated = b.create<memref::AllocaOp>(
      loc, type, dynShape, b.getI64IntegerAttr(bufferAlignment));
  return allocated;
}

void LinalgComprehensiveModuleBufferize::runOnOperation() {
  OneShotBufferizationOptions opt;
  if (!options) {
    // Make new bufferization options if none were provided when creating the
    // pass.
    if (useAlloca) {
      opt.allocationFn = allocationFnUsingAlloca;
      opt.deallocationFn = [](OpBuilder &b, Location loc, Value v) {
        return success();
      };
    }
    opt.allowReturnAllocs = allowReturnAllocs;
    opt.allowUnknownOps = allowUnknownOps;
    opt.analysisFuzzerSeed = analysisFuzzerSeed;
    opt.createDeallocs = createDeallocs;
    opt.fullyDynamicLayoutMaps = fullyDynamicLayoutMaps;
    opt.printConflicts = printConflicts;
    opt.testAnalysisOnly = testAnalysisOnly;
    opt.alwaysAliasingWithDest = alwaysAliasingWithDest;
    if (initTensorElimination) {
      opt.addPostAnalysisStep(insertSliceAnchoredInitTensorEliminationStep);
    }
  } else {
    opt = *options;
  }

  ModuleOp moduleOp = getOperation();
  applyEnablingTransformations(moduleOp);

  if (failed(runModuleBufferize(moduleOp, opt))) {
    signalPassFailure();
    return;
  }

  if (opt.testAnalysisOnly)
    return;

  OpPassManager cleanupPipeline("builtin.module");
  cleanupPipeline.addPass(createCanonicalizerPass());
  cleanupPipeline.addPass(createCSEPass());
  cleanupPipeline.addPass(createLoopInvariantCodeMotionPass());
  (void)runPipeline(cleanupPipeline, moduleOp);
}

std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass() {
  return std::make_unique<LinalgComprehensiveModuleBufferize>();
}

std::unique_ptr<Pass> mlir::createLinalgComprehensiveModuleBufferizePass(
    const OneShotBufferizationOptions &options) {
  return std::make_unique<LinalgComprehensiveModuleBufferize>(options);
}