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

//===------------------- DumpDispatchGraph.cpp ----------------------------===//
//
// Generate a graphviz graph for dispatches
//
//===----------------------------------------------------------------------===//

#include <utility>

#include "iree/compiler/Dialect/Flow/IR/FlowOps.h"
#include "iree/compiler/Dialect/Flow/Transforms/Passes.h"
#include "iree/compiler/Dialect/Stream/IR/StreamInterfaces.h"
#include "iree/compiler/Dialect/TensorExt/IR/TensorExtOps.h"
#include "iree/compiler/Dialect/Util/IR/UtilDialect.h"
#include "iree/compiler/Dialect/Util/IR/UtilOps.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/Format.h"
#include "llvm/Support/GraphWriter.h"
#include "llvm/Support/ToolOutputFile.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Linalg/IR/Linalg.h"
#include "mlir/IR/AsmState.h"
#include "mlir/IR/Block.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Operation.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Support/FileUtilities.h"
#include "mlir/Support/IndentedOstream.h"

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

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

namespace {

static const StringRef kLineStyleControlFlow = "dashed";
static const StringRef kLineStyleDataFlow = "solid";
static const StringRef kShapeNode = "box";
static const StringRef kShapeBox = "box";
static const StringRef kShapeNone = "plain";
static const StringRef kShapeEllipse = "ellipse";

static StringRef getShape(Operation *op) {
  if (isa<DispatchOp>(op)) {
    return kShapeBox;
  }

  return kShapeEllipse;
}

/// Return the size limits for eliding large attributes.
static int64_t getLargeAttributeSizeLimit() {
  // Use the default from the printer flags if possible.
  if (std::optional<int64_t> limit =
          OpPrintingFlags().getLargeElementsAttrLimit()) {
    return *limit;
  }
  return 16;
}

/// Return all values printed onto a stream as a string.
static std::string strFromOs(function_ref<void(raw_ostream &)> func) {
  std::string buf;
  llvm::raw_string_ostream os(buf);
  func(os);
  return os.str();
}

/// Escape special characters such as '\n' and quotation marks.
static std::string escapeString(std::string str) {
  return strFromOs([&](raw_ostream &os) {
    for (unsigned char c : str) {
      switch (c) {
      case '\\':
        os << '\\' << '\\';
        break;
      case '\t':
        os << '\\' << 't';
        break;
      case '\n':
        os << '\\' << 'n';
        break;
      case '"':
        os << '\\' << '"';
        break;
      case '\r': // translate "carriage return" as "\l"
        os << '\\' << 'l';
        break;
      default:
        if (llvm::isPrint(c)) {
          os << c;
          break;
        }

        // Always use a full 3-character octal escape.
        os << '\\';
        os << char('0' + ((c >> 6) & 7));
        os << char('0' + ((c >> 3) & 7));
        os << char('0' + ((c >> 0) & 7));
      }
    }
  });
}

/// Put quotation marks around a given string.
static std::string quoteString(const std::string &str) {
  return "\"" + str + "\"";
}

using AttributeMap = llvm::StringMap<std::string>;

/// This struct represents a node in the DOT language. Each node has an
/// identifier and an optional identifier for the cluster (subgraph) that
/// contains the node.
/// Note: In the DOT language, edges can be drawn only from nodes to nodes, but
/// not between clusters. However, edges can be clipped to the boundary of a
/// cluster with `lhead` and `ltail` attributes. Therefore, when creating a new
/// cluster, an invisible "anchor" node is created.
struct Node {
public:
  Node(int id = 0, std::optional<int> clusterId = std::nullopt)
      : id(id), clusterId(clusterId) {}

  int id;
  std::optional<int> clusterId;
};

class GraphPrinter {
public:
  GraphPrinter(raw_ostream &os, unsigned maxLabelLen, bool printAttrs,
               bool printControlFlowEdges, bool printDataFlowEdges,
               bool printResultTypes, bool emitDispatchBody,
               bool emitInitializers)
      : os(os), maxLabelLen(maxLabelLen), printAttrs(printAttrs),
        printControlFlowEdges(printControlFlowEdges),
        printDataFlowEdges(printDataFlowEdges),
        printResultTypes(printResultTypes) {}

