compact spatial IR through different new operations and dedicated syntax

fast spatial node merging with batch operations

src/PIM/Dialect/Pim/Pim.td

@@ -39,6 +39,22 @@ def PimCoreOp : PimOp<"core", [SingleBlock]> {
   }];
 }
 
+def PimCoreBatchOp : PimOp<"core_batch", [SingleBlock, AttrSizedOperandSegments]> {
+  let summary = "Execute equivalent batched core bodies";
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let arguments = (ins
+    I32Attr:$laneCount,
+    Variadic<PimTensor>:$weights,
+    Variadic<PimTensor>:$inputs
+  );
+
+  let assemblyFormat = [{
+    `lanes` $laneCount `(` $weights `)` `[` $inputs `]` attr-dict regions `:` type($weights) `[` type($inputs) `]` `->` `(` `)`
+  }];
+}
+
 def PimHaltOp : PimOp<"halt", [Terminator]> {
   let summary = "Halt execution of the core";
 
@@ -65,6 +81,20 @@ def PimSendOp : PimOp<"send", []> {
   }];
 }
 
+def PimSendBatchOp : PimOp<"send_batch", []> {
+  let summary = "Send a per-lane tensor to target cores from a batched core";
+
+  let arguments = (ins
+    PimTensor:$input,
+    I32Attr:$size,
+    DenseI32ArrayAttr:$targetCoreIds
+  );
+
+  let assemblyFormat = [{
+    `(` $input `)` attr-dict `:` type($input) `->` `(` `)`
+  }];
+}
+
 def PimReceiveOp : PimOp<"receive", [DestinationStyleOpInterface]> {
   let summary = "Receive a tensor from another core";
 
@@ -89,6 +119,30 @@ def PimReceiveOp : PimOp<"receive", [DestinationStyleOpInterface]> {
   }];
 }
 
+def PimReceiveBatchOp : PimOp<"receive_batch", [DestinationStyleOpInterface]> {
+  let summary = "Receive per-lane tensors from source cores into a batched core";
+
+  let arguments = (ins
+    PimTensor:$outputBuffer,
+    I32Attr:$size,
+    DenseI32ArrayAttr:$sourceCoreIds
+  );
+
+  let results = (outs
+    PimTensor:$output
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutputBufferMutable();
+    }
+  }];
+
+  let assemblyFormat = [{
+    `(` $outputBuffer `)` attr-dict `:` type($outputBuffer) `->` type($output)
+  }];
+}
+
 def PimMemCopyHostToDevOp : PimOp<"memcp_hd", [DestinationStyleOpInterface]> {
   let summary = "Copy a memory region from host memory into device memory";
 
@@ -115,6 +169,32 @@ def PimMemCopyHostToDevOp : PimOp<"memcp_hd", [DestinationStyleOpInterface]> {
   }];
 }
 
+def PimMemCopyHostToDevBatchOp : PimOp<"memcp_hd_batch", [DestinationStyleOpInterface]> {
+  let summary = "Copy a per-lane tensor from host memory into device memory inside a batched core";
+
+  let arguments = (ins
+    PimTensor:$deviceTarget,
+    PimTensor:$hostSource,
+    I32Attr:$deviceTargetOffset,
+    I32Attr:$hostSourceOffset,
+    I32Attr:$size
+  );
+
+  let results = (outs
+    PimTensor:$output
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getDeviceTargetMutable();
+    }
+  }];
+
+  let assemblyFormat = [{
+    `(` $deviceTarget `,` $hostSource `)` attr-dict `:` `(` type($deviceTarget) `,` type($hostSource) `)` `->` type($output)
+  }];
+}
+
 def PimMemCopyDevToHostOp : PimOp<"memcp_dh", [DestinationStyleOpInterface]> {
   let summary = "Copy a memory region from device memory into host memory";