[DAPHNE-#854] Add Passes for Matrix Result Property Recording and Insertion in IR

UserConfig.json

@@ -24,6 +24,9 @@
     "explain_vectorized": false,
     "explain_obj_ref_mgnt": false,
     "explain_mlir_codegen": false,
+    "enable_property_recording": false,
+    "enable_property_insert": false,
+    "properties_file_path": "properties.json",
     "taskPartitioningScheme": "STATIC",
     "numberOfThreads": -1,
     "minimumTaskSize": 1,
@@ -71,5 +74,5 @@
             "format": "%^[%n %L]:%$ %v"
         }
     ],
-    "sparsity_threshold": 0.25
+    "sparsity_threshold": 0.49
 }

doc/RunningDaphneLocally.md

@@ -80,6 +80,14 @@ Note that some of the more specific options are described in the documentation p
 
     Turns on the automatic selection of a suitable matrix representation (currently dense or sparse (CSR)). *Experimental feature.*
 
+- **`--enable_property_recording`**
+
+    Enables recording of certain properties (currently only sparsity) from all matrix result objects in the Intermediate Representation (IR). The recorded properties are saved in a JSON file named properties.json located in the root directory of Daphne.
+
+- **`--enable_property_insert`** and **`--properties_file_path=<path/to/propertiesFile>`**
+
+    Allows the insertion of previously recorded properties into operations with matrix results from a JSON file. By default, the file used is properties.json in the Daphne root directory. Optionally, a different file can be specified using the `--properties_file_path` argument.
+
 ## Return Codes
 
 If `daphne` terminates normally, one of the following status codes is returned:

src/api/cli/DaphneUserConfig.h

@@ -75,6 +75,9 @@ struct DaphneUserConfig {
     bool explain_obj_ref_mgnt = false;
     bool explain_mlir_codegen = false;
     bool statistics = false;
+    bool enable_property_recording = false;
+    bool enable_property_insert = false;
+    std::string properties_file_path = "properties.json";
 
     bool force_cuda = false;
 

src/api/internal/daphne_internal.cpp

@@ -32,6 +32,7 @@
 #include <util/DaphneLogger.h>
 #include <util/KernelDispatchMapping.h>
 #include <util/Statistics.h>
+#include <util/PropertyLogger.h>
 
 #include "mlir/ExecutionEngine/ExecutionEngine.h"
 #include "mlir/IR/Builders.h"
@@ -310,17 +311,39 @@ int startDAPHNE(int argc, const char **argv, DaphneLibResult *daphneLibRes, int
                                                    "e.g., `--args x=1,y=2.2`"),
                                               CommaSeparated);
     const std::string configFileInitValue = "-";
-    static opt<string> configFile("config", cat(daphneOptions),
-                                  desc("A JSON file that contains the DAPHNE configuration"), value_desc("filename"),
-                                  llvm::cl::init(configFileInitValue));
-
-    static opt<bool> enableStatistics("statistics", cat(daphneOptions), desc("Enables runtime statistics output."));
-
-    static opt<bool> enableProfiling("enable-profiling", cat(daphneOptions), desc("Enable profiling support"));
-    static opt<bool> timing("timing", cat(daphneOptions),
-                            desc("Enable timing of high-level steps (start-up, "
-                                 "parsing, compilation, execution) and print "
-                                 "the times to stderr in JSON format"));
+    static opt<string> configFile(
+        "config", cat(daphneOptions),
+        desc("A JSON file that contains the DAPHNE configuration"),
+        value_desc("filename"),
+        llvm::cl::init(configFileInitValue)
+    );
+
+    static opt<bool> enableStatistics(
+        "statistics", cat(daphneOptions),
+        desc("Enables runtime statistics output."));
+
+    static opt<bool> enablePropertyRecording(
+        "enable_property_recording", cat(daphneOptions),
+        desc("records runtime properties and outputs it in JSON."));
+
+    static opt<bool> enablePropertyInsert(
+        "enable_property_insert", cat(daphneOptions),
+        desc("inserts runtime properties from properties.json of previous run."));
+
+    static opt<std::string> properties_file_path(
+        "properties_file_path", cat(daphneOptions),
+        llvm::cl::desc("Path to the Properties File in JSON for the Property insert."),
+        llvm::cl::init("properties.json")
+    );
+
+    static opt<bool> enableProfiling (
+            "enable-profiling", cat(daphneOptions),
+            desc("Enable profiling support")
+    );
+    static opt<bool> timing (
+            "timing", cat(daphneOptions),
+            desc("Enable timing of high-level steps (start-up, parsing, compilation, execution) and print the times to stderr in JSON format")
+    );
 
     // Positional arguments ---------------------------------------------------
 
