compact pim IR

src/PIM/Dialect/Pim/Pim.td

@@ -50,9 +50,7 @@ def PimCoreBatchOp : PimOp<"core_batch", [SingleBlock, IsolatedFromAbove, AttrSi
     Variadic<PimTensor>:$inputs
   );
 
-  let assemblyFormat = [{
-    `lanes` $laneCount `(` $weights `)` `[` $inputs `]` attr-dict regions `:` type($weights) `[` type($inputs) `]` `->` `(` `)`
-  }];
+  let hasCustomAssemblyFormat = 1;
 }
 
 def PimHaltOp : PimOp<"halt", [Terminator]> {
@@ -63,6 +61,48 @@ def PimHaltOp : PimOp<"halt", [Terminator]> {
   }];
 }
 
+def PimYieldOp : PimOp<"yield", [Terminator]> {
+  let summary = "Yield results from a Pim region";
+
+  let arguments = (ins
+    Variadic<PimTensor>:$outputs
+  );
+
+  let hasCustomAssemblyFormat = 1;
+}
+
+def PimMapOp : PimOp<"map", [SingleBlock]> {
+  let summary = "Apply the same lane-local region to many independent tensors";
+
+  let arguments = (ins
+    Variadic<PimTensor>:$inputs
+  );
+
+  let results = (outs
+    Variadic<PimTensor>:$outputs
+  );
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// Tensor Utilities
+//===----------------------------------------------------------------------===//
+
+def PimEmptyManyOp : PimOp<"empty_many", []> {
+  let summary = "Create many identical empty tensors";
+
+  let results = (outs
+    Variadic<AnyRankedTensor>:$outputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Communication
 //===----------------------------------------------------------------------===//
@@ -81,6 +121,18 @@ def PimSendOp : PimOp<"send", []> {
   }];
 }
 
+def PimSendManyOp : PimOp<"send_many", []> {
+  let summary = "Send multiple tensors to target cores";
+
+  let arguments = (ins
+    DenseI32ArrayAttr:$targetCoreIds,
+    Variadic<PimTensor>:$inputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
 def PimSendBatchOp : PimOp<"send_batch", []> {
   let summary = "Send a per-lane tensor to target cores from a batched core";
 
@@ -90,9 +142,19 @@ def PimSendBatchOp : PimOp<"send_batch", []> {
     DenseI32ArrayAttr:$targetCoreIds
   );
 
-  let assemblyFormat = [{
-    `(` $input `)` attr-dict `:` type($input) `->` `(` `)`
-  }];
+  let hasCustomAssemblyFormat = 1;
+}
+
+def PimSendManyBatchOp : PimOp<"send_many_batch", []> {
+  let summary = "Send multiple per-lane tensors to target cores from a batched core";
+
+  let arguments = (ins
+    DenseI32ArrayAttr:$targetCoreIds,
+    Variadic<PimTensor>:$inputs
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
 }
 
 def PimReceiveOp : PimOp<"receive", [DestinationStyleOpInterface]> {
@@ -119,6 +181,28 @@ def PimReceiveOp : PimOp<"receive", [DestinationStyleOpInterface]> {
   }];
 }
 
+def PimReceiveManyOp : PimOp<"receive_many", [DestinationStyleOpInterface]> {
+  let summary = "Receive multiple tensors from source cores";
+
+  let arguments = (ins
+    Variadic<PimTensor>:$outputBuffers,
+    DenseI32ArrayAttr:$sourceCoreIds
+  );
+
+  let results = (outs
+    Variadic<PimTensor>:$outputs
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutputBuffersMutable();
+    }
+  }];
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
 def PimReceiveBatchOp : PimOp<"receive_batch", [DestinationStyleOpInterface]> {
   let summary = "Receive per-lane tensors from source cores into a batched core";
 
@@ -138,9 +222,29 @@ def PimReceiveBatchOp : PimOp<"receive_batch", [DestinationStyleOpInterface]> {
     }
   }];
 
-  let assemblyFormat = [{
-    `(` $outputBuffer `)` attr-dict `:` type($outputBuffer) `->` type($output)
+  let hasCustomAssemblyFormat = 1;
+}
+
+def PimReceiveManyBatchOp : PimOp<"receive_many_batch", [DestinationStyleOpInterface]> {
+  let summary = "Receive multiple per-lane tensors from source cores into a batched core";
+
+  let arguments = (ins
+    Variadic<PimTensor>:$outputBuffers,
+    DenseI32ArrayAttr:$sourceCoreIds
+  );
+
+  let results = (outs
+    Variadic<PimTensor>:$outputs
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutputBuffersMutable();
+    }
   }];
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
 }
 
 def PimMemCopyHostToDevOp : PimOp<"memcp_hd", [DestinationStyleOpInterface]> {
@@ -247,6 +351,55 @@ def PimMemCopyOp : PimOp<"memcp", [DestinationStyleOpInterface]> {
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Tensor utilities
+//===----------------------------------------------------------------------===//
+
+def PimExtractRowsOp : PimOp<"extract_rows", [DestinationStyleOpInterface]> {
+  let summary = "Extract every row of a rank-2 tensor as separate rank-2 row tensors";
+
+  let arguments = (ins
+    PimTensor:$input,
+    Variadic<PimTensor>:$outputBuffers
+  );
+
+  let results = (outs
+    Variadic<PimTensor>:$outputs
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutputBuffersMutable();
+    }
+  }];
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def PimConcatOp : PimOp<"concat", [DestinationStyleOpInterface]> {
+  let summary = "Concatenate tensors";
+
+  let arguments = (ins
+    I64Attr:$axis,
+    Variadic<PimTensor>:$inputs,
+    PimTensor:$outputBuffer
+  );
+
+  let results = (outs
+    PimTensor:$output
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutputBufferMutable();
+    }
+  }];
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Math
 //===----------------------------------------------------------------------===//