  void emitFunctions(ModuleOp module) {
    auto funcOps = module.getOps<mlir::FunctionOpInterface>();
    if (funcOps.empty()) {
      return;
    }

    emitGraph([&]() {
      for (auto funcOp : funcOps) {
        if (!emitInitializers && isa<IREE::Util::InitializerOp>(funcOp)) {
          continue;
        }
        processOperation(funcOp);
      }
      emitAllEdgeStmts();
    });
  }

private:
  /// Create a CFG graph for a region. Used in `Region::viewGraph`.
  void emitRegionCFG(Region &region) {
    printControlFlowEdges = true;
    printDataFlowEdges = false;
    emitGraph([&]() { processRegion(region); });
  }

  /// Emit all edges. This function should be called after all nodes have been
  /// emitted.
  void emitAllEdgeStmts() {
    for (const std::string &edge : edges) {
      os << edge << ";\n";
    }
    edges.clear();
  }

  /// Emit a cluster (subgraph). The specified builder generates the body of the
  /// cluster. Return the anchor node of the cluster.
  Node emitClusterStmt(function_ref<void()> builder, std::string label = "") {
    int clusterId = ++counter;
    os << "subgraph cluster_" << clusterId << " {\n";
    os.indent();
    // Emit invisible anchor node from/to which arrows can be drawn.
    Node anchorNode = emitNodeStmt(" ", kShapeNone);
    os << attrStmt("label", quoteString(escapeString(std::move(label))))
       << ";\n";
    builder();
    os.unindent();
    os << "}\n";
    return Node(anchorNode.id, clusterId);
  }

  /// Generate an attribute statement.
  std::string attrStmt(const Twine &key, const Twine &value) {
    return (key + " = " + value).str();
  }

  /// Emit an attribute list.
  void emitAttrList(raw_ostream &os, const AttributeMap &map) {
    os << "[";
    interleaveComma(map, os, [&](const auto &it) {
      os << this->attrStmt(it.getKey(), it.getValue());
    });
    os << "]";
  }

  // Print an MLIR attribute to `os`. Large attributes are truncated.
  void emitMlirAttr(raw_ostream &os, Attribute attr) {
    // A value used to elide large container attribute.
    int64_t largeAttrLimit = getLargeAttributeSizeLimit();

    // Always emit splat attributes.
    if (isa<SplatElementsAttr>(attr)) {
      attr.print(os);
      return;
    }

    // Elide "big" elements attributes.
    auto elements = dyn_cast<ElementsAttr>(attr);
    if (elements && elements.getNumElements() > largeAttrLimit) {
      auto type = cast<ShapedType>(elements.getType());
      os << std::string(type.getRank(), '[') << "..."
         << std::string(type.getRank(), ']') << " : " << type;
      return;
    }

    auto array = dyn_cast<ArrayAttr>(attr);
    if (array && static_cast<int64_t>(array.size()) > largeAttrLimit) {
      os << "[...]";
      return;
    }

    // Print all other attributes.
    std::string buf;
    llvm::raw_string_ostream ss(buf);
    attr.print(ss);
    os << truncateString(ss.str());
  }

  /// Append an edge to the list of edges.
  /// Note: Edges are written to the output stream via `emitAllEdgeStmts`.
  void emitEdgeStmt(Node n1, Node n2, std::string label, StringRef style) {
    AttributeMap attrs;
    attrs["style"] = style.str();
    // Do not label edges that start/end at a cluster boundary. Such edges are
    // clipped at the boundary, but labels are not. This can lead to labels
    // floating around without any edge next to them.
    if (!n1.clusterId && !n2.clusterId) {
      attrs["label"] = quoteString(escapeString(std::move(label)));
    }
    // Use `ltail` and `lhead` to draw edges between clusters.
    if (n1.clusterId) {
      attrs["ltail"] = "cluster_" + std::to_string(*n1.clusterId);
    }
    if (n2.clusterId) {
      attrs["lhead"] = "cluster_" + std::to_string(*n2.clusterId);
    }

    edges.push_back(strFromOs([&](raw_ostream &os) {
      os << llvm::format("v%i -> v%i ", n1.id, n2.id);
      emitAttrList(os, attrs);
    }));
  }

  /// Emit a graph. The specified builder generates the body of the graph.
  void emitGraph(function_ref<void()> builder) {
    os << "digraph G {\n";
    os.indent();
    // Edges between clusters are allowed only in compound mode.
    os << attrStmt("compound", "true") << ";\n";
    builder();
    os.unindent();
    os << "}\n";
  }

