promote weight inputs to actual weights in spat compute nodes

src/PIM/Conversion/ONNXToSpatial/Common.hpp

@@ -1,16 +1,19 @@
 #pragma once
 
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Block.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/ValueRange.h"
-#include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 #include <cassert>
 #include <type_traits>
 #include <utility>
 
+#include "llvm/ADT/SmallPtrSet.h"
+
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -102,6 +105,42 @@ inline auto getTensorShape(mlir::Value tensor) {
   return mlir::cast<mlir::RankedTensorType>(tensor.getType()).getShape();
 }
 
+inline bool isWeightLikeComputeOperand(mlir::Value value) {
+  auto rankedType = mlir::dyn_cast<mlir::RankedTensorType>(value.getType());
+  if (!rankedType || !isMatrixShape(rankedType.getShape()))
+    return false;
+
+  llvm::SmallPtrSet<mlir::Operation*, 8> visited;
+
+  while (auto* definingOp = value.getDefiningOp()) {
+    if (!visited.insert(definingOp).second)
+      return false;
+    if (hasWeightAlways(definingOp))
+      return true;
+
+    if (auto extractSliceOp = mlir::dyn_cast<mlir::tensor::ExtractSliceOp>(definingOp)) {
+      value = extractSliceOp.getSource();
+      continue;
+    }
+    if (auto expandShapeOp = mlir::dyn_cast<mlir::tensor::ExpandShapeOp>(definingOp)) {
+      value = expandShapeOp.getSrc();
+      continue;
+    }
+    if (auto collapseShapeOp = mlir::dyn_cast<mlir::tensor::CollapseShapeOp>(definingOp)) {
+      value = collapseShapeOp.getSrc();
+      continue;
+    }
+    if (auto transposeOp = mlir::dyn_cast<mlir::ONNXTransposeOp>(definingOp)) {
+      value = transposeOp.getData();
+      continue;
+    }
+
+    return false;
+  }
+
+  return false;
+}
+
 namespace detail {
 
 inline mlir::ValueRange getBlockArgs(mlir::Block* block) { return mlir::ValueRange(block->getArguments()); }
@@ -111,6 +150,11 @@ decltype(auto) invokeWithBlockArgs(Fn&& fn, mlir::Block* block, std::index_seque
   return std::forward<Fn>(fn)(block->getArgument(Is)...);
 }
 
+template <typename Fn, size_t... Is>
+decltype(auto) invokeWithValues(Fn&& fn, mlir::ArrayRef<mlir::Value> values, std::index_sequence<Is...>) {
+  return std::forward<Fn>(fn)(values[Is]...);
+}
+
 template <size_t>
 using ValueArg = mlir::Value;