[Transforms] Add constant_tensors_folding pass

include/gc/Analysis/DataFlow/ConstantSubgraphAnalyser.h

@@ -0,0 +1,125 @@
+//===-- ConstantSubgraphAnalyser.h - Constant subgraph ----------*- C++ -*-===//
+//
+// This file is 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
+//
+//===----------------------------------------------------------------------===//
+///
+/// This file implements constant subgraph analysis. In this file are:
+/// 1. the lattice value class that represents operations with constant inputs
+/// and outputs in the program, and
+/// 2. a sparse constant subgraph analysis.
+///
+///===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_DATAFLOW_CONSTANTSUBGRAPHANALYSER_H
+#define MLIR_ANALYSIS_DATAFLOW_CONSTANTSUBGRAPHANALYSER_H
+
+#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
+
+namespace mlir {
+namespace dataflow {
+
+//===----------------------------------------------------------------------===//
+// IsConstantTensor
+//===----------------------------------------------------------------------===//
+
+/// This lattice represents a boolean indicating if a value is constant.
+class IsConstantTensor {
+public:
+  /// Construct as uninitialized.
+  explicit IsConstantTensor() = default;
+
+  /// Construct with a known state.
+  explicit IsConstantTensor(bool initialized, bool isConstantTensor)
+      : initialized(initialized), isConstantTensor(isConstantTensor) {}
+
+  /// Get the state. Must be initialized before.
+  bool getIsConstantTensor() const {
+    assert(!isUninitialized());
+    return isConstantTensor;
+  }
+
+  /// Compare.
+  bool operator==(const IsConstantTensor &rhs) const {
+    return initialized == rhs.initialized &&
+           isConstantTensor == rhs.isConstantTensor;
+  }
+
+  void print(raw_ostream &os) const;
+
+  /// Get uninitialized state. This happens when the
+  /// state hasn't been set during the analysis.
+  static IsConstantTensor getUninitialized() { return IsConstantTensor{}; }
+
+  /// Whether the state is uninitialized.
+  bool isUninitialized() const { return !initialized; }
+
+  /// Get unknown state.
+  static IsConstantTensor getUnknown() {
+    return IsConstantTensor{/*initialized=*/false,
+                            /*isConstantTensor*/ false};
+  }
+
+  // Join two states.
+  static IsConstantTensor join(const IsConstantTensor &lhs,
+                               const IsConstantTensor &rhs) {
+    // if one is uninitialized, use another
+    if (lhs.isUninitialized())
+      return rhs;
+    if (rhs.isUninitialized())
+      return lhs;
+
+    // both are initialized, intersect them
+    if (!lhs.isUninitialized() && !rhs.isUninitialized()) {
+      return IsConstantTensor(true, lhs.getIsConstantTensor() &&
+                                        rhs.getIsConstantTensor());
+    }
+    return getUninitialized();
+  }
+
+private:
+  bool initialized = false;
+  bool isConstantTensor = false;
+};
+
+//===----------------------------------------------------------------------===//
+// ConstantSubgraphAnalyser
+//===----------------------------------------------------------------------===//
+
+class ConstantSubgraphAnalyser
+    : public SparseForwardDataFlowAnalysis<Lattice<IsConstantTensor>> {
+public:
+  using SparseForwardDataFlowAnalysis::SparseForwardDataFlowAnalysis;
+
+  LogicalResult visitOperation(Operation *op,
+                      ArrayRef<const Lattice<IsConstantTensor> *> operands,
+                      ArrayRef<Lattice<IsConstantTensor> *> results) override;
+
+  void setToEntryState(Lattice<IsConstantTensor> *lattice) override;
+};
+
+//===----------------------------------------------------------------------===//
+// RunConstantSubgraphAnalyser
+//===----------------------------------------------------------------------===//
+
+/// Runs constant subgraph analysis on the IR defined by `op`.
+struct RunConstantSubgraphAnalyser {
+public:
+  RunConstantSubgraphAnalyser();
+
+  void run(Operation *op);
+
+  bool getIsConstantTensor(Value val);
+
+private:
+  /// Stores the result of the analysis.
+  DataFlowSolver solver;
+
+  void getConstantSubgraph(DataFlowSolver &solver, Operation *topFunc);
+};
+} // end namespace dataflow
+} // end namespace mlir
+
+#endif // MLIR_ANALYSIS_DATAFLOW_CONSTANTSUBGRAPHANALYSER_H