@@ -385,6 +408,7 @@ int startDAPHNE(int argc, const char **argv, DaphneLibResult *daphneLibRes, int
         user_config.matmul_tile = true;
     }
     user_config.use_mlir_hybrid_codegen = performHybridCodegen;
+
 
     if (!libDir.getValue().empty())
         user_config.libdir = libDir.getValue();
@@ -478,6 +502,11 @@ int startDAPHNE(int argc, const char **argv, DaphneLibResult *daphneLibRes, int
     }
 
     user_config.statistics = enableStatistics;
+    user_config.enable_property_recording = enablePropertyRecording;
+    user_config.enable_property_insert = enablePropertyInsert;
+
+    if (user_config.enable_property_insert)
+        user_config.properties_file_path = properties_file_path.getValue();
 
     if (user_config.use_distributed && distributedBackEndSetup == ALLOCATION_TYPE::DIST_MPI) {
 #ifndef USE_MPI
@@ -652,9 +681,12 @@ int startDAPHNE(int argc, const char **argv, DaphneLibResult *daphneLibRes, int
     if (user_config.statistics)
         Statistics::instance().dumpStatistics(KernelDispatchMapping::instance());
 
-    // explicitly destroying the moduleOp here due to valgrind complaining about
-    // a memory leak otherwise.
+    if (user_config.enable_property_recording)
+        PropertyLogger::instance().savePropertiesAsJson("properties.json");
+
+    // explicitly destroying the moduleOp here due to valgrind complaining about a memory leak otherwise.
     moduleOp->destroy();
+
     return StatusCode::SUCCESS;
 }
 

src/api/python/daphne/utils/consts.py

@@ -36,7 +36,8 @@
 VALID_ARITHMETIC_TYPES = Union['DAGNode', int, float]
 VALID_COMPUTED_TYPES = Union['Matrix', 'Frame', 'Scalar']
 
-TMP_PATH = os.path.join("/tmp/", "DaphneLib")
+#TMP_PATH = os.path.join("/tmp/", "DaphneLib")
+TMP_PATH = os.path.join("/home", "dominic", "daphne", "tmp", "DaphneLib")
 os.makedirs(TMP_PATH, exist_ok=True)
 
 _PROTOTYPE_PATH_ENV_VAR_NAME = "DAPHNELIB_DIR_PATH"

src/compiler/execution/DaphneIrExecutor.cpp

@@ -115,6 +115,16 @@ bool DaphneIrExecutor::runPasses(mlir::ModuleOp module) {
     pm.addNestedPass<mlir::func::FuncOp>(mlir::daphne::createInferencePass());
     pm.addPass(mlir::createCanonicalizerPass());
 
+    if(userConfig_.enable_property_recording)
+        pm.addPass(mlir::daphne::createRecordPropertiesPass());
+
+    if(userConfig_.enable_property_insert)
+    {
+        pm.addPass(mlir::daphne::createInsertPropertiesPass(userConfig_.properties_file_path));
+        pm.addNestedPass<mlir::func::FuncOp>(mlir::daphne::createInferencePass());
+        pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
+    }
+
     if (selectMatrixRepresentations_) {
         pm.addNestedPass<mlir::func::FuncOp>(mlir::daphne::createSelectMatrixRepresentationsPass(userConfig_));
         pm.addNestedPass<mlir::func::FuncOp>(mlir::createCanonicalizerPass());
@@ -175,22 +185,23 @@ bool DaphneIrExecutor::runPasses(mlir::ModuleOp module) {
     pm.addPass(mlir::createCanonicalizerPass());
     pm.addPass(mlir::createCSEPass());
 
+
+
     if (userConfig_.use_obj_ref_mgnt)
         pm.addNestedPass<mlir::func::FuncOp>(mlir::daphne::createManageObjRefsPass());
     if (userConfig_.explain_obj_ref_mgnt)
         pm.addPass(mlir::daphne::createPrintIRPass("IR after managing object references:"));
 
     pm.addNestedPass<mlir::func::FuncOp>(mlir::daphne::createRewriteToCallKernelOpPass(userConfig_, usedLibPaths));
     if (userConfig_.explain_kernels)
-        pm.addPass(mlir::daphne::createPrintIRPass("IR after kernel lowering:"));
-
+        pm.addPass(
+            mlir::daphne::createPrintIRPass("IR after kernel lowering:"));
     pm.addPass(mlir::createConvertSCFToCFPass());
     pm.addNestedPass<mlir::func::FuncOp>(mlir::LLVM::createRequestCWrappersPass());
     pm.addPass(mlir::daphne::createLowerToLLVMPass(userConfig_));
     pm.addPass(mlir::createReconcileUnrealizedCastsPass());
     if (userConfig_.explain_llvm)
         pm.addPass(mlir::daphne::createPrintIRPass("IR after llvm lowering:"));
-
     // Initialize the use of each distinct kernels library to false.
     usedLibPaths = userConfig_.kernelCatalog.getLibPaths();
 

src/compiler/explanation/CMakeLists.txt

@@ -14,6 +14,8 @@
 
 add_mlir_dialect_library(MLIRDaphneExplain
     PrintIRPass.cpp
+    RecordPropertiesPass.cpp
+    InsertPropertiesPass.cpp
 
     DEPENDS
     MLIRDaphneOpsIncGen

src/compiler/explanation/InsertPropertiesPass.cpp

@@ -0,0 +1,183 @@
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Value.h"
+#include "mlir/IR/Visitors.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/raw_ostream.h"
+#include <algorithm>
+#include <cstdint>
+#include <fstream>
+#include <ir/daphneir/Daphne.h>
+#include <ir/daphneir/Passes.h>
+#include <mlir/Dialect/SCF/IR/SCF.h>
+#include <nlohmannjson/json.hpp>
+#include <string>
+
+namespace mlir {
+#define GEN_PASS_DECL_INSERTPROPERTIESPASS
+#define GEN_PASS_DEF_INSERTPROPERTIESPASS
+#include "ir/daphneir/Passes.h.inc"
+} // namespace mlir
+
+using namespace mlir;
+
+nlohmann::json readPropertiesFromFile(const std::string &filename) {
+  std::ifstream file(filename);
+  if (!file.is_open()) {
+    throw std::runtime_error("Failed to open file: " + filename);
+  }
+
+  nlohmann::json properties;
+  file >> properties;
+  return properties;
+}
+
+struct PropertyEntry {
+  uint32_t valueID;
+  nlohmann::json properties;
+
+  bool operator<(const PropertyEntry &other) const {
+    return valueID < other.valueID;
+  }
+};
+
+namespace {
+struct InsertPropertiesPass
+    : public impl::InsertPropertiesPassBase<InsertPropertiesPass> {
+  InsertPropertiesPass() = default;
+
+public:
+  explicit InsertPropertiesPass(std::string properties_file_path)
+      : impl::InsertPropertiesPassBase<InsertPropertiesPass>() {
+    this->properties_file_path = properties_file_path;
+  };
+  void runOnOperation() override;
+};
+} // end of anonymous namespace
+
+void InsertPropertiesPass::runOnOperation() {
+  func::FuncOp func = getOperation();
+  OpBuilder builder(func.getContext());
+
+  nlohmann::json propertiesJson =
+      readPropertiesFromFile(this->properties_file_path);
+  llvm::DenseMap<std::pair<Value, Operation *>, Value> castOpMap;
+
+  std::vector<PropertyEntry> properties;
+  for (auto it = propertiesJson.begin(); it != propertiesJson.end(); ++it) {
+    uint32_t valueID = std::stoi(it.key());
+    properties.push_back({valueID, it.value()});
+  }
+  std::sort(properties.begin(), properties.end());
+  size_t propertyIndex = 0;
+
+  auto insertRecordedProperties = [&](Operation *op) {
+    size_t numResults = op->getNumResults();
+    for (unsigned i = 0; i < numResults && propertyIndex < properties.size();
+         ++i) {
+      Value res = op->getResult(i);
+      nlohmann::json &prop = properties[propertyIndex].properties;
+      auto it = prop.begin();
+      while (it != prop.end()) {
+        const std::string &key = it.key();
+        const nlohmann::json &value = it.value();
+        if (key == "sparsity") {
+          if (value.is_null()) {
+            llvm::errs() << "Error: 'sparsity' is null for property index "
+                         << propertyIndex << "\n";
+            ++it;
+            continue;
+          } else if (!value.is_number()) {
+            llvm::errs()
+                << "Error: 'sparsity' is not a number for property index "
+                << propertyIndex << "\n";
+            ++it;
+            continue;
+          }
+
+          if (res.getType().isa<daphne::MatrixType>()) {
+            auto mt = res.getType().dyn_cast<daphne::MatrixType>();
+            double sparsity = value.get<double>();
+            if (mt) {
+              if ((llvm::isa<scf::ForOp>(op) || llvm::isa<scf::WhileOp>(op) ||
+                   llvm::isa<scf::IfOp>(op))) {
+                builder.setInsertionPointAfter(op);
+                builder.create<daphne::CastOp>(op->getLoc(),
+                                               mt.withSparsity(sparsity), res);
+              }
+
+              else {
+                for (auto &use : res.getUses()) {
+                  Operation *userOp = use.getOwner();
+                  if (isa<scf::ForOp>(userOp) || isa<scf::IfOp>(userOp) ||
+                      isa<scf::WhileOp>(userOp)) {
+                    auto key = std::make_pair(res, userOp);
+
+                    Value castOpValue;
+                    if (castOpMap.count(key)) {
+                      castOpValue = castOpMap[key];
+                    } else {
+                      builder.setInsertionPoint(userOp);
+                      castOpValue =
+                          builder.create<daphne::CastOp>(op->getLoc(), mt, res);
+                      castOpMap[key] = castOpValue;
+                    }
+
+                    userOp->setOperand(use.getOperandNumber(), castOpValue);
+                  }
+                }
+              }
+
+              ++propertyIndex;
+            }
+          }
+        }
+        ++it;
+      }
+    }
+  };
+
+  func.walk<WalkOrder::PreOrder>([&](Operation *op) -> WalkResult {
+    if (propertyIndex >= properties.size())
+      return WalkResult::advance();
+
+    // Skip specific ops that should not be processed
+    if (isa<daphne::RecordPropertiesOp>(op) || op->hasAttr("daphne.value_ids"))
+      return WalkResult::advance();
+
+    if (auto castOp = dyn_cast<daphne::CastOp>(op)) {
+      if (castOp.isRemovePropertyCast()) {
+        return WalkResult::advance();
+      }
+    }
+
+    // Handle loops (scf.for and scf.while) and If blocks as black boxes
+    if (isa<scf::ForOp>(op) || isa<scf::WhileOp>(op) || isa<scf::IfOp>(op)) {
+      insertRecordedProperties(op);
+      return WalkResult::skip();
+    }
+
+    if (auto funcOp = llvm::dyn_cast<func::FuncOp>(op)) {
+      // Check if this is the @main function or a UDF
+      if (funcOp.getName() == "main") {
+        return WalkResult::advance();
+      } else {
+        return WalkResult::skip();
+      }
+    }
+
+    // Process all other operations that output matrix types
+    insertRecordedProperties(op);
+    return WalkResult::advance();
+  });
+
+  if (propertyIndex < properties.size()) {
+    llvm::errs() << "Warning: Not all properties were applied." << "\n";
+  }
+}
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+mlir::daphne::createInsertPropertiesPass(std::string propertiesFilePath) {
+  return std::make_unique<InsertPropertiesPass>(propertiesFilePath);
+}