refactorone

src/PIM/Dialect/Spatial/Spatial.td

@@ -2,8 +2,12 @@
 #define SPATIAL_DIALECT_H
 
 include "mlir/IR/OpBase.td"
+include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/BuiltinTypes.td"
 include "mlir/IR/AttrTypeBase.td"
+include "mlir/IR/RegionKindInterface.td"
+include "mlir/Interfaces/ParallelCombiningOpInterface.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
 
 def SpatialDialect : Dialect {
   let name = "spat";
@@ -22,7 +26,9 @@ def SpatTensor :
 // Execution
 //===----------------------------------------------------------------------===//
 
-def SpatCompute : SpatOp<"compute", [SingleBlock, AttrSizedOperandSegments]> {
+def SpatCompute : SpatOp<"compute",
+    [SingleBlock, AttrSizedOperandSegments,
+     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
   let summary = "Compute region with attached constant weights";
 
   let arguments = (ins
@@ -36,14 +42,20 @@ def SpatCompute : SpatOp<"compute", [SingleBlock, AttrSizedOperandSegments]> {
 
   let regions = (region SizedRegion<1>:$body);
 
+  let extraClassDeclaration = [{
+    ::mlir::BlockArgument getWeightArgument(unsigned idx);
+    ::mlir::BlockArgument getInputArgument(unsigned idx);
+  }];
+
   let hasVerifier = 1;
   let hasFolder = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
 def SpatComputeBatch : SpatOp<"compute_batch",
-    [SingleBlock, AttrSizedOperandSegments]> {
-  let summary = "Compressed batch of independent equivalent compute lanes";
+    [SingleBlock, AttrSizedOperandSegments,
+     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
+  let summary = "Tensor-native batch of equivalent compute lanes with shared weights and packed inputs";
 
   let arguments = (ins
     I32Attr:$laneCount,
@@ -57,10 +69,41 @@ def SpatComputeBatch : SpatOp<"compute_batch",
 
   let regions = (region SizedRegion<1>:$body);
 
+  let extraClassDeclaration = [{
+    ::mlir::BlockArgument getLaneArgument();
+    ::mlir::BlockArgument getWeightArgument(unsigned idx);
+    ::mlir::BlockArgument getInputArgument(unsigned idx);
+    ::mlir::BlockArgument getOutputArgument(unsigned idx);
+  }];
+
   let hasVerifier = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
+def SpatInParallelOp : SpatOp<"in_parallel", [
+       Pure,
+       Terminator,
+       DeclareOpInterfaceMethods<InParallelOpInterface>,
+       HasParent<"SpatComputeBatch">,
+      ] # GraphRegionNoTerminator.traits> {
+  let summary = "Parallel combining terminator for resultful spat.compute_batch";
+
+  let regions = (region SizedRegion<1>:$region);
+
+  let hasCustomAssemblyFormat = 1;
+  let hasVerifier = 1;
+
+  let skipDefaultBuilders = 1;
+  let builders = [
+    OpBuilder<(ins)>,
+  ];
+
+  let extraClassDeclaration = [{
+    ::llvm::iterator_range<::mlir::Block::iterator> getYieldingOps();
+    ::mlir::OpResult getParentResult(int64_t idx);
+  }];
+}
+
 def SpatYieldOp : SpatOp<"yield", [Terminator]> {
   let summary = "Yield results from a compute region";
 
@@ -110,14 +153,14 @@ def SpatChannelSendOp : SpatOp<"channel_send", []> {
   let summary = "Send a tensor through a logical channel";
 
   let arguments = (ins
-    I64Attr:$channelId,
-    I32Attr:$sourceCoreId,
-    I32Attr:$targetCoreId,
+    Index:$channelId,
+    Index:$sourceCoreId,
+    Index:$targetCoreId,
     SpatTensor:$input
   );
 
   let assemblyFormat = [{
-    $input attr-dict `:` type($input)
+    $input `channel` $channelId `from` $sourceCoreId `to` $targetCoreId attr-dict `:` type($input)
   }];
 }
 
@@ -125,9 +168,9 @@ def SpatChannelReceiveOp : SpatOp<"channel_receive", []> {
   let summary = "Receive a tensor from a logical channel";
 
   let arguments = (ins
-    I64Attr:$channelId,
-    I32Attr:$sourceCoreId,
-    I32Attr:$targetCoreId
+    Index:$channelId,
+    Index:$sourceCoreId,
+    Index:$targetCoreId
   );
 
   let results = (outs
@@ -135,31 +178,33 @@ def SpatChannelReceiveOp : SpatOp<"channel_receive", []> {
   );
 
   let assemblyFormat = [{
-    attr-dict `:` type($output)
+    `channel` $channelId `from` $sourceCoreId `to` $targetCoreId attr-dict `:` type($output)
   }];
 }
 
-def SpatChannelSendTensorOp : SpatOp<"channel_send_tensor", []> {
+def SpatChannelSendTensorOp : SpatOp<"channel_send_tensor", [AttrSizedOperandSegments]> {
   let summary = "Send equal contiguous chunks of one tensor through logical channels";
 
   let arguments = (ins
-    DenseI64ArrayAttr:$channelIds,
-    DenseI32ArrayAttr:$sourceCoreIds,
-    DenseI32ArrayAttr:$targetCoreIds,
+    Variadic<Index>:$channelIds,
+    Variadic<Index>:$sourceCoreIds,
+    Variadic<Index>:$targetCoreIds,
     SpatTensor:$input
   );
 
   let hasVerifier = 1;
-  let hasCustomAssemblyFormat = 1;
+  let assemblyFormat = [{
+    $input `channels` `(` $channelIds `)` `from` `(` $sourceCoreIds `)` `to` `(` $targetCoreIds `)` attr-dict `:` type($input)
+  }];
 }
 
-def SpatChannelReceiveTensorOp : SpatOp<"channel_receive_tensor", []> {
+def SpatChannelReceiveTensorOp : SpatOp<"channel_receive_tensor", [AttrSizedOperandSegments]> {
   let summary = "Receive equal contiguous chunks of one tensor from logical channels";
 
   let arguments = (ins
-    DenseI64ArrayAttr:$channelIds,
-    DenseI32ArrayAttr:$sourceCoreIds,
-    DenseI32ArrayAttr:$targetCoreIds
+    Variadic<Index>:$channelIds,
+    Variadic<Index>:$sourceCoreIds,
+    Variadic<Index>:$targetCoreIds
   );
 
   let results = (outs
@@ -167,44 +212,50 @@ def SpatChannelReceiveTensorOp : SpatOp<"channel_receive_tensor", []> {
   );
 
   let hasVerifier = 1;
-  let hasCustomAssemblyFormat = 1;
+  let assemblyFormat = [{
+    `channels` `(` $channelIds `)` `from` `(` $sourceCoreIds `)` `to` `(` $targetCoreIds `)` attr-dict `:` type($output)
+  }];
 }
 
-def SpatChannelSendBatchOp : SpatOp<"channel_send_batch", []> {
+def SpatChannelSendBatchOp : SpatOp<"channel_send_batch", [AttrSizedOperandSegments]> {
   let summary = "Send per-lane tensors through logical channels in a batch body";
 
   let arguments = (ins
-    DenseI64ArrayAttr:$channelIds,
-    DenseI32ArrayAttr:$sourceCoreIds,
-    DenseI32ArrayAttr:$targetCoreIds,
+    Variadic<Index>:$channelIds,
+    Variadic<Index>:$sourceCoreIds,
+    Variadic<Index>:$targetCoreIds,
     SpatTensor:$input
   );
 
   let hasVerifier = 1;
-  let hasCustomAssemblyFormat = 1;
+  let assemblyFormat = [{
+    $input `channels` `(` $channelIds `)` `from` `(` $sourceCoreIds `)` `to` `(` $targetCoreIds `)` attr-dict `:` type($input)
+  }];
 }
 
-def SpatChannelSendTensorBatchOp : SpatOp<"channel_send_tensor_batch", []> {
+def SpatChannelSendTensorBatchOp : SpatOp<"channel_send_tensor_batch", [AttrSizedOperandSegments]> {
   let summary = "Send equal contiguous chunks of one per-lane tensor through logical channels in a batch body";
 
   let arguments = (ins
-    DenseI64ArrayAttr:$channelIds,
-    DenseI32ArrayAttr:$sourceCoreIds,
-    DenseI32ArrayAttr:$targetCoreIds,
+    Variadic<Index>:$channelIds,
+    Variadic<Index>:$sourceCoreIds,
+    Variadic<Index>:$targetCoreIds,
     SpatTensor:$input
   );
 
   let hasVerifier = 1;
-  let hasCustomAssemblyFormat = 1;
+  let assemblyFormat = [{
+    $input `channels` `(` $channelIds `)` `from` `(` $sourceCoreIds `)` `to` `(` $targetCoreIds `)` attr-dict `:` type($input)
+  }];
 }
 
-def SpatChannelReceiveBatchOp : SpatOp<"channel_receive_batch", []> {
+def SpatChannelReceiveBatchOp : SpatOp<"channel_receive_batch", [AttrSizedOperandSegments]> {
   let summary = "Receive a per-lane tensor through logical channels in a batch body";
 
   let arguments = (ins
-    DenseI64ArrayAttr:$channelIds,
-    DenseI32ArrayAttr:$sourceCoreIds,
-    DenseI32ArrayAttr:$targetCoreIds
+    Variadic<Index>:$channelIds,
+    Variadic<Index>:$sourceCoreIds,
+    Variadic<Index>:$targetCoreIds
   );
 
   let results = (outs
@@ -212,16 +263,18 @@ def SpatChannelReceiveBatchOp : SpatOp<"channel_receive_batch", []> {
   );
 
   let hasVerifier = 1;
-  let hasCustomAssemblyFormat = 1;
+  let assemblyFormat = [{
+    `channels` `(` $channelIds `)` `from` `(` $sourceCoreIds `)` `to` `(` $targetCoreIds `)` attr-dict `:` type($output)
+  }];
 }
 
-def SpatChannelReceiveTensorBatchOp : SpatOp<"channel_receive_tensor_batch", []> {
+def SpatChannelReceiveTensorBatchOp : SpatOp<"channel_receive_tensor_batch", [AttrSizedOperandSegments]> {
   let summary = "Receive equal contiguous chunks of one per-lane tensor through logical channels in a batch body";
 
   let arguments = (ins
-    DenseI64ArrayAttr:$channelIds,
-    DenseI32ArrayAttr:$sourceCoreIds,
-    DenseI32ArrayAttr:$targetCoreIds
+    Variadic<Index>:$channelIds,
+    Variadic<Index>:$sourceCoreIds,
+    Variadic<Index>:$targetCoreIds
   );
 
   let results = (outs
@@ -229,7 +282,9 @@ def SpatChannelReceiveTensorBatchOp : SpatOp<"channel_receive_tensor_batch", []>
   );
 
   let hasVerifier = 1;
-  let hasCustomAssemblyFormat = 1;
+  let assemblyFormat = [{
+    `channels` `(` $channelIds `)` `from` `(` $sourceCoreIds `)` `to` `(` $targetCoreIds `)` attr-dict `:` type($output)
+  }];
 }
 
 //===----------------------------------------------------------------------===//
@@ -240,7 +295,7 @@ def SpatVMMOp : SpatOp<"wvmm", []> {
   let summary = "Vector-matrix multiplication within a weighted compute operation";
 
   let arguments = (ins
-    I32Attr:$weightIndex,
+    SpatTensor:$weight,
     SpatTensor:$input
   );
 
@@ -251,7 +306,7 @@ def SpatVMMOp : SpatOp<"wvmm", []> {
   let hasVerifier = 1;
 
   let assemblyFormat = [{
-    `(` $input `)` attr-dict `:` type($input) `->` type($output)
+    `[` $weight `]` `(` $input `)` attr-dict `:` `(` type($weight) `,` type($input) `)` `->` type($output)
   }];
 }
 
@@ -259,7 +314,7 @@ def SpatMVMOp : SpatOp<"Wmvm", []> {
   let summary = "Matrix-vector multiplication within a weighted compute operation";
 
   let arguments = (ins
-    I32Attr:$weightIndex,
+    SpatTensor:$weight,
     SpatTensor:$input
   );
 
@@ -270,7 +325,7 @@ def SpatMVMOp : SpatOp<"Wmvm", []> {
   let hasVerifier = 1;
 
   let assemblyFormat = [{
-    `(` $input `)` attr-dict `:` type($input) `->` type($output)
+    `[` $weight `]` `(` $input `)` attr-dict `:` `(` type($weight) `,` type($input) `)` `->` type($output)
   }];
 }