include/gc/Dialect/OneDNNGraph/OneDNNGraphDialect.td

@@ -22,6 +22,8 @@ def OneDNNGraphDialect : Dialect {
         This dialect follows oneDNN Graph Specification.
     }];
     let cppNamespace = "::mlir::onednn_graph";
+
+    let hasOperationAttrVerify = 1;
 }
 
 #endif // ONEDNNGRAPH_DIALECT

include/gc/Transforms/Passes.h

@@ -120,8 +120,13 @@ void populateGPUPipeline(mlir::OpPassManager &);
 #endif
 
 #define GEN_PASS_DECL
+#define GEN_PASS_DECL_CONSTANTSUBGRAPHANALYSIS
+#define GEN_PASS_DECL_CONSTANTTENSORFOLDING
 #include "gc/Transforms/Passes.h.inc"
 
+std::unique_ptr<Pass> createConstantSubgraphAnalysisPass();
+std::unique_ptr<Pass> createConstantTensorFoldingPass();
+
 #define GEN_PASS_REGISTRATION
 #include "gc/Transforms/Passes.h.inc"
 } // namespace gc

include/gc/Transforms/Passes.td

@@ -169,6 +169,26 @@ def MergeNestedForall : Pass<"merge-nested-forall"> {
   let dependentDialects = ["scf::SCFDialect"];
 }
 
