mirror of
https://github.com/zama-ai/concrete.git
synced 2026-04-17 03:00:54 -04:00
7a295f89bd
This makes the trip counts of operations in the TFHE statistics pass as well as the per-location memory usage statistics in the memory usage statistics pass optional. These values are unset if the trip count could not be determined statically.
62 lines
1.6 KiB
C++
62 lines
1.6 KiB
C++
// Part of the Concrete Compiler Project, under the BSD3 License with Zama
|
|
// Exceptions. See
|
|
// https://github.com/zama-ai/concrete/blob/main/LICENSE.txt
|
|
// for license information.
|
|
|
|
#ifndef CONCRETELANG_ANALYSIS_UTILS_H
|
|
#define CONCRETELANG_ANALYSIS_UTILS_H
|
|
|
|
#include <boost/outcome.h>
|
|
#include <concretelang/Common/Error.h>
|
|
#include <limits>
|
|
#include <mlir/Dialect/SCF/IR/SCF.h>
|
|
#include <mlir/IR/Location.h>
|
|
|
|
namespace mlir {
|
|
namespace concretelang {
|
|
class TripCountTracker {
|
|
public:
|
|
void pushTripCount(mlir::Operation *op, std::optional<int64_t> n) {
|
|
if (tripCount.has_value()) {
|
|
if (n.has_value()) {
|
|
assert(std::numeric_limits<int64_t>::max() / n.value() >
|
|
tripCount.value());
|
|
|
|
tripCount = tripCount.value() * n.value();
|
|
} else {
|
|
savedTripCount = *tripCount;
|
|
tripCount = std::nullopt;
|
|
firstDynamicTripCountOp = op;
|
|
}
|
|
}
|
|
}
|
|
|
|
void popTripCount(mlir::Operation *op, std::optional<int64_t> n) {
|
|
if (n.has_value()) {
|
|
if (tripCount.has_value()) {
|
|
tripCount = tripCount.value() / n.value();
|
|
}
|
|
} else {
|
|
if (firstDynamicTripCountOp == op) {
|
|
tripCount = savedTripCount;
|
|
}
|
|
firstDynamicTripCountOp = nullptr;
|
|
}
|
|
}
|
|
|
|
std::optional<int64_t> getTripCount() { return tripCount; }
|
|
|
|
protected:
|
|
std::optional<int64_t> tripCount = 1;
|
|
size_t savedTripCount;
|
|
mlir::Operation *firstDynamicTripCountOp = nullptr;
|
|
};
|
|
|
|
/// Get the string representation of a location
|
|
std::string locationString(mlir::Location loc);
|
|
|
|
} // namespace concretelang
|
|
} // namespace mlir
|
|
|
|
#endif
|