// Copyright 2024 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/Dialect/Flow/Conversion/MeshToFlow/Patterns.h"
#include "iree/compiler/Dialect/Flow/Transforms/Passes.h"
#include "iree/compiler/Utils/Folding.h"
#include "iree/compiler/Utils/Indexing.h"
#include "iree/compiler/Utils/OpVisitor.h"
#include "iree/compiler/Utils/Permutation.h"
#include "iree/compiler/Utils/SmallVectorDenseMapInfo.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/Dialect/MemRef/Transforms/Transforms.h"
#include "mlir/Dialect/Mesh/IR/MeshDialect.h"
#include "mlir/Dialect/Mesh/IR/MeshOps.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Tensor/IR/Tensor.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"

namespace mlir::iree_compiler::IREE::Flow {

#define GEN_PASS_DEF_CONVERTMESHTOFLOWPASS
#include "iree/compiler/Dialect/Flow/Transforms/Passes.h.inc"

namespace {

static bool hasMoreThanOneMesh(Operation *op) {
  int meshCount = 0;
  op->walk([&meshCount](mesh::MeshOp mesh) {
    ++meshCount;
    return meshCount > 1 ? WalkResult::interrupt() : WalkResult::advance();
  });
  return meshCount > 1;
}

static SmallVector<mesh::MeshAxis> getAllMeshAxes(mesh::MeshOp mesh) {
  SmallVector<mesh::MeshAxis> res(mesh.getRank());
  std::iota(res.begin(), res.end(), 0);
  return res;
}

using MeshAndAxesSet =
    DenseSet<std::tuple<mesh::MeshOp, SmallVector<mesh::MeshAxis>>>;

template <typename Op>
struct CollectiveOpVisitor {
  CollectiveOpVisitor(MeshAndAxesSet &meshAndAxesSet,
                      SymbolTableCollection &symbolTableCollection)
      : meshAndAxesSet(meshAndAxesSet),
        symbolTableCollection(symbolTableCollection) {}
  void operator()(Op op) {
    meshAndAxesSet.insert(std::make_tuple(
        symbolTableCollection.lookupNearestSymbolFrom<mesh::MeshOp>(
            op, op.getMeshAttr()),
        llvm::to_vector(op.getMeshAxes())));
  }

private:
  MeshAndAxesSet &meshAndAxesSet;
  SymbolTableCollection &symbolTableCollection;
};

template <typename Op>
struct CollectiveOpWithoutMeshAxesVisitor {
  CollectiveOpWithoutMeshAxesVisitor(
      MeshAndAxesSet &meshAndAxesSet,
      SymbolTableCollection &symbolTableCollection)
      : meshAndAxesSet(meshAndAxesSet),
        symbolTableCollection(symbolTableCollection) {}
  void operator()(Op op) {
    mesh::MeshOp mesh =
        symbolTableCollection.lookupNearestSymbolFrom<mesh::MeshOp>(
            op, op.getMeshAttr());
    meshAndAxesSet.insert(std::make_tuple(mesh, getAllMeshAxes(mesh)));
  }

private:
  MeshAndAxesSet &meshAndAxesSet;
  SymbolTableCollection &symbolTableCollection;
};

void populateMeshAndAxes(Operation *op, MeshAndAxesSet &meshAndAxesSet,
                         SymbolTableCollection &symbolTableCollection) {
  OpVisitorCollection opVisitors;
  opVisitors.emplaceVisitors<
      CollectiveOpVisitor<mesh::AllGatherOp>,
      CollectiveOpVisitor<mesh::AllReduceOp>,
      CollectiveOpVisitor<mesh::AllToAllOp>,
      CollectiveOpVisitor<mesh::ReduceScatterOp>,
      CollectiveOpWithoutMeshAxesVisitor<mesh::ProcessLinearIndexOp>>(
      meshAndAxesSet, symbolTableCollection);

  op->walk([&opVisitors](Operation *op) {
    opVisitors(op);
    return WalkResult::advance();
  });
}

// Derives a channel symbol name from the given mesh axes.
static SmallString<64> getMeshChannelName(mesh::MeshOp mesh,
                                          ArrayRef<mesh::MeshAxis> axes) {
  SmallString<64> res;
  llvm::raw_svector_ostream stream(res);
  stream << "_mesh_" << mesh.getSymName();
  if (axes.empty()) {
    return res;
  }
  stream << "_axes";
  for (mesh::MeshAxis axis : axes) {
    stream << "_" << axis;
  }
  return res;
}

static bool isDefaultChannel(mesh::MeshOp mesh,
                             ArrayRef<mesh::MeshAxis> meshAxes) {
  if (mesh.getRank() != static_cast<int64_t>(meshAxes.size())) {
    return false;
  }
  return isIdentityPermutation(meshAxes);
}

static Value getDefaultChannel(Location loc, mesh::MeshOp mesh,
                               bool useNamedDefaultChannels,
                               OpBuilder &builder) {
  if (useNamedDefaultChannels)
    return builder.create<IREE::Flow::ChannelDefaultOp>(loc, mesh.getSymName());
  else
    return builder.create<IREE::Flow::ChannelDefaultOp>(loc);
}

static Value buildCachedChannelLoading(Location loc, mesh::MeshOp meshOp,
                                       ArrayRef<mesh::MeshAxis> meshAxes,
                                       bool useNamedDefaultChannels,
                                       OpBuilder &builder) {
  if (isDefaultChannel(meshOp, meshAxes)) {
    return getDefaultChannel(loc, meshOp, useNamedDefaultChannels, builder);
  }
  // TODO: lookup the mesh name instead of generating it again - today this will
  // fail if there are any conflicting names during channel creation.
  return builder.create<IREE::Util::GlobalLoadOp>(
      loc, builder.getType<IREE::Flow::ChannelType>(),
      getMeshChannelName(meshOp, meshAxes));
}

// Remove from `values` elements that have indices present in filter.
static SmallVector<Value> filterOutByIndex(ArrayRef<Value> values,
                                           ArrayRef<mesh::MeshAxis> filter) {
  SmallVector<Value> res;
  for (size_t i = 0; i < values.size(); ++i) {
    if (!llvm::is_contained(filter, i)) {
      res.push_back(values[i]);
    }
  }
  return res;
}

static Value buildChannelCreation(mesh::MeshOp mesh,
                                  ArrayRef<mesh::MeshAxis> meshAxes,
                                  bool useNamedDefaultChannels,
                                  ImplicitLocOpBuilder &builder) {
  assert(mesh);
  Value meshChannel = getDefaultChannel(builder.getLoc(), mesh,
                                        useNamedDefaultChannels, builder);
  SmallVector<Value> meshProcessMultiIndex =
      builder.create<mesh::ProcessMultiIndexOp>(mesh).getResults();
  SmallVector<Value> meshShape =
      builder.create<mesh::MeshShapeOp>(mesh).getResults();
  SmallVector<Value> reorderedMeshIndex =
      permute(ArrayRef<Value>(meshProcessMultiIndex), meshAxes);
  SmallVector<Value> reorderedMeshShape =
      permute(ArrayRef<Value>(meshShape), meshAxes);
  SmallVector<Value> groupIndex =
      filterOutByIndex(meshProcessMultiIndex, meshAxes);
  SmallVector<Value> groupsShape = filterOutByIndex(meshShape, meshAxes);
  OpFoldResult reorderedProcessLinearIndex =
      linearIndexFromShape(toOpFoldResults(reorderedMeshIndex),
                           toOpFoldResults(reorderedMeshShape), builder);
  OpFoldResult color = linearIndexFromShape(
      toOpFoldResults(groupIndex), toOpFoldResults(groupsShape), builder);
  return builder.create<IREE::Flow::ChannelSplitOp>(
      meshChannel,
      getValueOrCreateConstantIndexOp(builder, builder.getLoc(), color),
      getValueOrCreateConstantIndexOp(builder, builder.getLoc(),
                                      reorderedProcessLinearIndex));
}

static void buildChannelInitializer(mesh::MeshOp mesh,
                                    ArrayRef<mesh::MeshAxis> meshAxes,
                                    bool useNamedDefaultChannels,
                                    ImplicitLocOpBuilder &builder) {
  IREE::Util::InitializerOp initOp =
      builder.create<IREE::Util::InitializerOp>();
  Block *block = builder.createBlock(&initOp.getBody());
  ImplicitLocOpBuilder::InsertionGuard insertionGuard(builder);
  builder.setInsertionPointToStart(block);
  Value channel =
      buildChannelCreation(mesh, meshAxes, useNamedDefaultChannels, builder);
  builder.create<IREE::Util::GlobalStoreOp>(channel,
                                            getMeshChannelName(mesh, meshAxes));
  builder.create<IREE::Util::ReturnOp>();
}

// Construct a Flow channel inside `module` using
// util.global and util.initializer.
static void buildGlobalChannelCreation(mesh::MeshOp mesh,
                                       ArrayRef<mesh::MeshAxis> meshAxes,
                                       bool useNamedDefaultChannels,
                                       ModuleOp module, OpBuilder &opBuilder) {
  if (isDefaultChannel(mesh, meshAxes)) {
    return;
  }

  ImplicitLocOpBuilder builder(mesh.getLoc(), opBuilder);
  ImplicitLocOpBuilder::InsertionGuard insertionGuard(builder);
  builder.setInsertionPointToStart(&module.getBodyRegion().getBlocks().front());

  auto channelName = getMeshChannelName(mesh, meshAxes);
  builder.create<IREE::Util::GlobalOp>(
      builder.getStringAttr("private"), channelName,
      builder.getType<IREE::Flow::ChannelType>(), false, TypedAttr(),
      builder.getAttr<IREE::Util::InlineNeverAttr>());
  buildChannelInitializer(mesh, meshAxes, useNamedDefaultChannels, builder);
}

static void createChannels(ModuleOp moduleOp,
                           SymbolTableCollection &symbolTableCollection,
                           MeshAndAxesSet &meshAndAxesSet,
                           bool useNamedDefaultChannels) {
  populateMeshAndAxes(moduleOp, meshAndAxesSet, symbolTableCollection);

  OpBuilder builder(moduleOp->getContext());

  // Sort for deterministic testing with FileCheck.
  auto meshAndAxesSetSorted = llvm::to_vector(meshAndAxesSet);
  llvm::sort(meshAndAxesSetSorted, [](auto &a, auto &b) {
    int nameCompareRes =
        std::get<0>(a).getSymName().compare(std::get<0>(b).getSymName());
    if (nameCompareRes == 0)
      return std::get<1>(a) < std::get<1>(b);
    return nameCompareRes < 0;
  });
  for (auto &[mesh, meshAxes] : llvm::make_range(meshAndAxesSetSorted.rbegin(),
                                                 meshAndAxesSetSorted.rend())) {
    buildGlobalChannelCreation(mesh, meshAxes, useNamedDefaultChannels,
                               moduleOp, builder);
  }
}

static LogicalResult
convertCollectives(ModuleOp moduleOp,
                   SymbolTableCollection &symbolTableCollection,
                   bool useNamedDefaultChannels) {
  RewritePatternSet patterns(moduleOp->getContext());
  IREE::Flow::populateMeshToFlowCollectivesPatterns(
      patterns, symbolTableCollection,
      [&](Location loc, mesh::MeshOp meshOp,
          std::optional<ArrayRef<mesh::MeshAxis>> meshAxes,
          OpBuilder &builder) {
        if (meshAxes.has_value()) {
          return buildCachedChannelLoading(loc, meshOp, *meshAxes,
                                           useNamedDefaultChannels, builder);
        } else {
          return buildCachedChannelLoading(loc, meshOp, getAllMeshAxes(meshOp),
                                           useNamedDefaultChannels, builder);
        }
      });
  return applyPatternsGreedily(moduleOp, std::move(patterns));
}

static void removeMeshOps(MeshAndAxesSet &meshAndAxesSet) {
  auto meshRange =
      llvm::map_range(meshAndAxesSet, [](auto &v) { return std::get<0>(v); });
  DenseSet<mesh::MeshOp> meshOpsSet(std::begin(meshRange), std::end(meshRange));
  for (mesh::MeshOp op : meshOpsSet) {
    if (op)
      op.erase();
  }
}

struct ConvertMeshToFlowPass
    : public IREE::Flow::impl::ConvertMeshToFlowPassBase<
          ConvertMeshToFlowPass> {
  void runOnOperation() override {
    SymbolTableCollection symbolTableCollection;
    MeshAndAxesSet meshAndAxesSet;
    const bool useNamedDefaultChannels = hasMoreThanOneMesh(getOperation());

    createChannels(getOperation(), symbolTableCollection, meshAndAxesSet,
                   useNamedDefaultChannels);
    if (failed(convertCollectives(getOperation(), symbolTableCollection,
                                  useNamedDefaultChannels))) {
      return signalPassFailure();
    }

    // Cleanup mesh definition ops that are no longer referenced.
    removeMeshOps(meshAndAxesSet);
  }
};

} // namespace

} // namespace mlir::iree_compiler::IREE::Flow