  /// Emit a node statement.
  Node emitNodeStmt(Operation *op) {
    int nodeId = ++counter;
    AttributeMap attrs;
    auto label = getLabel(op);
    auto shape = getShape(op);
    attrs["label"] = quoteString(escapeString(std::move(label)));
    attrs["shape"] = shape.str();
    os << llvm::format("v%i ", nodeId);
    emitAttrList(os, attrs);
    os << ";\n";
    return Node(nodeId);
  }
  Node emitNodeStmt(std::string label, StringRef shape = kShapeNode) {
    int nodeId = ++counter;
    AttributeMap attrs;
    attrs["label"] = quoteString(escapeString(std::move(label)));
    attrs["shape"] = shape.str();
    os << llvm::format("v%i ", nodeId);
    emitAttrList(os, attrs);
    os << ";\n";
    return Node(nodeId);
  }

  void printResults(raw_ostream &os, Operation *op, AsmState &state) {
    for (auto result : op->getResults()) {
      result.printAsOperand(os, state);
    }
  }

  void printResultsAndName(raw_ostream &os, Operation *op, AsmState &state) {
    printResults(os, op, state);
    os << " = " << op->getName();
  }

  void printDispatchTensorLoad(raw_ostream &os,
                               IREE::TensorExt::DispatchTensorLoadOp op,
                               AsmState &state) {
    printResultsAndName(os, op.getOperation(), state);
    os << " ";
    op.getSource().printAsOperand(os, state);
    os << " -> " << op.getResult().getType();
    os << "\r";
  }

  void printDispatchTensorStore(raw_ostream &os,
                                IREE::TensorExt::DispatchTensorStoreOp op,
                                AsmState &state) {
    os << op->getName() << " ";
    op.getValue().printAsOperand(os, state);
    os << ", ";
    op.getTarget().printAsOperand(os, state);
    os << "\r";
  }

  void printGeneric(raw_ostream &os, linalg::GenericOp op, AsmState &state) {
    printLinalgInsOuts(os, op, state);
    for (Operation &operation : *op.getBlock()) {
      os.indent(8);
      annotateOperation(os, &operation, state);
    }
  }

  template <typename T>
  void printLinalgInsOuts(raw_ostream &os, T op, AsmState &state) {
    printResultsAndName(os, op.getOperation(), state);
    os << "[";
    llvm::interleaveComma(op.getIteratorTypesArray(), os);
    os << "] (";
    printOperands(os, op.getInputs(), state);
    os << ") -> (";
    printOperands(os, op.getOutputs(), state);
    os << ")\r";
  }

  void annotateOperation(raw_ostream &os, Operation *op, AsmState &state) {
    if (isa<arith::ConstantOp>(op)) {
      return;
    }

    if (op->hasTrait<OpTrait::ReturnLike>() &&
        isa<mlir::FunctionOpInterface>(op->getParentOp())) {
      return;
    }

    if (auto load = dyn_cast<IREE::TensorExt::DispatchTensorLoadOp>(op)) {
      printDispatchTensorLoad(os, load, state);
      return;
    }

    if (auto store = dyn_cast<IREE::TensorExt::DispatchTensorStoreOp>(op)) {
      printDispatchTensorStore(os, store, state);
      return;
    }

    if (auto generic = dyn_cast<linalg::GenericOp>(op)) {
      printGeneric(os, generic, state);
      return;
    }

    if (auto linalgOp = dyn_cast<linalg::MatmulOp>(op)) {
      printLinalgInsOuts(os, linalgOp, state);
      return;
    }

    if (auto linalgOp = dyn_cast<linalg::BatchMatmulOp>(op)) {
      printLinalgInsOuts(os, linalgOp, state);
      return;
    }

    os << *op << "\r";
  }

