//===- VectorizerTestPass.cpp - VectorizerTestPass Pass Impl --------------===//
//
// Copyright 2019 The MLIR Authors.
//
// 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
//
//   http://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.
// =============================================================================
//
// This file implements a simple testing pass for vectorization functionality.
//
//===----------------------------------------------------------------------===//

#include "mlir/AffineOps/AffineOps.h"
#include "mlir/Analysis/AffineAnalysis.h"
#include "mlir/Analysis/NestedMatcher.h"
#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Analysis/VectorAnalysis.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/StandardTypes.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/Functional.h"
#include "mlir/Support/STLExtras.h"
#include "mlir/Transforms/Passes.h"
#include "third_party/llvm/llvm/include/llvm/ADT/STLExtras.h"

#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"

#define DEBUG_TYPE "vectorizer-test"

using namespace mlir;

using llvm::outs;
using llvm::SetVector;

using functional::map;

static llvm::cl::OptionCategory clOptionsCategory(DEBUG_TYPE " options");

static llvm::cl::list<int> clTestVectorShapeRatio(
    "vector-shape-ratio",
    llvm::cl::desc("Specify the HW vector size for vectorization"),
    llvm::cl::ZeroOrMore, llvm::cl::cat(clOptionsCategory));
static llvm::cl::opt<bool> clTestForwardSlicingAnalysis(
    "forward-slicing",
    llvm::cl::desc("Enable testing forward static slicing and topological sort "
                   "functionalities"),
    llvm::cl::cat(clOptionsCategory));
static llvm::cl::opt<bool> clTestBackwardSlicingAnalysis(
    "backward-slicing",
    llvm::cl::desc("Enable testing backward static slicing and "
                   "topological sort functionalities"),
    llvm::cl::cat(clOptionsCategory));
static llvm::cl::opt<bool> clTestSlicingAnalysis(
    "slicing",
    llvm::cl::desc("Enable testing static slicing and topological sort "
                   "functionalities"),
    llvm::cl::cat(clOptionsCategory));
static llvm::cl::opt<bool> clTestComposeMaps(
    "compose-maps",
    llvm::cl::desc(
        "Enable testing the composition of AffineMap where each "
        "AffineMap in the composition is specified as the affine_map attribute "
        "in a constant op."),
    llvm::cl::cat(clOptionsCategory));
static llvm::cl::opt<bool> clTestNormalizeMaps(
    "normalize-maps",
    llvm::cl::desc(
        "Enable testing the normalization of AffineAffineApplyOp "
        "where each AffineAffineApplyOp in the composition is a single output "
        "instruction."),
    llvm::cl::cat(clOptionsCategory));

namespace {

struct VectorizerTestPass : public FunctionPass<VectorizerTestPass> {
  static constexpr auto kTestAffineMapOpName = "test_affine_map";
  static constexpr auto kTestAffineMapAttrName = "affine_map";

  void runOnFunction() override;
  void testVectorShapeRatio(Function *f);
  void testForwardSlicing(Function *f);
  void testBackwardSlicing(Function *f);
  void testSlicing(Function *f);
  void testComposeMaps(Function *f);
  void testNormalizeMaps(Function *f);
};

} // end anonymous namespace

void VectorizerTestPass::testVectorShapeRatio(Function *f) {
  using matcher::Op;
  SmallVector<int64_t, 8> shape(clTestVectorShapeRatio.begin(),
                                clTestVectorShapeRatio.end());
  auto subVectorType =
      VectorType::get(shape, FloatType::getF32(f->getContext()));
  // Only filter instructions that operate on a strict super-vector and have one
  // return. This makes testing easier.
  auto filter = [subVectorType](Instruction &inst) {
    assert(subVectorType.getElementType() ==
               FloatType::getF32(subVectorType.getContext()) &&
           "Only f32 supported for now");
    if (!matcher::operatesOnSuperVectors(inst, subVectorType)) {
      return false;
    }
    if (inst.getNumResults() != 1) {
      return false;
    }
    return true;
  };
  auto pat = Op(filter);
  SmallVector<NestedMatch, 8> matches;
  pat.match(f, &matches);
  for (auto m : matches) {
    auto *opInst = m.getMatchedInstruction();
    // This is a unit test that only checks and prints shape ratio.
    // As a consequence we write only Ops with a single return type for the
    // purpose of this test. If we need to test more intricate behavior in the
    // future we can always extend.
    auto superVectorType = opInst->getResult(0)->getType().cast<VectorType>();
    auto ratio = shapeRatio(superVectorType, subVectorType);
    if (!ratio.hasValue()) {
      opInst->emitNote("NOT MATCHED");
    } else {
      outs() << "\nmatched: " << *opInst << " with shape ratio: ";
      interleaveComma(MutableArrayRef<unsigned>(*ratio), outs());
    }
  }
}

static std::string toString(Instruction *inst) {
  std::string res;
  llvm::raw_string_ostream os(res);
  inst->print(os);
  return res;
}

static NestedPattern patternTestSlicingOps() {
  // Just use a custom op name for this test, it makes life easier.
  constexpr auto kTestSlicingOpName = "slicing-test-op";
  using functional::map;
  using matcher::Op;
  // Match all OpInstructions with the kTestSlicingOpName name.
  auto filter = [](Instruction &inst) {
    return inst.getName().getStringRef() == kTestSlicingOpName;
  };
  return Op(filter);
}

