add constant folding and verification pass for pim host operations

better validation scripts output
big refactors

src/PIM/Dialect/Pim/Pim.td

@@ -0,0 +1,391 @@
+#ifndef PIM_DIALECT_H
+#define PIM_DIALECT_H
+
+include "mlir/IR/OpBase.td"
+include "mlir/IR/AttrTypeBase.td"
+include "mlir/Dialect/MemRef/IR/MemRefBase.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/DestinationStyleOpInterface.td"
+include "mlir/Dialect/Bufferization/IR/BufferViewFlowOpInterface.td"
+
+def PimDialect : Dialect {
+  let name = "pim";
+  let summary = "A low-level dialect for the PIM coprocessors on ReRAM crossbars";
+  let cppNamespace = "::onnx_mlir::pim";
+}
+
+class PimOp<string mnemonic, list<Trait> traits = []> :
+    Op<PimDialect, mnemonic, traits>;
+
+def PimTensor :
+  AnyTypeOf<[AnyMemRef, AnyRankedTensor], "", "::mlir::ShapedType">;
+
+// Communication
+
+def PimSendOp: PimOp<"send", []> {
+  let arguments = (ins
+    PimTensor: $src,
+    I32Attr: $size,
+    I32Attr: $targetCoreId
+  );
+
+  let assemblyFormat = [{
+    `(` $src `)` attr-dict `:` type($src) `->` `(` `)`
+  }];
+}
+
+def PimReceiveOp: PimOp<"receive", [DestinationStyleOpInterface]> {
+  let arguments = (ins
+    PimTensor: $dst,
+    I32Attr: $size,
+    I32Attr: $srcCoreId
+  );
+
+  let results = (outs
+    PimTensor: $out
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getDstMutable();
+    }
+  }];
+
+  let assemblyFormat = [{
+    `(` $dst `)` attr-dict `:` type($dst) `->` type($out)
+  }];
+}
+
+// Core
+
+def PimCoreOp: PimOp<"core", [SingleBlock]> {
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let arguments = (ins
+    Variadic<PimTensor>:$weights,
+    I32Attr: $coreId
+  );
+
+  let assemblyFormat = [{
+    `(` $weights `)` attr-dict regions `:` type($weights) `->` `(` `)`
+  }];
+}
+
+// Memory
+
+def PimConstantOp: PimOp<"constant", []> {
+  let description = [{
+    Allocate a constant value in global memory
+  }];
+
+  let arguments = (ins
+    AnyAttr: $value,
+    BoolAttr: $shouldAllocate
+  );
+
+  let results = (outs
+    PimTensor: $out
+  );
+}
+
+def PimMemCopyHostToDevOp: PimOp<"memcp_hd", [DestinationStyleOpInterface]> {
+  let description = [{
+    Copy a memory region from host memory into device memory
+  }];
+
+  let arguments = (ins
+    PimTensor: $deviceDst,
+    PimTensor: $hostSrc,
+    I32Attr: $deviceDstOffset,
+    I32Attr: $hostSrcOffset,
+    I32Attr: $size
+  );
+
+  let results = (outs
+    PimTensor: $deviceDstOut
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getDeviceDstMutable();
+    }
+  }];
+
+  let assemblyFormat = [{
+    `(` $deviceDst `,` $hostSrc `)` attr-dict `:` `(` type($deviceDst) `,` type($hostSrc) `)` `->` type($deviceDstOut)
+  }];
+}
+
+def PimMemCopyDevToHostOp: PimOp<"memcp_dh", [DestinationStyleOpInterface]> {
+  let description = [{
+    Copy a memory region from device memory into host memory
+  }];
+
+  let arguments = (ins
+    PimTensor: $hostDst,
+    PimTensor: $deviceSrc,
+    I32Attr: $hostDstOffset,
+    I32Attr: $deviceSrcOffset,
+    I32Attr: $size
+  );
+
+  let results = (outs
+    PimTensor: $hostDstOut
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getHostDstMutable();
+    }
+  }];
+
+
+  let assemblyFormat = [{
+    `(` $hostDst `,` $deviceSrc `)` attr-dict `:` `(` type($hostDst) `,` type($deviceSrc) `)` `->` type($hostDstOut)
+  }];
+}
+
+def PimMemCopyOp: PimOp<"memcp", [DestinationStyleOpInterface]> {
+  let description = [{
+    Copy a memory region from and to the same memory
+  }];
+
+  let arguments = (ins
+    PimTensor: $dst,
+    PimTensor: $src,
+    I32Attr: $dstOffset,
+    I32Attr: $srcOffset,
+    I32Attr: $size
+  );
+
+  let results = (outs
+    PimTensor: $dstOut
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getDstMutable();
+    }
+  }];
+
+
+  let assemblyFormat = [{
+    `(` $dst `,` $src `)` attr-dict `:` `(` type($dst) `,` type($src) `)` `->` type($dstOut)
+  }];
+}
+
+// Algebra
+
+def PimTransposeOp: PimOp<"transpose", [DestinationStyleOpInterface]> {
+  let description = [{
+    Matrix transpose
+  }];
+
+  let arguments = (ins
+    PimTensor: $data,
+    I64ArrayAttr: $perms,
+    PimTensor: $outBuf
+  );
+
+  let results = (outs
+    PimTensor: $outRes
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutBufMutable();
+    }
+  }];
+
+  let assemblyFormat = [{
+    `(` $data `,` $outBuf `)` attr-dict `:` `(` type($data) `,` type($outBuf) `)` `->` type($outRes)
+  }];
+}
+
+def PimVMMOp: PimOp<"vmm", [DestinationStyleOpInterface]> {
+  let description = [{
+    Vector-matrix multiplication: c = a * b
+  }];
+
+  let arguments = (ins
+    I32Attr: $weightIndex,
+    PimTensor: $vectorInput,
+    PimTensor: $outBuf
+  );
+
+  let results = (outs
+    PimTensor: $outRes
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutBufMutable();
+    }
+  }];
+
+  let assemblyFormat = [{
+    `(` $vectorInput `,` $outBuf `)` attr-dict `:` `(` type($vectorInput) `,` type($outBuf) `)` `->` type($outRes)
+  }];
+}
+
+def PimMVMOp: PimOp<"mvm", [DestinationStyleOpInterface]> {
+  let description = [{
+    Matrix-vector multiplication: c = a * b
+  }];
+
+  let arguments = (ins
+    I32Attr: $weightIndex,
+    PimTensor: $vectorInput,
+    PimTensor: $outBuf
+  );
+
+  let results = (outs
+    PimTensor: $outRes
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutBufMutable();
+    }
+  }];
+}
+
+def PimVAddOp: PimOp<"vadd", [DestinationStyleOpInterface]> {
+  let description = [{
+    Element-wise addition: c = a + b
+  }];
+
+  let arguments = (ins
+    PimTensor: $a,
+    PimTensor: $b,
+    PimTensor: $outBuf
+  );
+
+  let results = (outs
+    PimTensor: $outRes
+  );
+
+  let extraClassDeclaration = [{
+    mlir::MutableOperandRange getDpsInitsMutable() {
+      return getOutBufMutable();
+    }
+  }];
+
+  let assemblyFormat = [{
+    `(` $a `,` $b `,` $outBuf `)` attr-dict `:` `(` type($a) `,` type($b) `,` type($outBuf) `)` `->` type($outRes)
+  }];
+}
+
+def PimVMaxOp: PimOp<"vmax", [DeclareOpInterfaceMethods<BufferViewFlowOpInterface>]> {
+  let description = [{
+    Element-wise max: c = max(a, b)
+  }];
+
+  let arguments = (ins
+    PimTensor: $a,
+    PimTensor: $b,
+    PimTensor: $outBuf
+  );
+
+  let results = (outs
+    PimTensor: $outRes
+  );
+}
+
+def PimApplyFiltersOp: PimOp<"apply_filters", [DeclareOpInterfaceMethods<BufferViewFlowOpInterface>]> {
+  let description = [{
+    Apply filters to a tensor
+  }];
+
+  let arguments = (ins
+    I64ArrayAttr: $weightIndices,
+    I64ArrayAttr: $xKernelPositions,
+    I64ArrayAttr: $yKernelPositions,
+    PimTensor: $input,
+    PimTensor: $outBuf,
+    PimTensor: $accumBuf
+  );
+
+  let results = (outs
+    PimTensor: $outRes
+  );
+
+  let assemblyFormat = [{
+    `(` `input` `=` $input `,` `outBuf` `=` $outBuf `,` `accumBuf` `=` $accumBuf `)` attr-dict `:` 
+    type($input) `,` type($outBuf) `,` type($accumBuf) `->` type($outRes)
+  }];
+}
+
+def PimSumOp: PimOp<"sum", [DeclareOpInterfaceMethods<BufferViewFlowOpInterface>]> {
+  let description = [{
+    Sum all elements into a single one
+  }];
+
+  let arguments = (ins
+    PimTensor: $a,
+    PimTensor: $outBuf
+  );
+
+  let results = (outs
+    PimTensor: $outRes
+  );
+}
+
+def PimVSDivOp: PimOp<"vsdiv", [DeclareOpInterfaceMethods<BufferViewFlowOpInterface>]> {
+  let description = [{
+    Element-wise division between each element of a vector, and a scalar (wrapped in a tensor for convenience)
+  }];
+
+  let arguments = (ins
+    PimTensor: $dividend,
+    PimTensor: $divisor,
+    PimTensor: $outBuf
+  );
+
+  let results = (outs
+    PimTensor: $outRes
+  );
+}
+
+def PimVReluOp: PimOp<"vrelu", [DeclareOpInterfaceMethods<BufferViewFlowOpInterface>]> {
+  let description = [{
+    Element-wise ReLU: c = max(a, 0)
+  }];
+
+  let arguments = (ins
+    PimTensor: $a,
+    PimTensor: $outBuf
+  );
+
+  let results = (outs
+    PimTensor: $outRes
+  );
+}
+
+def PimVExpOp: PimOp<"vexp", [DeclareOpInterfaceMethods<BufferViewFlowOpInterface>]> {
+  let description = [{
+    Element-wise exp: c = exp(a)
+  }];
+
+  let arguments = (ins
+    PimTensor: $a,
+    PimTensor: $outBuf
+  );
+
+  let results = (outs
+    PimTensor: $outRes
+  );
+}
+
+def PimHaltOp: PimOp<"halt", [Terminator]> {
+  let description = [{
+    Halts the execution of the core
+  }];
+
+  let assemblyFormat = [{
+    attr-dict
+  }];
+}
+
+#endif // PIM_DIALECT_H