compact spatial IR through different new operations and dedicated syntax

fast spatial node merging with batch operations

src/PIM/Dialect/Spatial/Channels.hpp

@@ -0,0 +1,43 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/StringRef.h"
+
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+namespace onnx_mlir::spatial {
+
+struct ChannelEndpoints {
+  SpatChannelSendOp send;
+  SpatChannelReceiveOp receive;
+};
+
+class Channels {
+public:
+  using ChannelId = int64_t;
+
+  explicit Channels(mlir::func::FuncOp funcOp);
+
+  ChannelId allocate();
+
+  void insertSend(SpatChannelSendOp sendOp);
+  void insertReceive(SpatChannelReceiveOp receiveOp);
+  void eraseSend(SpatChannelSendOp sendOp);
+  void eraseReceive(SpatChannelReceiveOp receiveOp);
+
+  llvm::FailureOr<ChannelEndpoints> lookup(ChannelId id) const;
+  llvm::FailureOr<SpatChannelReceiveOp> getReceiveFor(SpatChannelSendOp sendOp) const;
+  llvm::FailureOr<SpatChannelSendOp> getSendFor(SpatChannelReceiveOp receiveOp) const;
+
+  mlir::LogicalResult verify() const;
+
+private:
+  ChannelId nextChannelId = 0;
+  llvm::DenseMap<ChannelId, ChannelEndpoints> endpoints;
+};
+
+} // namespace onnx_mlir::spatial