/**
 * \file src/jit/impl/mlir/ir/lower_to_affine_pass.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2020 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or
 * implied.
 */

#include "megbrain_build_config.h"
#if MGB_JIT && MGB_JIT_MLIR

#include "./common.h"
#include "./each_mode.h"

#include "megbrain/common.h"
#include "megbrain/jit/mlir/ir/dialect.h"
#include "megbrain/jit/mlir/ir/passes.h"
#include "megbrain/jit/mlir/ir/utils.h"

#include <llvm/ADT/Sequence.h>
#include <mlir/Dialect/Affine/IR/AffineOps.h>
#include <mlir/IR/StandardTypes.h>
#include <mlir/Pass/Pass.h>
#include <mlir/Transforms/DialectConversion.h>

using namespace mgb;
using namespace jit;

namespace {

using LoopIterationFn = function_ref<Value(
        OpBuilder& rewriter, ValueRange memRefOperands, ValueRange loopIvs)>;

void lower_op_to_loops(Operation* op, ValueRange operands,
                       PatternRewriter& rewriter,
                       LoopIterationFn process_iteration) {
    auto memref_type = (*op->result_type_begin()).cast<MemRefType>();
    auto loc = op->getLoc();

    auto alloc = jit::insert_alloc_and_dealloc(memref_type, loc, rewriter);

    llvm::SmallVector<int64_t, 4> lower_bounds(memref_type.getRank(), 0);
    llvm::SmallVector<int64_t, 4> steps(memref_type.getRank(), 1);
    buildAffineLoopNest(
            rewriter, loc, lower_bounds, memref_type.getShape(), steps,
            [&](OpBuilder& nested_builder, Location loc, ValueRange ivs) {
                Value value_to_store =
                        process_iteration(nested_builder, operands, ivs);
                nested_builder.create<AffineStoreOp>(loc, value_to_store, alloc,
                                                     ivs);
            });

    // Replace this operation with the generated alloc.
    rewriter.replaceOp(op, alloc);
}

struct ElemwiseLowering : public ConversionPattern {
    ElemwiseLowering(MLIRContext* ctx)
            : ConversionPattern(mgb::dialect::Elemwise::getOperationName(), 1,
                                ctx) {}
    LogicalResult matchAndRewrite(
            Operation* op, ArrayRef<Value> operands,
            ConversionPatternRewriter& rewriter) const final {
        auto loc = op->getLoc();
        auto dst_memref_type = (*op->result_type_begin()).cast<MemRefType>();
        megdnn::TensorLayout dst_layout = mlir_type_to_layout(dst_memref_type);
        dst_layout.init_contiguous_stride();
        lower_op_to_loops(
                op, operands, rewriter,
                [dst_layout, loc, op](OpBuilder& builder,
                                      ValueRange memref_operands,
                                      ValueRange loop_ivs) {
                    auto inputs = llvm::to_vector<4>(llvm::map_range(
                            memref_operands, [&](mlir::Value val) {
                                return get_affine_load_op(builder, loc, val,
                                                          loop_ivs, dst_layout);
                            }));
                    return lower_elemwise_to_std(op, builder, loc, inputs);
                });
        return success();
    }
};

struct TypeCvtLowering : public ConversionPattern {
    TypeCvtLowering(MLIRContext* ctx)
            : ConversionPattern(mgb::dialect::TypeCvt::getOperationName(), 1,
                                ctx) {}
    LogicalResult matchAndRewrite(
            Operation* op, ArrayRef<Value> operands,
            ConversionPatternRewriter& rewriter) const final {
        auto loc = op->getLoc();
        lower_op_to_loops(
                op, operands, rewriter,
                [loc, op](OpBuilder& builder, ValueRange memref_operands,
                          ValueRange loop_ivs) {
                    mlir::Value input = get_operand<AffineLoadOp>(
                            builder, loc, memref_operands[0], loop_ivs);
                    return lower_typecvt_to_std(op, builder, loc, input);
                });
        return success();
    }
};

struct AssignOpLowering : public ConversionPattern {
    AssignOpLowering(MLIRContext* ctx)
            : ConversionPattern(dialect::AssignOp::getOperationName(), 1, ctx) {
    }

