add PIM accelerator

src/PIM/Dialect/Spatial/Spatial.td

@@ -0,0 +1,355 @@
+#ifndef SPATIAL_DIALECT_H
+#define SPATIAL_DIALECT_H
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/BuiltinTypes.td"
+include "mlir/IR/AttrTypeBase.td"
+
+def SpatialDialect : Dialect {
+  let name = "spat";
+  let summary = "Dialect designed for deep learning computation in a spatial architecture";
+  let cppNamespace = "::onnx_mlir::spatial";
+  let useDefaultTypePrinterParser = 1;
+}
+
+class SpatOp<string mnemonic, list<Trait> traits = []> :
+    Op<SpatialDialect, mnemonic, traits>;
+
+// TODO maybe remove and use AnyRankedTensor directly
+def SpatTensor:
+  AnyTypeOf<[AnyMemRef, AnyRankedTensor], "", "::mlir::ShapedType">;
+
+class SpatType<string name, string typeMnemonic, list<Trait> traits = []>
+    : TypeDef<SpatialDialect, name, traits> {
+  let mnemonic = typeMnemonic;
+}
+
+def SpatChannelType : SpatType<"SpatChannel", "ch"> {
+  let summary = "Virtual channel type";
+}
+
+def SpatWeightedCompute: SpatOp<"compute", [SingleBlock, AttrSizedOperandSegments]> {
+  let summary = "Compute operation, with constant weights already attached";
+
+  let arguments = (ins
+    Variadic<SpatTensor>:$weights,
+    Variadic<SpatTensor>:$inputs
+  );
+
+  let results = (outs
+    Variadic<SpatTensor>:$outputs
+  );
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let hasVerifier = 1;
+
+  let assemblyFormat = [{
+    `[` $weights `]` `(` $inputs `)` attr-dict `:` `[` type($weights) `]` `(` type($inputs) `)` `->` type($outputs) $body
+  }];
+}
+
+def SpatYieldOp: SpatOp<"yield", [Terminator]> {
+  let arguments = (ins
+    Variadic<SpatTensor>:$outputs
+  );
+
+  let assemblyFormat = [{
+    $outputs attr-dict `:` type($outputs)
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// Data movement operations
+//===----------------------------------------------------------------------===//
+
+def SpatChannelNewOp: SpatOp<"channel_new", []> {
+  let results = (outs
+    SpatChannelType:$new_channel
+  );
+
+  let builders = [
+    OpBuilder<(ins ), [{
+      $_state.addTypes(SpatChannelType());
+    }]>
+  ];
+
+  let assemblyFormat = [{
+    attr-dict
+  }];
+}
+
+def SpatChannelSendOp: SpatOp<"channel_send", []> {
+  let arguments = (ins
+    SpatChannelType: $channel,
+    SpatTensor: $data
+  );
+
+  let assemblyFormat = [{
+    $data `to` $channel attr-dict `:` `(` type($data) `->` type($channel) `)`
+  }];
+}
+
+def SpatChannelReceiveOp: SpatOp<"channel_receive", []> {
+  let arguments = (ins
+    SpatChannelType: $channel
+  );
+
+  let results = (outs
+    SpatTensor: $data
+  );
+
+  let assemblyFormat = [{
+    $channel attr-dict `:` `(` type($channel) `->` type($data) `)`
+  }];
+}
+
+def SpatChannelBroadcastSendOp : SpatOp<"channel_broadcast_send", []> {
+  let arguments = (ins
+    SpatChannelType: $channel,
+    SpatTensor: $data
+  );
+}
+
+def SpatChannelBroadcastReceiveOp : SpatOp<"channel_broadcast_receive", []> {
+  let arguments = (ins
+    SpatChannelType: $channel
+  );
+
+  let results = (outs
+    SpatTensor: $data
+  );
+}
+
+//===----------------------------------------------------------------------===//
+// Math operations
+//===----------------------------------------------------------------------===//
+
+def SpatConstantOp: SpatOp<"constant", []> {
+  let description = [{
+    "Constant value, should be used for weights and biases"
+  }];
+
+  let arguments = (ins
+    AnyAttr: $value,
+    BoolAttr: $shouldAllocate
+  );
+
+  let results = (outs
+    SpatTensor: $out
+  );
+}
+
+def SpatWeightedVMMOp: SpatOp<"Wvmm", []> {
+  let summary = "Vector-matrix-Multiplication within a WeightedCompute operation. The matrix is found in the weights of the WeightedCompute operation, indexed by the weightIndex attribute.";
+
+  let arguments = (ins
+    I32Attr: $weightIndex,
+    SpatTensor:$vector
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  // TODO: Verifier that checks it is within a WeightedCompute operation,
+  // that the weightIndex is valid, and that the matrix is of the right size.
+  let hasVerifier = 1;
+}
+
+def SpatWeightedMVMOp: SpatOp<"Wmvm", []> {
+  let summary = "Matrix-vector multiplication within a WeightedCompute operation. The matrix is found in the weights of the WeightedCompute operation, indexed by the weightIndex attribute.";
+
+  let arguments = (ins
+    I32Attr: $weightIndex,
+    SpatTensor:$vector
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  // TODO: Verifier that checks it is within a WeightedCompute operation,
+  // that the weightIndex is valid, and that the matrix is of the right size.
+  let hasVerifier = 1;
+}
+
+
+def SpatVAddOp: SpatOp<"vadd", []> {
+  let summary = "Element-wise add between tensors a and b. Tensor b must have the same size of tensor b or be a 1x1";
+
+  let arguments = (ins
+    SpatTensor: $a,
+    SpatTensor: $b
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+
+  let assemblyFormat = [{
+    $a `,` $b attr-dict `:` `(` type($a) `,` type($b) `)` `->` type($output)
+  }];
+}
+
+def SpatVMulOp: SpatOp<"vmul", []> {
+  let summary = "Element-wise multiplication between tensors a and b. Tensor b must have the same size of tensor b or be a 1x1";
+
+    let arguments = (ins
+      SpatTensor: $a,
+      SpatTensor: $b
+    );
+
+    let results = (outs
+      SpatTensor:$output
+    );
+
+    //let hasVerifier = 1;
+
+    let assemblyFormat = [{
+      $a `,` $b attr-dict `:` `(` type($a) `,` type($b) `)` `->` type($output)
+    }];
+}
+
+def SpatVDivOp: SpatOp<"vdiv", []> {
+  let summary = "Element-wise division between tensors a and b. Tensor b must have the same size of tensor b or be a 1x1";
+
+  let arguments = (ins
+    SpatTensor:$a,
+    SpatTensor:$b
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  //let hasVerifier = 1;
+
+  let assemblyFormat = [{
+    $a `,` $b attr-dict `:` `(` type($a) `,` type($b) `)` `->` type($output)
+  }];
+}
+
+//TODO: remove
+def SpatVSDivOp: SpatOp<"vsdiv", []> {
+
+  let summary = "Element-wise division between each element of a vector, and a scalar (wrapped in a tensor for convenience)";
+
+  let arguments = (ins
+    SpatTensor:$dividend,
+    SpatTensor:$divisor
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+}
+
+def SpatSumOp: SpatOp<"sum", []> {
+  let summary = "Sum all the elements in the input tensors into a single scalar wrapped in tensor for convenience";
+
+  let arguments = (ins
+    SpatTensor: $input
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+}
+
+def SpatSigmoidOp: SpatOp<"sigmoid", []> {
+  let arguments = (ins
+    SpatTensor:$input
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+}
+
+def SpatReluOp: SpatOp<"relu", []> {
+  let arguments = (ins
+    SpatTensor:$input
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+}
+
+def SpatVMaxOp: SpatOp<"vmax", []> {
+
+  let summary = "Element-wise max function";
+
+  let arguments = (ins
+    SpatTensor: $a,
+    SpatTensor: $b
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+}
+
+def SpatApplyFiltersOp : SpatOp<"apply_filters", []> {
+  let summary = "Apply multiple crossbar weights to a convolutional input tile.";
+  let description = [{
+    Applies a variable number of crossbar weights to a single large image tensor tile,
+    producing a corresponding output tile. This essentially encapsulates a big for loop
+    over all pixels in the input tile, where each pixel is multiplied by all the weights
+    in the operation.
+  }];
+
+  let arguments = (ins
+    I64ArrayAttr: $weightIndices,
+    I64ArrayAttr: $xKernelPositions,
+    I64ArrayAttr: $yKernelPositions,
+    SpatTensor: $input
+  );
+  let results = (outs SpatTensor);
+
+  let assemblyFormat = [{
+    $input attr-dict `:` type($input) `->` type(results)
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// Other operations
+//===----------------------------------------------------------------------===//
+
+def SpatImgConcatOp: SpatOp<"img_concat", []> {
+
+  let summary = "Concatenate pixel tiles into a single image";
+
+  let description = [{
+    Concatenate pixel tiles into a single image:
+    1. First, concatenate the pixel tiles along the "channel" axis (axis 1).
+    2. Next, concatenate the pixel tiles along the "width" axis (axis 2).
+    3. Finally, concatenate the pixel tiles along the "height" axis (axis 3).
+
+    The input tiles should be provided in a specific order:
+      start from the top left pixel,
+      then continue with the pixel on its right,
+      and once you finish the first row of pixels, go to the next row.
+  }];
+
+  let arguments = (ins
+    Variadic<SpatTensor>:$inputs
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+
+  let extraClassDeclaration = [{
+      mlir::Value getInputTile(size_t x, size_t y, size_t tile);
+  }];
+}
+
+#endif // SPATIAL_DIALECT_H