reduce spatial compile-times in convolutions using a scf.for instead of materializing a huge number of instructions

src/PIM/Common/PimCommon.hpp

@@ -6,6 +6,7 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Value.h"
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 
@@ -20,6 +21,13 @@ struct ResolvedContiguousAddress {
   int64_t byteOffset = 0;
 };
 
+struct StaticValueKnowledge {
+  llvm::DenseMap<mlir::Value, int64_t> indexValues;
+  llvm::DenseMap<mlir::Value, mlir::Value> aliases;
+
+  StaticValueKnowledge() {}
+};
+
 std::string getOutputDir();
 
 void createDirectory(const std::string& directory);
@@ -52,5 +60,29 @@ bool isMemoryContiguous(llvm::ArrayRef<int64_t> srcShape,
                         llvm::ArrayRef<int64_t> strides);
 
 llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(mlir::Value value);
+llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddress(mlir::Value value,
+                                                                    const StaticValueKnowledge& knowledge);
+
+llvm::FailureOr<int64_t> resolveIndexValue(mlir::Value value);
+llvm::FailureOr<int64_t> resolveIndexValue(mlir::Value value, const StaticValueKnowledge& knowledge);
+
+/// Follows alias and view/DPS chains using `knowledge` to find the value an scf.for
+/// iter-arg is ultimately backed by. Used when interpreting scf.for loop carries.
+mlir::Value resolveLoopCarriedAlias(mlir::Value value, const StaticValueKnowledge& knowledge);
+
+/// Returns true for ops inside a pim.core body that do not emit any PIM instruction and
+/// only contribute to static addressing or index computations (arith integer math,
+/// memref view ops, memref.alloc, arith.constant).
+bool isCoreStaticAddressOp(mlir::Operation* op);
+
+/// Walks `block` (the body of a pim.core region or an scf.for nested in it), statically
+/// unrolling any scf.for with resolvable bounds using `knowledge`. For each remaining op
+/// that is not skipped (pim.halt, scf.yield, or isCoreStaticAddressOp), `callback` is
+/// invoked with the op and the in-scope knowledge. The walker keeps going after a callback
+/// failure so callers can collect multiple diagnostics, but propagates the overall result.
+mlir::LogicalResult
+walkPimCoreBlock(mlir::Block& block,
+                 const StaticValueKnowledge& knowledge,
+                 llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback);
 
 } // namespace onnx_mlir