+def ConstantSubgraphAnalysis : Pass<"constant-subgraph-analysis"> {
+  let summary = "Constant Subgraph Analysis";
+  let description = [{
+    This pass implements a constant subgraph analysis.
+  }];
+  let constructor = "mlir::gc::createConstantSubgraphAnalysisPass()";
+}
+
+def ConstantTensorFolding : Pass<"constant-tensor-folding"> {
+  let summary = "Constant Tensor Folding Transform";
+  let description = [{
+    This pass implements a constant tensor folding transform.
+  }];
+  let constructor = "mlir::gc::createConstantTensorFoldingPass()";
+  let dependentDialects = [
+      "tensor::TensorDialect",
+      "linalg::LinalgDialect",
+      "LLVM::LLVMDialect"];
+}
+
 def FoldTensorOperation : Pass<"fold-tensor-operation"> {
   let summary = "Fold some tensor operation";
   let description = [{

lib/gc/Analysis/CMakeLists.txt

@@ -4,6 +4,7 @@ gc_set_mlir_link_components(MLIR_LINK_COMPONENTS
 
 gc_add_mlir_library(GcAnalysis
   TargetDescriptionAnalysis.cpp
+  DataFlow/ConstantSubgraphAnalyser.cpp
   MatmulConfigAnalysis.cpp
 
   DEPENDS

lib/gc/Analysis/DataFlow/ConstantSubgraphAnalyser.cpp

@@ -0,0 +1,187 @@
+//===-- ConstantSubgraphAnalyser.cpp - Constant subgraph  -------*- C++ -*-===//
+//
+// This file is 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 <cassert>
+#include <unordered_set>
+
+#include "gc/Analysis/DataFlow/ConstantSubgraphAnalyser.h"
+#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/Analysis/DataFlow/SparseAnalysis.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/SCF/IR/SCF.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Dialect.h"
+#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/LogicalResult.h"
+#include "mlir/Transforms/Passes.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "in-constant-subgraph"
+
+using namespace mlir;
+using namespace mlir::dataflow;
+
+//===----------------------------------------------------------------------===//
+// IsConstantTensor
+//===----------------------------------------------------------------------===//
+
+void IsConstantTensor::print(raw_ostream &os) const {
+  if (isUninitialized()) {
+    os << "<UNINITIALIZED>";
+    return;
+  }
+  os << getIsConstantTensor();
+}
+
+//===----------------------------------------------------------------------===//
+// ConstantSubgraphAnalyser
+//===----------------------------------------------------------------------===//
+
+LogicalResult ConstantSubgraphAnalyser::visitOperation(
+    Operation *op, ArrayRef<const Lattice<IsConstantTensor> *> operands,
+    ArrayRef<Lattice<IsConstantTensor> *> results) {
+  LLVM_DEBUG(llvm::dbgs() << "ConstantSubgraphAnalyser: Visiting operation:\n"
+                          << *op << "\n");
+
+  bool in = true;
+  if (op->hasTrait<OpTrait::ConstantLike>()) {
+    LLVM_DEBUG(llvm::dbgs() << "Curr op is a Constant op\n");
+    in = true;
+  } else if (operands.empty()) { // For example, tensor.empty()
+    LLVM_DEBUG(llvm::dbgs() << "Curr op has 0 operand, constant\n");
+    in = true;
+  } else {
+    LLVM_DEBUG(llvm::dbgs() << "Curr op has " << operands.size()
+                            << " operands, check if constant\n");
+    for (auto *operandLattice : operands) {
+      auto operandState = operandLattice->getValue().getIsConstantTensor();
+      LLVM_DEBUG(llvm::dbgs() << "Operand: " << operandLattice->getPoint()
+                              << ", lattice value: " << operandState << "\n");
+      if (!operandState) {
+        in = false;
+        break;
+      }
+    }
+  }
+
+  // lattice in results should be in unintialized state.
+  if (!in) {
+    LLVM_DEBUG(llvm::dbgs() << "Curr op not in constant subgraph\n");
+    for (auto lattice : results) {
+      propagateIfChanged(lattice, lattice->join(IsConstantTensor(true, false)));
+    }
+  } else {
+    LLVM_DEBUG(llvm::dbgs() << "Curr op in constant subgraph\n");
+    for (auto lattice : results) {
+      propagateIfChanged(lattice, lattice->join(IsConstantTensor(true, true)));
+    }
+  }
+  return LogicalResult::success();
+}
+
+void ConstantSubgraphAnalyser::setToEntryState(
+    Lattice<IsConstantTensor> *lattice) {
+  if (auto blockArg = cast<BlockArgument>(lattice->getPoint())) {
+    auto parentOp = blockArg.getParentBlock()->getParentOp();
+    auto parentOpAttr = parentOp->getAttrDictionary();
+
+    std::unordered_set<int> constArgsIndexes;
+    std::optional<NamedAttribute> compiletimeConstArgs =
+        parentOpAttr.getNamed("compiletime_const_args_index");
+    if (compiletimeConstArgs.has_value()) {
+      for (auto id :
+           llvm::dyn_cast<ArrayAttr>(compiletimeConstArgs->getValue())) {
+        constArgsIndexes.insert(llvm::cast<IntegerAttr>(id).getInt());
+      }
+    }
+    std::optional<NamedAttribute> runtimeConstArgs =
+        parentOpAttr.getNamed("runtime_const_args_index");
+    if (runtimeConstArgs.has_value()) {
+      for (auto id : llvm::dyn_cast<ArrayAttr>(runtimeConstArgs->getValue())) {
+        constArgsIndexes.insert(llvm::cast<IntegerAttr>(id).getInt());
+      }
+    }
+
+    if (constArgsIndexes.count(blockArg.getArgNumber())) {
+      LLVM_DEBUG(llvm::dbgs() << "Block argument: " << blockArg
+                              << " is marked as constant\n");
+      propagateIfChanged(lattice, lattice->join(IsConstantTensor(true, true)));
+      return;
+    }
+    propagateIfChanged(lattice, lattice->join(IsConstantTensor(true, false)));
+  } else {
+    propagateIfChanged(lattice,
+                       lattice->join(IsConstantTensor::getUninitialized()));
+  }
+}
+
+//===----------------------------------------------------------------------===//
+// RunConstantSubgraphAnalyser
+//===----------------------------------------------------------------------===//
+
+/// Get the operations whose inputs and outputs are all constant values.
+/// These operations will be put into a seperate subgraph.
+void RunConstantSubgraphAnalyser::getConstantSubgraph(DataFlowSolver &solver,
+                                                      Operation *topFunc) {
+  OpBuilder builder(topFunc->getContext());
+  SmallVector<Operation *> constantOperations;
+
+  Block &block = topFunc->getRegions().front().getBlocks().front();
+  for (Operation &op : llvm::make_early_inc_range(block)) {
+    // If all the result values of a op are const, we mark this op as const.
+    bool resultsAllConstant = true;
+    if (op.getNumResults() == 0)
+      continue;
+
+    for (Value res : op.getResults()) {
+      auto *lattice = solver.lookupState<Lattice<IsConstantTensor>>(res);
+      if (!lattice || lattice->getValue().isUninitialized()) {
+        resultsAllConstant = false;
+        break;
+      }
+      const IsConstantTensor &latticeValue = lattice->getValue();
+      if (!latticeValue.getIsConstantTensor()) {
+        resultsAllConstant = false;
+        break;
+      }
+    }
+    if (resultsAllConstant) {
+      op.setAttr("onednn_graph.in_const_subgraph", builder.getBoolAttr(true));
+      constantOperations.push_back(&op);
+    }
+  }
+
+  if (constantOperations.empty())
+    return;
+}
+
+RunConstantSubgraphAnalyser::RunConstantSubgraphAnalyser() {
+  solver.load<DeadCodeAnalysis>();
+  solver.load<ConstantSubgraphAnalyser>();
+}
+
+void RunConstantSubgraphAnalyser::run(Operation *op) {
+  if (failed(solver.initializeAndRun(op)))
+    return;
+
+  getConstantSubgraph(solver, op);
+}
+
+bool RunConstantSubgraphAnalyser::getIsConstantTensor(Value val) {
+  auto *lattice = solver.lookupState<Lattice<IsConstantTensor>>(val);
+  const IsConstantTensor &latticeValue = lattice->getValue();
+  return latticeValue.getIsConstantTensor();
+}

lib/gc/Dialect/OneDNNGraph/OneDNNGraphDialect.cpp

@@ -25,3 +25,9 @@ void OneDNNGraphDialect::initialize() {
 #include "gc/Dialect/OneDNNGraph/OneDNNGraphOps.cpp.inc"
       >();
 }
+
+LogicalResult
+OneDNNGraphDialect::verifyOperationAttribute(Operation *op,
+                                             NamedAttribute attr) {
+  return success();
+}

lib/gc/Transforms/CMakeLists.txt

@@ -16,6 +16,8 @@ gc_add_mlir_library(GcPasses
   IterativeTilingAndFusion.cpp
   TilingUsingInterfaceX.cpp
   VerifyTargetDescription.cpp
+  ConstantSubgraphAnalysis.cpp
+  ConstantTensorFolding.cpp
   DecomposeAggregatedOps.cpp
   DeepTileContractionOp.cpp
   TilingUtil.cpp
@@ -36,6 +38,7 @@ gc_add_mlir_library(GcPasses
     ${MLIR_LINK_COMPONENTS}
     ${GC_ONEDNN_DIALECT_LIB_NAME}
     GcInterface
+    GcAnalysis
     MLIRMicrokernelTransforms
   )