[Mlir-commits] [llvm] [mlir] [mlir] Add Normalize pass (PR #162266)

[Shourya Goel]

================
@@ -0,0 +1,436 @@
+//===- Normalize.cpp - Conversion from MLIR to its canonical form ---------===//
+//
+// 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 "mlir/Conversion/Normalize/Normalize.h"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
+#include "mlir/IR/AsmState.h"
+#include "mlir/IR/TypeUtilities.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <iomanip>
+#include <sstream>
+
+namespace mlir {
+#define GEN_PASS_DEF_NORMALIZE
+#include "mlir/Conversion/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+#define DEBUG_TYPE "normalize"
+
+namespace {
+struct NormalizePass : public impl::NormalizeBase<NormalizePass> {
+  NormalizePass() = default;
+
+  void runOnOperation() override;
+
+private:
+  const uint64_t MagicHashConstant = 0x6acaa36bef8325c5ULL;
+  void
+  collectOutputOperations(Block &block,
+                          SmallVector<Operation *, 16> &Output) const noexcept;
+  bool isOutput(mlir::Operation &op) const noexcept;
+
+  void reorderOperations(const SmallVector<mlir::Operation *, 16> &Outputs);
+  void
+  reorderOperation(mlir::Operation *used, mlir::Operation *user,
+                   llvm::SmallPtrSet<const mlir::Operation *, 32> &visited);
+
+  void renameOperations(const SmallVector<Operation *, 16> &Outputs);
+  void RenameOperation(mlir::Operation *op,
+                       SmallPtrSet<const mlir::Operation *, 32> &visited);
+
+  bool isInitialOperation(mlir::Operation *const op) const noexcept;
+  void nameAsInitialOperation(
+      mlir::Operation *op,
+      llvm::SmallPtrSet<const mlir::Operation *, 32> &visited);
+  void nameAsRegularOperation(
+      mlir::Operation *op,
+      llvm::SmallPtrSet<const mlir::Operation *, 32> &visited);
+  bool hasOnlyImmediateOperands(mlir::Operation *const op) const noexcept;
+  llvm::SetVector<int> getOutputFootprint(
+      mlir::Operation *op,
+      llvm::SmallPtrSet<const mlir::Operation *, 32> &visited) const;
+  void foldOperation(mlir::Operation *op);
+  void reorderOperationOperandsByName(mlir::Operation *op);
+  mlir::OpPrintingFlags flags{};
+};
+} // namespace
+
+void NormalizePass::runOnOperation() {
+  flags.printNameLocAsPrefix(true);
+
+  ModuleOp module = getOperation();
+
+  for (auto &op : module.getOps()) {
+    SmallVector<Operation *, 16> Outputs;
+
+    for (auto &region : op.getRegions())
+      for (auto &block : region)
+        collectOutputOperations(block, Outputs);
+
+    reorderOperations(Outputs);
+    renameOperations(Outputs);
+  }
+}
+
+void NormalizePass::renameOperations(
+    const SmallVector<Operation *, 16> &Outputs) {
+  llvm::SmallPtrSet<const mlir::Operation *, 32> visited;
+
+  for (auto *op : Outputs)
+    RenameOperation(op, visited);
+}
+
+void NormalizePass::RenameOperation(
+    Operation *op, SmallPtrSet<const mlir::Operation *, 32> &visited) {
+  if (!visited.count(op)) {
+    visited.insert(op);
+
+    if (isInitialOperation(op)) {
+      nameAsInitialOperation(op, visited);
+    } else {
+      nameAsRegularOperation(op, visited);
+    }
+    foldOperation(op);
+    reorderOperationOperandsByName(op);
+  }
+}
+
+bool NormalizePass::isInitialOperation(
+    mlir::Operation *const op) const noexcept {
+  return !op->use_empty() and hasOnlyImmediateOperands(op);
+}
+
+bool NormalizePass::hasOnlyImmediateOperands(
+    mlir::Operation *const op) const noexcept {
+  for (mlir::Value operand : op->getOperands())
+    if (mlir::Operation *defOp = operand.getDefiningOp())
+      if (!(defOp->hasTrait<OpTrait::ConstantLike>()))
+        return false;
+  return true;
+}
+
+std::string inline to_string(uint64_t const hash) noexcept {
+  std::ostringstream oss;
+  oss << std::hex << std::setw(5) << std::setfill('0') << hash;
+  std::string tmp = oss.str();
+  return tmp.size() > 5 ? tmp.substr(tmp.size() - 5, 5) : tmp;
+}
+
+uint64_t inline strHash(std::string_view data) noexcept {
+  const static uint64_t FNV_OFFSET = 0xcbf29ce484222325ULL;
+  const static uint64_t FNV_PRIME = 0x100000001b3ULL;
+  uint64_t hash = FNV_OFFSET;
+  for (const auto &c : data) {
+    hash ^= static_cast<uint64_t>(c);
+    hash *= FNV_PRIME;
+  }
+  return hash;
+}
+
+std::string inline split(std::string_view str, const char &delimiter,
+                         int indx = 0) noexcept {
+  std::stringstream ss{std::string{str}};
+  std::string item;
+  int cnt = 0;
+  while (std::getline(ss, item, delimiter)) {
+    if (cnt == indx) {
+      std::replace(item.begin(), item.end(), ':', '_');
+      return item;
+    } else {
+      cnt++;
+    }
+  }
+  return nullptr;
+}
+
+void NormalizePass::nameAsInitialOperation(
+    mlir::Operation *op,
+    llvm::SmallPtrSet<const mlir::Operation *, 32> &visited) {
+
+  for (mlir::Value operand : op->getOperands())
+    if (mlir::Operation *defOp = operand.getDefiningOp())
+      RenameOperation(defOp, visited);
+
+  uint64_t Hash = MagicHashConstant;
+
+  uint64_t opcodeHash = strHash(op->getName().getStringRef().str());
----------------
Sh0g0-1758 wrote:

I don't think so. iirc, set_fixed_execution_hash_seed has been removed in the latest tree which means llvm::hash_value currently gives different hash in subsequent runs (unless you rebuild llvm with certain flags). I believe the user would want the SSA names to be consistent throughout runs. 

https://github.com/llvm/llvm-project/pull/162266