    LogicalResult matchAndRewrite(
            Operation* op, ArrayRef<Value> operands,
            ConversionPatternRewriter& rewriter) const final {
        auto loc = op->getLoc();
        auto memref_type = operands[0].getType().cast<MemRefType>();
        dialect::AssignOpAdaptor assign_adaptor(operands);

        llvm::SmallVector<int64_t, 4> lower_bounds(memref_type.getRank(), 0);
        llvm::SmallVector<int64_t, 4> steps(memref_type.getRank(), 1);
        buildAffineLoopNest(
                rewriter, loc, lower_bounds, memref_type.getShape(), steps,
                [&](OpBuilder& nested_builder, Location loc, ValueRange ivs) {
                    auto loaded_lhs = nested_builder.create<AffineLoadOp>(
                            loc, assign_adaptor.lhs(), ivs);
                    nested_builder.create<AffineStoreOp>(
                            loc, loaded_lhs, assign_adaptor.rhs(), ivs);
                });

        rewriter.eraseOp(op);
        return success();
    }
};

struct ReturnOpLowering : public OpRewritePattern<dialect::ReturnOp> {
    using OpRewritePattern<dialect::ReturnOp>::OpRewritePattern;

    LogicalResult matchAndRewrite(dialect::ReturnOp op,
                                  PatternRewriter& rewriter) const final {
        // We lower "mgb.return" directly to "std.return".
        rewriter.replaceOpWithNewOp<mlir::ReturnOp>(op);
        return success();
    }
};

struct ConstantScalarOpLowering
        : public OpRewritePattern<dialect::ConstantScalarOp> {
    using OpRewritePattern<dialect::ConstantScalarOp>::OpRewritePattern;

    LogicalResult matchAndRewrite(dialect::ConstantScalarOp op,
                                  PatternRewriter& rewriter) const final {
        dialect::ConstantScalarOpAdaptor constant_scalar_adaptor(op);
        rewriter.replaceOpWithNewOp<mlir::ConstantOp>(
                op, constant_scalar_adaptor.value());
        return success();
    }
};

class MgbToAffineLoweringPass
        : public PassWrapper<MgbToAffineLoweringPass, FunctionPass> {
public:
    void getDependentDialects(mlir::DialectRegistry& registry) const override {
        registry.insert<mlir::AffineDialect>();
        registry.insert<mlir::StandardOpsDialect>();
    }

    void runOnFunction() override final {
        ConversionTarget target(getContext());
        target.addLegalDialect<AffineDialect, StandardOpsDialect>();
        // target.addLegalDialect<AffineDialect>();
        target.addIllegalDialect<MgbDialect>();

        OwningRewritePatternList patterns;
        patterns.insert<ElemwiseLowering, TypeCvtLowering, ReturnOpLowering,
                        AssignOpLowering, ConstantScalarOpLowering>(
                &getContext());

        if (failed(applyPartialConversion(getFunction(), target,
                                          std::move(patterns)))) {
            signalPassFailure();
        }
    }
};

}  // namespace

std::unique_ptr<mlir::Pass> mgb::jit::create_lower_to_affine_pass() {
    return std::make_unique<MgbToAffineLoweringPass>();
}

namespace mgb {
namespace jit {
void register_test_mgb_to_affine_lowering_pass() {
    PassRegistration<MgbToAffineLoweringPass>(
            "mgb-convert-to-affine",
            "Perform conversion from MGB Dialect to Affine Dialect ",
            [] { return std::make_unique<MgbToAffineLoweringPass>(); });
}
}  // namespace jit
}  // namespace mgb
#endif  // MGB_JIT && MGB_JIT_MLIR

// vim: syntax=cpp.doxygen