#ifndef PIM_DIALECT_H
#define PIM_DIALECT_H

include "mlir/IR/OpBase.td"
include "mlir/IR/OpAsmInterface.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,
    DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
  let summary = "Execute a block on a PIM core";

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

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

  let extraClassDeclaration = [{
    ::mlir::BlockArgument getWeightArgument(unsigned idx);
  }];

  let hasVerifier = 1;
  let hasCustomAssemblyFormat = 1;
}

def PimCoreBatchOp : PimOp<"core_batch", [SingleBlock, AttrSizedOperandSegments,
    DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
  let summary = "Execute equivalent batched core bodies";

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

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

  let extraClassDeclaration = [{
    ::mlir::BlockArgument getLaneArgument();
    ::mlir::BlockArgument getWeightArgument(unsigned idx);
    ::mlir::BlockArgument getInputArgument(unsigned idx);
  }];

  let hasVerifier = 1;
  let hasCustomAssemblyFormat = 1;
}

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

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

def PimYieldOp : PimOp<"yield", [Terminator]> {
  let summary = "Yield results from a Pim region";

  let arguments = (ins
    Variadic<PimTensor>:$outputs
  );

  let hasCustomAssemblyFormat = 1;
}

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

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

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

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

def PimSendTensorOp : PimOp<"send_tensor", []> {
  let summary = "Send equal contiguous chunks of one tensor to target cores";

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

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

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

  let hasCustomAssemblyFormat = 1;
}

def PimSendTensorBatchOp : PimOp<"send_tensor_batch", []> {
  let summary = "Send equal contiguous chunks of one per-lane tensor from a batched core";

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

  let hasVerifier = 1;
  let hasCustomAssemblyFormat = 1;
}

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

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

  let results = (outs
    PimTensor:$output
  );

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

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

def PimReceiveTensorOp : PimOp<"receive_tensor", [DestinationStyleOpInterface]> {
  let summary = "Receive equal contiguous chunks from source cores into one tensor";

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

  let results = (outs
    PimTensor:$output
  );

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

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

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

  let results = (outs
    PimTensor:$output
  );

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

  let hasCustomAssemblyFormat = 1;
}

def PimReceiveTensorBatchOp : PimOp<"receive_tensor_batch", [DestinationStyleOpInterface]> {
  let summary = "Receive equal contiguous chunks into one per-lane tensor inside a batched core";

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

  let results = (outs
    PimTensor:$output
  );

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

  let hasVerifier = 1;
  let hasCustomAssemblyFormat = 1;
}

def PimMemCopyHostToDevOp : PimOp<"memcp_hd", [DestinationStyleOpInterface]> {
  let summary = "Copy a memory region from host memory into device memory";

  let arguments = (ins
    Index:$deviceTargetOffset,
    Index:$hostSourceOffset,
    PimTensor:$deviceTarget,
    PimTensor:$hostSource,
    I32Attr:$size
  );

  let results = (outs
    PimTensor:$output
  );

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

  let assemblyFormat = [{
    `[` $deviceTargetOffset `,` $hostSourceOffset `]`
    `(` $deviceTarget `,` $hostSource `)` attr-dict
    `:` type($deviceTarget) `,` type($hostSource) `->` type($output)
  }];
}

def PimMemCopyHostToDevBatchOp : PimOp<"memcp_hd_batch", [DestinationStyleOpInterface]> {
  let summary = "Copy a per-lane tensor from host memory into device memory inside a batched core";

  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
    Index:$hostTargetOffset,
    Index:$deviceSourceOffset,
    PimTensor:$hostTarget,
    PimTensor:$deviceSource,
    I32Attr:$size
  );

  let results = (outs
    PimTensor:$output
  );

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

  let assemblyFormat = [{
    `[` $hostTargetOffset `,` $deviceSourceOffset `]`
    `(` $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)
  }];
}

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
//===----------------------------------------------------------------------===//

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
    PimTensor:$weight,
    PimTensor:$input,
    PimTensor:$outputBuffer
  );

  let results = (outs
    PimTensor:$output
  );

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

  let hasVerifier = 1;
  let assemblyFormat = [{
    `[` $weight `]` `(` $input `,` $outputBuffer `)` attr-dict `:` `(` type($weight) `,` 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 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)
  }];
}

def PimVSoftmaxOp : PimOp<"vsoftmax", [DestinationStyleOpInterface]> {
  let summary = "Softmax over the full input vector";

  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