#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">;

//===----------------------------------------------------------------------===//
// Execution
//===----------------------------------------------------------------------===//

def PimCoreOp : PimOp<"core", [SingleBlock]> {
  let summary = "Execute a block on a PIM core";

  let regions = (region SizedRegion<1>:$body);

  let arguments = (ins
    Variadic<PimTensor>:$weights,
    I32Attr:$coreId
  );

  let assemblyFormat = [{
    `(` $weights `)` attr-dict regions `:` type($weights) `->` `(` `)`
  }];
}

def PimHaltOp : PimOp<"halt", [Terminator]> {
  let summary = "Halt execution of the core";

  let assemblyFormat = [{
    attr-dict
  }];
}

//===----------------------------------------------------------------------===//
// Communication
//===----------------------------------------------------------------------===//

def PimSendOp : PimOp<"send", []> {
  let summary = "Send a tensor to another core";

  let arguments = (ins
    PimTensor:$input,
    I32Attr:$size,
    I32Attr:$targetCoreId
  );

  let assemblyFormat = [{
    `(` $input `)` attr-dict `:` type($input) `->` `(` `)`
  }];
}

def PimReceiveOp : PimOp<"receive", [DestinationStyleOpInterface]> {
  let summary = "Receive a tensor from another core";

  let arguments = (ins
    PimTensor:$outputBuffer,
    I32Attr:$size,
    I32Attr:$sourceCoreId
  );

  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";

  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";

  let arguments = (ins
    PimTensor:$hostTarget,
    PimTensor:$deviceSource,
    I32Attr:$hostTargetOffset,
    I32Attr:$deviceSourceOffset,
    I32Attr:$size
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getHostTargetMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $hostTarget `,` $deviceSource `)` attr-dict `:` `(` type($hostTarget) `,` type($deviceSource) `)` `->` type($output)
  }];
}

def PimMemCopyOp : PimOp<"memcp", [DestinationStyleOpInterface]> {
  let summary = "Copy a memory region within the same memory space";

  let arguments = (ins
    PimTensor:$target,
    PimTensor:$source,
    I32Attr:$targetOffset,
    I32Attr:$sourceOffset,
    I32Attr:$size
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getTargetMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $target `,` $source `)` attr-dict `:` `(` type($target) `,` type($source) `)` `->` type($output)
  }];
}

//===----------------------------------------------------------------------===//
// Math
//===----------------------------------------------------------------------===//

def PimTransposeOp : PimOp<"transpose", [DestinationStyleOpInterface]> {
  let summary = "Transpose a matrix";

  let arguments = (ins
    PimTensor:$input,
    I64ArrayAttr:$permutation,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getOutputBufferMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $input `,` $outputBuffer `)` attr-dict `:` `(` type($input) `,` type($outputBuffer) `)` `->` type($output)
  }];
}

def PimVMMOp : PimOp<"vmm", [DestinationStyleOpInterface]> {
  let summary = "Vector-matrix multiplication: c = a * b";

  let arguments = (ins
    I32Attr:$weightIndex,
    PimTensor:$input,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getOutputBufferMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $input `,` $outputBuffer `)` attr-dict `:` `(` type($input) `,` type($outputBuffer) `)` `->` type($output)
  }];
}

def PimMVMOp : PimOp<"mvm", [DestinationStyleOpInterface]> {
  let summary = "Matrix-vector multiplication: c = a * b";

  let arguments = (ins
    I32Attr:$weightIndex,
    PimTensor:$input,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getOutputBufferMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $input `,` $outputBuffer `)` attr-dict `:` `(` type($input) `,` type($outputBuffer) `)` `->` type($output)
  }];
}

def PimVVAddOp : PimOp<"vvadd", [DestinationStyleOpInterface]> {
  let summary = "Element-wise addition: c = a + b";

  let arguments = (ins
    PimTensor:$lhs,
    PimTensor:$rhs,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getOutputBufferMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $lhs `,` $rhs `,` $outputBuffer `)` attr-dict `:` `(` type($lhs) `,` type($rhs) `,` type($outputBuffer) `)` `->` type($output)
  }];
}