  void printDispatchBody(raw_ostream &os, DispatchOp &dispatchOp) {
    // Find the entry point function from the dispatch entry point symbol
    // attribute.
    auto entryPoint = *dispatchOp.getEntryPointRefs().begin();
    auto executableOp = cast<ExecutableOp>(SymbolTable::lookupNearestSymbolFrom(
        dispatchOp, entryPoint.getRootReference()));
    if (!executableOp) {
      return;
    }

    auto calleeNameAttr = entryPoint.getLeafReference();
    auto innerModule = executableOp.getInnerModule();
    if (!innerModule) {
      return;
    }
    auto funcOps = innerModule.getOps<mlir::FunctionOpInterface>();
    auto funcIt = llvm::find_if(
        funcOps, [&](auto op) { return op.getNameAttr() == calleeNameAttr; });
    if (funcIt == funcOps.end()) {
      return;
    }

    auto callee = *funcIt;

    AsmState state(callee);

    // Iterate the operations of the function body and print important
    // operation.
    for (auto &block : callee.getBlocks()) {
      for (auto &op : block.getOperations()) {
        annotateOperation(os, &op, state);
      }
    }
  }

  void printOperands(raw_ostream &os, ::mlir::Operation::operand_range operands,
                     AsmState &state) {
    auto numOperands = operands.size();

    for (auto it : llvm::enumerate(operands)) {
      auto operand = it.value();
      auto op = operand.getDefiningOp();

      if (op && isScalarConstantOp(op)) {
        auto ty = operand.getType();
        if (isa<IntegerType>(ty)) {
          os << cast<arith::ConstantIntOp>(op).value();
        } else if (isa<FloatType>(ty)) {
          cast<arith::ConstantFloatOp>(op).value().print(os);
        } else {
          os << cast<arith::ConstantIndexOp>(op).value();
        }
      } else {
        operand.printAsOperand(os, state);
      }

      if (it.index() != numOperands - 1) {
        os << ", ";
      }
    }
  }

  /// Generate a label for an operation.
  std::string getLabel(Operation *op) {
    return strFromOs([&](raw_ostream &os) {
      if (op->getNumRegions() == 0) {
        auto funcOp = op->getParentOfType<mlir::FunctionOpInterface>();
        AsmState state(funcOp);
        printResults(os, op, state);
        os << " = " << op->getName();

        if (auto dispatch = dyn_cast<DispatchOp>(op)) {
          // print workload
          os << "[";
          printOperands(os, dispatch.getWorkload(), state);
          os << "]\n";

          // Print entry function name, if there is only one entry function,
          // then the name space and the entry function names are the same,
          // and we can just print the function name to save space.
          auto entryPoint = *dispatch.getEntryPointRefs().begin();
          auto rootName = entryPoint.getRootReference();
          auto leafName = entryPoint.getLeafReference();
          if (rootName == leafName) {
            os << leafName;
          } else {
            os << entryPoint; // print the full name
          }

          // print entry function args
          os << "(";
          printOperands(os, dispatch.getArguments(), state);
          os << ")\n";

          if (emitDispatchBody) {
            printDispatchBody(os, dispatch);
          }

        } else {
          os << "\n";
        }
      } else {
        os << op->getName() << "\n";
      }
      if (auto affinityOp = dyn_cast<IREE::Stream::AffinityOpInterface>(op)) {
        os << affinityOp.getAffinityAttr() << "\n";
      } else if (auto transferOp = dyn_cast<IREE::Flow::TensorTransferOp>(op)) {
        os << transferOp.getTarget() << "\n";
      } else if (auto barrierOp = dyn_cast<IREE::Flow::TensorBarrierOp>(op)) {
        os << barrierOp.getTarget() << "\n";
      }

      if (printResultTypes) {
        std::string buf;
        llvm::raw_string_ostream ss(buf);
        interleave(op->getResultTypes(), ss, "\n");
        os << ss.str();
      }
    });
  }

  /// Generate a label for a block argument.
  std::string getLabel(BlockArgument arg) {
    return "arg" + std::to_string(arg.getArgNumber());
  }

  /// Process a block. Emit a cluster and one node per block argument and
  /// operation inside the cluster.
  void processBlock(Block &block) {
    emitClusterStmt([&]() {
      for (BlockArgument &blockArg : block.getArguments()) {
        valueToNode[blockArg] = emitNodeStmt(getLabel(blockArg));
      }

      // Emit a node for each operation.
      std::optional<Node> prevNode;
      for (Operation &op : block) {
        Node nextNode = processOperation(&op);
        if (printControlFlowEdges && prevNode) {
          emitEdgeStmt(*prevNode, nextNode, /*label=*/"",
                       kLineStyleControlFlow);
        }
        prevNode = nextNode;
      }
    });
  }