void VectorizerTestPass::testBackwardSlicing(Function *f) {
  SmallVector<NestedMatch, 8> matches;
  patternTestSlicingOps().match(f, &matches);
  for (auto m : matches) {
    SetVector<Instruction *> backwardSlice;
    getBackwardSlice(m.getMatchedInstruction(), &backwardSlice);
    auto strs = map(toString, backwardSlice);
    outs() << "\nmatched: " << *m.getMatchedInstruction()
           << " backward static slice: ";
    for (const auto &s : strs) {
      outs() << "\n" << s;
    }
  }
}

void VectorizerTestPass::testForwardSlicing(Function *f) {
  SmallVector<NestedMatch, 8> matches;
  patternTestSlicingOps().match(f, &matches);
  for (auto m : matches) {
    SetVector<Instruction *> forwardSlice;
    getForwardSlice(m.getMatchedInstruction(), &forwardSlice);
    auto strs = map(toString, forwardSlice);
    outs() << "\nmatched: " << *m.getMatchedInstruction()
           << " forward static slice: ";
    for (const auto &s : strs) {
      outs() << "\n" << s;
    }
  }
}

void VectorizerTestPass::testSlicing(Function *f) {
  SmallVector<NestedMatch, 8> matches;
  patternTestSlicingOps().match(f, &matches);
  for (auto m : matches) {
    SetVector<Instruction *> staticSlice = getSlice(m.getMatchedInstruction());
    auto strs = map(toString, staticSlice);
    outs() << "\nmatched: " << *m.getMatchedInstruction() << " static slice: ";
    for (const auto &s : strs) {
      outs() << "\n" << s;
    }
  }
}

static bool customOpWithAffineMapAttribute(Instruction &inst) {
  return inst.getName().getStringRef() ==
         VectorizerTestPass::kTestAffineMapOpName;
}

void VectorizerTestPass::testComposeMaps(Function *f) {
  using matcher::Op;
  auto pattern = Op(customOpWithAffineMapAttribute);
  SmallVector<NestedMatch, 8> matches;
  pattern.match(f, &matches);
  SmallVector<AffineMap, 4> maps;
  maps.reserve(matches.size());
  for (auto m : llvm::reverse(matches)) {
    auto *opInst = m.getMatchedInstruction();
    auto map = opInst->getAttr(VectorizerTestPass::kTestAffineMapAttrName)
                   .cast<AffineMapAttr>()
                   .getValue();
    maps.push_back(map);
  }
  AffineMap res;
  for (auto m : maps) {
    res = res ? res.compose(m) : m;
  }
  simplifyAffineMap(res).print(outs() << "\nComposed map: ");
}

static bool affineApplyOp(Instruction &inst) {
  return inst.isa<AffineApplyOp>();
}

static bool singleResultAffineApplyOpWithoutUses(Instruction &inst) {
  auto app = inst.dyn_cast<AffineApplyOp>();
  return app && app.use_empty();
}

void VectorizerTestPass::testNormalizeMaps(Function *f) {
  using matcher::Op;

  // Save matched AffineApplyOp that all need to be erased in the end.
  auto pattern = Op(affineApplyOp);
  SmallVector<NestedMatch, 8> toErase;
  pattern.match(f, &toErase);
  {
    // Compose maps.
    auto pattern = Op(singleResultAffineApplyOpWithoutUses);
    SmallVector<NestedMatch, 8> matches;
    pattern.match(f, &matches);
    for (auto m : matches) {
      auto app = m.getMatchedInstruction()->cast<AffineApplyOp>();
      FuncBuilder b(m.getMatchedInstruction());
      SmallVector<Value *, 8> operands(app.getOperands());
      makeComposedAffineApply(&b, app.getLoc(), app.getAffineMap(), operands);
    }
  }
  // We should now be able to erase everything in reverse order in this test.
  for (auto m : llvm::reverse(toErase)) {
    m.getMatchedInstruction()->erase();
  }
}

void VectorizerTestPass::runOnFunction() {
  // Thread-safe RAII local context, BumpPtrAllocator freed on exit.
  NestedPatternContext mlContext;

  // Only support single block functions at this point.
  Function *f = getFunction();
  if (f->getBlocks().size() != 1)
    return;

  if (!clTestVectorShapeRatio.empty()) {
    testVectorShapeRatio(f);
  }
  if (clTestForwardSlicingAnalysis) {
    testForwardSlicing(f);
  }
  if (clTestBackwardSlicingAnalysis) {
    testBackwardSlicing(f);
  }
  if (clTestSlicingAnalysis) {
    testSlicing(f);
  }
  if (clTestComposeMaps) {
    testComposeMaps(f);
  }
  if (clTestNormalizeMaps) {
    testNormalizeMaps(f);
  }
}

FunctionPassBase *mlir::createVectorizerTestPass() {
  return new VectorizerTestPass();
}

static PassRegistration<VectorizerTestPass>
    pass("vectorizer-test", "Tests vectorizer standalone functionality.");

#undef DEBUG_TYPE