def PimVVSubOp : PimOp<"vvsub", [DestinationStyleOpInterface]> {
  let summary = "Element-wise subtraction: c = a - b";

  let arguments = (ins
    PimTensor:$lhs,
    PimTensor:$rhs,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getOutputBufferMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $lhs `,` $rhs `,` $outputBuffer `)` attr-dict `:` `(` type($lhs) `,` type($rhs) `,` type($outputBuffer) `)` `->` type($output)
  }];
}

def PimVVMulOp : PimOp<"vvmul", [DestinationStyleOpInterface]> {
  let summary = "Element-wise multiplication: c = a * b";

  let arguments = (ins
    PimTensor:$lhs,
    PimTensor:$rhs,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getOutputBufferMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $lhs `,` $rhs `,` $outputBuffer `)` attr-dict `:` `(` type($lhs) `,` type($rhs) `,` type($outputBuffer) `)` `->` type($output)
  }];
}

def PimVVMaxOp : PimOp<"vvmax", [DestinationStyleOpInterface]> {
  let summary = "Element-wise max: c = max(a, b)";

  let arguments = (ins
    PimTensor:$lhs,
    PimTensor:$rhs,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getOutputBufferMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $lhs `,` $rhs `,` $outputBuffer `)` attr-dict `:` `(` type($lhs) `,` type($rhs) `,` type($outputBuffer) `)` `->` type($output)
  }];
}

def PimVVDMulOp : PimOp<"vvdmul", [DestinationStyleOpInterface]> {
  let summary = "Dot product: c = dot(a, b)";

  let arguments = (ins
    PimTensor:$lhs,
    PimTensor:$rhs,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getOutputBufferMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $lhs `,` $rhs `,` $outputBuffer `)` attr-dict `:` `(` type($lhs) `,` type($rhs) `,` type($outputBuffer) `)` `->` type($output)
  }];
}

def PimSumOp : PimOp<"sum", [DestinationStyleOpInterface]> {
  let summary = "Reduce all elements to a single value";

  let arguments = (ins
    PimTensor:$input,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getOutputBufferMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $input `,` $outputBuffer `)` attr-dict `:` `(` type($input) `,` type($outputBuffer) `)` `->` type($output)
  }];
}

def PimVAvgOp : PimOp<"vavg", [DestinationStyleOpInterface]> {
  let summary = "Average all elements into a single value";

  let arguments = (ins
    PimTensor:$input,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getOutputBufferMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $input `,` $outputBuffer `)` attr-dict `:` `(` type($input) `,` type($outputBuffer) `)` `->` type($output)
  }];
}

def PimVReluOp : PimOp<"vrelu", [DestinationStyleOpInterface]> {
  let summary = "Element-wise ReLU: c = max(a, 0)";

  let arguments = (ins
    PimTensor:$input,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getOutputBufferMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $input `,` $outputBuffer `)` attr-dict `:` `(` type($input) `,` type($outputBuffer) `)` `->` type($output)
  }];
}

def PimVTanhOp : PimOp<"vtanh", [DestinationStyleOpInterface]> {
  let summary = "Element-wise tanh activation";

  let arguments = (ins
    PimTensor:$input,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getOutputBufferMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $input `,` $outputBuffer `)` attr-dict `:` `(` type($input) `,` type($outputBuffer) `)` `->` type($output)
  }];
}

def PimVSigmOp : PimOp<"vsigm", [DestinationStyleOpInterface]> {
  let summary = "Element-wise sigmoid activation";

  let arguments = (ins
    PimTensor:$input,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

  let extraClassDeclaration = [{
    mlir::MutableOperandRange getDpsInitsMutable() {
      return getOutputBufferMutable();
    }
  }];

  let assemblyFormat = [{
    `(` $input `,` $outputBuffer `)` attr-dict `:` `(` type($input) `,` type($outputBuffer) `)` `->` type($output)
  }];
}

#endif // PIM_DIALECT_H