  bool isScalarConstantOp(Operation *op) {
    if (auto constOp = dyn_cast<mlir::arith::ConstantOp>(op)) {
      if (constOp.getResult().getType().isIntOrIndexOrFloat()) {
        return true;
      }
    }

    return false;
  }

  /// Process an operation. If the operation has regions, emit a cluster.
  /// Otherwise, emit a node.
  Node processOperation(Operation *op) {
    Node node;

    // Do not handle some noisy Operations. Typically, the arith ops are used to
    // compute offsets, sizes, etc.
    if (op->getDialect() ==
        op->getContext()->getLoadedDialect<arith::ArithDialect>()) {
      return node;
    }
    if (isa<IREE::Util::AssumeIntOp>(op)) {
      return node;
    }
    if (op->getNumResults() &&
        llvm::all_of(op->getResultTypes(), llvm::IsaPred<IndexType>)) {
      return node;
    }

    if (op->getNumRegions() == 1) {
      // do not generate a cluster when there is one region.
      processRegion(op->getRegion(0));
    } else if (op->getNumRegions() > 1) {
      // Emit cluster for op with regions.
      node = emitClusterStmt(
          [&]() {
            for (Region &region : op->getRegions()) {
              processRegion(region);
            }
          },
          getLabel(op));
    } else {
      node = emitNodeStmt(op);
    }

    // Insert data flow edges originating from each operand.
    if (printDataFlowEdges) {
      unsigned numOperands = op->getNumOperands();
      for (unsigned i = 0; i < numOperands; i++) {
        auto operand = op->getOperand(i);

        // a constant operand is not going to be available in the map.
        if (valueToNode.count(operand)) {
          emitEdgeStmt(valueToNode[op->getOperand(i)], node,
                       /*label=*/numOperands == 1 ? "" : std::to_string(i),
                       kLineStyleDataFlow);
        }
      }
    }

    for (Value result : op->getResults()) {
      valueToNode[result] = node;
    }

    return node;
  }

  /// Process a region.
  void processRegion(Region &region) {
    for (Block &block : region.getBlocks()) {
      processBlock(block);
    }
  }

  /// Truncate long strings.
  std::string truncateString(std::string str) {
    if (str.length() <= maxLabelLen) {
      return str;
    }
    return str.substr(0, maxLabelLen) + "...";
  }

  /// Output stream to write DOT file to.
  raw_indented_ostream os;
  /// A list of edges. For simplicity, should be emitted after all nodes were
  /// emitted.
  std::vector<std::string> edges;
  /// Mapping of SSA values to Graphviz nodes/clusters.
  DenseMap<Value, Node> valueToNode;
  /// Counter for generating unique node/subgraph identifiers.
  int counter = 0;

  /// Pass options.
  unsigned maxLabelLen = 20;
  bool printAttrs = true;
  bool printControlFlowEdges = false;
  bool printDataFlowEdges = true;
  bool printResultTypes = true;
  bool emitDispatchBody = false;
  bool emitInitializers = false;
};

/// This pass generates a Graphviz dataflow visualization of an MLIR operation.
/// Note: See https://www.graphviz.org/doc/info/lang.html for more information
/// about the Graphviz DOT language.
class DumpDispatchGraphPass
    : public IREE::Flow::impl::DumpDispatchGraphPassBase<
          DumpDispatchGraphPass> {
public:
  using IREE::Flow::impl::DumpDispatchGraphPassBase<
      DumpDispatchGraphPass>::DumpDispatchGraphPassBase;

  void runOnOperation() override {
    auto modOp = dyn_cast<ModuleOp>(getOperation());
    if (!modOp) {
      return;
    }

    // Open the output file we'll be streaming to.
    // Since we are processing the entire module at once we overwrite the file.
    std::string errorMessage;
    auto file = openOutputFile(outputFile, &errorMessage);
    if (!file) {
      llvm::errs() << errorMessage << "\n";
      return signalPassFailure();
    }

    GraphPrinter printer(file->os(), maxLabelLen, printAttrs,
                         printControlFlowEdges, printDataFlowEdges,
                         printResultTypes, emitDispatchBody, emitInitializers);
    printer.emitFunctions(modOp);

    file->keep();
  }

private:
};

} // namespace

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