#ifndef SPATIAL_DIALECT_H
#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/EnumAttr.td"
include "mlir/IR/RegionKindInterface.td"
include "mlir/Interfaces/ControlFlowInterfaces.td"
include "mlir/Interfaces/ParallelCombiningOpInterface.td"
include "mlir/Interfaces/SideEffectInterfaces.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 useDefaultAttributePrinterParser = 0;
  let extraClassDeclaration = [{
    ::mlir::Attribute parseAttribute(::mlir::DialectAsmParser &parser,
                                     ::mlir::Type type) const override;
    void printAttribute(::mlir::Attribute attr,
                        ::mlir::DialectAsmPrinter &printer) const override;
  }];
}

def SpatialLayoutCapabilityInterface : OpInterface<"SpatialLayoutCapabilityInterface"> {
  let description = [{
    Contract implemented by logical Spatial planning operations that expose
    their legal physical layout alternatives to the Spatial planner.
  }];

  let methods = [
    InterfaceMethod<
      "Return legal physical layout alternatives for this operation and its current operand layouts.",
      "::llvm::SmallVector<::onnx_mlir::spatial::LayoutAlternative>",
      "getLayoutAlternatives",
      (ins "const ::onnx_mlir::spatial::SpatialTargetResources &":$target,
           "::llvm::ArrayRef<::onnx_mlir::spatial::PhysicalLayout>":$operandLayouts)>
  ];

  let cppNamespace = "::onnx_mlir::spatial";
}

def SpatLogicalLayoutNCHW : I32EnumAttrCase<"NCHW", 0, "nchw">;
def SpatLogicalLayout : I32EnumAttr<"LogicalLayout", "Logical tensor layout", [
  SpatLogicalLayoutNCHW
]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::onnx_mlir::spatial";
}

def SpatLogicalLayoutAttr : EnumAttr<SpatialDialect, SpatLogicalLayout, "logical_layout"> {
  let assemblyFormat = "$value";
}

def SpatPhysicalLayoutDenseNCHW : I32EnumAttrCase<"DenseNCHW", 0, "dense_nchw">;
def SpatPhysicalLayoutNCHWRowStrip : I32EnumAttrCase<"NCHWRowStrip", 1, "nchw_row_strip">;
def SpatPhysicalLayoutNHWCRowStrip : I32EnumAttrCase<"NHWCRowStrip", 2, "nhwc_row_strip">;
def SpatPhysicalLayoutFragmented : I32EnumAttrCase<"Fragmented", 3, "fragmented">;
def SpatPhysicalLayout : I32EnumAttr<"PhysicalLayout", "Physical tensor layout", [
  SpatPhysicalLayoutDenseNCHW,
  SpatPhysicalLayoutNCHWRowStrip,
  SpatPhysicalLayoutNHWCRowStrip,
  SpatPhysicalLayoutFragmented
]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::onnx_mlir::spatial";
}

def SpatPhysicalLayoutAttr : EnumAttr<SpatialDialect, SpatPhysicalLayout, "physical_layout"> {
  let assemblyFormat = "$value";
}

def SpatBlueprintModePhysicalView : I32EnumAttrCase<"PhysicalView", 0, "physical_view">;
def SpatBlueprintModeFragmentAssembly : I32EnumAttrCase<"FragmentAssembly", 1, "fragment_assembly">;
def SpatBlueprintMode : I32EnumAttr<"BlueprintMode", "Blueprint reconstruction mode", [
  SpatBlueprintModePhysicalView,
  SpatBlueprintModeFragmentAssembly
]> {
  let genSpecializedAttr = 0;
  let cppNamespace = "::onnx_mlir::spatial";
}

def SpatBlueprintModeAttr : EnumAttr<SpatialDialect, SpatBlueprintMode, "blueprint_mode"> {
  let assemblyFormat = "$value";
}

class SpatOp<string mnemonic, list<Trait> traits = []> :
    Op<SpatialDialect, mnemonic, traits>;

class SpatLayoutPlanOp<string mnemonic> : SpatOp<mnemonic,
    [SpatialLayoutCapabilityInterface,
     DeclareOpInterfaceMethods<SpatialLayoutCapabilityInterface>]>;

// TODO maybe remove and use AnyRankedTensor directly
def SpatTensor :
  AnyTypeOf<[AnyMemRef, AnyRankedTensor], "", "::mlir::ShapedType">;

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

class SpatComputeLikeBase<string mnemonic> : SpatOp<mnemonic,
    [AttrSizedOperandSegments,
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
  let summary = "Compute region with attached constant weights";

  let arguments = (ins
    Variadic<SpatTensor>:$weights,
    Variadic<SpatTensor>:$inputs
  );

  let results = (outs
    Variadic<SpatTensor>:$outputs
  );

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

  let hasVerifier = 1;
  let hasFolder = 1;
  let hasCustomAssemblyFormat = 1;
}

def SpatGraphCompute : SpatComputeLikeBase<"graph_compute"> {
  let hasCanonicalizer = 1;
  let extraClassDeclaration = [{
    std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
    std::optional<::mlir::BlockArgument> getInputArgument(unsigned idx);
    std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
      insertWeight(unsigned idx, ::mlir::Value weight, ::mlir::Location loc);
    std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
      insertInput(unsigned idx, ::mlir::Value input, ::mlir::Location loc);
    ::llvm::SetVector<::mlir::Value, ::llvm::SmallVector<::mlir::Value, 4>, ::llvm::SmallDenseSet<::mlir::Value, 4>> getCrossbarWeights();
    ::mlir::FailureOr<std::tuple<::mlir::OpResult, SpatGraphCompute>>
      insertOutput(::mlir::RewriterBase &rewriter, unsigned idx, ::mlir::Type type, ::mlir::Location loc);
  }];
}

def SpatScheduledCompute : SpatComputeLikeBase<"scheduled_compute"> {
  let extraClassDeclaration = [{
    std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
    std::optional<::mlir::BlockArgument> getInputArgument(unsigned idx);
    std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
      insertWeight(unsigned idx, ::mlir::Value weight, ::mlir::Location loc);
    std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
      insertInput(unsigned idx, ::mlir::Value input, ::mlir::Location loc);
    ::llvm::SetVector<::mlir::Value, ::llvm::SmallVector<::mlir::Value, 4>, ::llvm::SmallDenseSet<::mlir::Value, 4>> getCrossbarWeights();
    ::mlir::FailureOr<std::tuple<::mlir::OpResult, SpatScheduledCompute>>
      insertOutput(::mlir::RewriterBase &rewriter, unsigned idx, ::mlir::Type type, ::mlir::Location loc);
  }];
}

class SpatComputeBatchLikeBase<string mnemonic> : SpatOp<mnemonic,
    [AttrSizedOperandSegments,
     DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
  let summary = "Tensor-native batch of equivalent compute lanes with shared weights and packed inputs";

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

  let results = (outs
    Variadic<SpatTensor>:$outputs
  );

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

  let hasVerifier = 1;
  let hasCustomAssemblyFormat = 1;
}

def SpatGraphComputeBatch : SpatComputeBatchLikeBase<"graph_compute_batch"> {
  let hasCanonicalizer = 1;
  let extraClassDeclaration = [{
    std::optional<::mlir::BlockArgument> getLaneArgument();
    std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
    std::optional<::mlir::BlockArgument> getInputArgument(unsigned idx);
    std::optional<::mlir::BlockArgument> getOutputArgument(unsigned idx);
    std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
      insertWeight(unsigned idx, ::mlir::Value weight, ::mlir::Location loc);
    std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
      insertInput(unsigned idx, ::mlir::Value input, ::mlir::Location loc);
    ::llvm::SetVector<::mlir::Value, ::llvm::SmallVector<::mlir::Value, 4>, ::llvm::SmallDenseSet<::mlir::Value, 4>> getCrossbarWeights();
    ::mlir::FailureOr<std::tuple<::mlir::OpResult, ::mlir::BlockArgument, SpatGraphComputeBatch>>
      insertOutput(::mlir::RewriterBase &rewriter, unsigned idx, ::mlir::Type type, ::mlir::Location loc);
  }];
}

def SpatScheduledComputeBatch : SpatComputeBatchLikeBase<"scheduled_compute_batch"> {
  let hasCanonicalizer = 1;
  let extraClassDeclaration = [{
    std::optional<::mlir::BlockArgument> getLaneArgument();
    std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
    std::optional<::mlir::BlockArgument> getInputArgument(unsigned idx);
    std::optional<::mlir::BlockArgument> getOutputArgument(unsigned idx);
    std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
      insertWeight(unsigned idx, ::mlir::Value weight, ::mlir::Location loc);
    std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
      insertInput(unsigned idx, ::mlir::Value input, ::mlir::Location loc);
    ::llvm::SetVector<::mlir::Value, ::llvm::SmallVector<::mlir::Value, 4>, ::llvm::SmallDenseSet<::mlir::Value, 4>> getCrossbarWeights();
    ::mlir::FailureOr<std::tuple<::mlir::OpResult, ::mlir::BlockArgument, SpatScheduledComputeBatch>>
      insertOutput(::mlir::RewriterBase &rewriter, unsigned idx, ::mlir::Type type, ::mlir::Location loc);
  }];
}

def SpatInParallelOp : SpatOp<"in_parallel", [
       Pure,
       Terminator,
       DeclareOpInterfaceMethods<InParallelOpInterface>,
      ] # GraphRegionNoTerminator.traits> {
  let summary = "Parallel combining terminator for resultful Spatial compute batches";

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

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

  let hasCustomAssemblyFormat = 1;
}

def SpatBlockYieldOp : SpatOp<"block_yield", [
    Terminator,
    DeclareOpInterfaceMethods<BranchOpInterface, ["getSuccessorForOperands"]>
  ]> {
  let summary = "Terminate a scheduled structural compute block";

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

  let successors = (successor
    VariadicSuccessor<AnySuccessor>:$next
  );

  let hasVerifier = 1;
  let hasCustomAssemblyFormat = 1;
}

def SpatDeferredCommunicationOp : SpatOp<"deferred_communication", [SingleBlock]> {
  let summary = "Temporary scheduled payload derivation placeholder";

  let arguments = (ins
    Variadic<SpatTensor>:$sources,
    OptionalAttr<I64Attr>:$specialization_count
  );

  let results = (outs
    SpatTensor:$output
  );

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

  let hasVerifier = 1;
  let hasCustomAssemblyFormat = 1;
}

def SpatDeferredSourceSelectOp : SpatOp<"deferred_source_select", []> {
  let summary = "Select a deferred tensor source with a statically analyzable index";

  let arguments = (ins
    Index:$selector,
    Variadic<SpatTensor>:$sources
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
  let hasCustomAssemblyFormat = 1;
}

def SpatExtractRowsOp : SpatOp<"extract_rows", []> {
  let summary = "Extract every row of a rank-2 tensor as separate rank-2 row tensors";

  let arguments = (ins
    SpatTensor:$input
  );

  let results = (outs
    Variadic<SpatTensor>:$outputs
  );

  let hasVerifier = 1;
  let hasCustomAssemblyFormat = 1;
}

def SpatConcatOp : SpatOp<"concat", []> {
  let summary = "Concatenate tensors with compact Spatial operand syntax";

  let arguments = (ins
    I64Attr:$axis,
    Variadic<SpatTensor>:$inputs
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
  let hasCustomAssemblyFormat = 1;
}

//===----------------------------------------------------------------------===//
// Planning
//===----------------------------------------------------------------------===//

def SpatConv2DPlanOp : SpatLayoutPlanOp<"conv2d_plan"> {
  let summary = "Structured Conv2D planning op that preserves logical ONNX geometry";

  let arguments = (ins
    SpatTensor:$input,
    SpatTensor:$weight,
    Optional<SpatTensor>:$bias,
    DenseI64ArrayAttr:$pads,
    DenseI64ArrayAttr:$strides,
    DenseI64ArrayAttr:$dilations,
    I64Attr:$group,
    SpatLogicalLayoutAttr:$logicalLayout
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
}

def SpatFlattenPlanOp : SpatLayoutPlanOp<"flatten_plan"> {
  let summary = "Layout-aware static Flatten planning op";

  let arguments = (ins
    SpatTensor:$input,
    I64Attr:$axis,
    SpatLogicalLayoutAttr:$logicalLayout
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
}

def SpatReluPlanOp : SpatLayoutPlanOp<"relu_plan"> {
  let summary = "Layout-aware ReLU planning op";

  let arguments = (ins
    SpatTensor:$input,
    SpatLogicalLayoutAttr:$logicalLayout
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
}

def SpatSiluPlanOp : SpatLayoutPlanOp<"silu_plan"> {
  let summary = "Layout-aware SiLU planning op";

  let arguments = (ins
    SpatTensor:$input,
    SpatLogicalLayoutAttr:$logicalLayout
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
}

def SpatResizeNearestPlanOp : SpatLayoutPlanOp<"resize_nearest_plan"> {
  let summary = "Layout-aware nearest asymmetric Resize planning op";

  let arguments = (ins
    SpatTensor:$input,
    SpatLogicalLayoutAttr:$logicalLayout
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
}

def SpatMaxPool2DPlanOp : SpatLayoutPlanOp<"max_pool2d_plan"> {
  let summary = "Layout-aware 2D NCHW MaxPool planning op";

  let arguments = (ins
    SpatTensor:$input,
    DenseI64ArrayAttr:$kernelShape,
    DenseI64ArrayAttr:$pads,
    DenseI64ArrayAttr:$strides,
    DenseI64ArrayAttr:$dilations,
    SpatLogicalLayoutAttr:$logicalLayout
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
}

def SpatGlobalAveragePoolPlanOp : SpatLayoutPlanOp<"global_average_pool_plan"> {
  let summary = "Layout-aware NCHW global average-pool planning op";

  let arguments = (ins
    SpatTensor:$input,
    SpatLogicalLayoutAttr:$logicalLayout
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
}

def SpatBiasAddPlanOp : SpatLayoutPlanOp<"bias_add_plan"> {
  let summary = "Layout-aware Conv-style bias add planning op";

  let arguments = (ins
    SpatTensor:$input,
    SpatTensor:$bias,
    SpatLogicalLayoutAttr:$logicalLayout
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
}

def SpatAddPlanOp : SpatLayoutPlanOp<"add_plan"> {
  let summary = "Layout-aware elementwise add planning op";

  let arguments = (ins
    SpatTensor:$lhs,
    SpatTensor:$rhs,
    SpatLogicalLayoutAttr:$logicalLayout
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
}

def SpatConcatPlanOp : SpatLayoutPlanOp<"concat_plan"> {
  let summary = "Layout-aware tensor concatenation planning op";

  let arguments = (ins
    Variadic<SpatTensor>:$inputs,
    I64Attr:$axis,
    SpatLogicalLayoutAttr:$logicalLayout
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
}

def SpatBlueprintOp : SpatOp<"blueprint", []> {
  let summary = "Blueprint for assembling logical tensors from published fragments";

  let arguments = (ins
    SpatTensor:$input,
    Variadic<SpatTensor>:$fragments,
    SpatLogicalLayoutAttr:$logicalLayout,
    SpatPhysicalLayoutAttr:$physicalLayout,
    DenseI64ArrayAttr:$fragmentOffsets,
    DenseI64ArrayAttr:$fragmentSizes,
    StrAttr:$indexMap,
    OptionalAttr<SpatBlueprintModeAttr>:$mode,
    OptionalAttr<DenseI64ArrayAttr>:$fragmentOperandIndices,
    OptionalAttr<DenseI64ArrayAttr>:$fragmentSourceSlots,
    OptionalAttr<DenseI64ArrayAttr>:$fragmentSourceOffsets,
    OptionalAttr<DenseI64ArrayAttr>:$fragmentStrides,
    OptionalAttr<StrAttr>:$conflictPolicy,
    OptionalAttr<StrAttr>:$coveragePolicy
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
  let hasCustomAssemblyFormat = 1;
}

def SpatMaterializeLayoutOp : SpatOp<"materialize_layout", []> {
  let summary = "Explicit layout conversion or materialization barrier";

  let arguments = (ins
    SpatTensor:$input,
    SpatLogicalLayoutAttr:$logicalLayout,
    SpatPhysicalLayoutAttr:$sourcePhysicalLayout,
    SpatPhysicalLayoutAttr:$targetPhysicalLayout
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;
}

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

def SpatChannelSendOp : SpatOp<"channel_send", []> {
  let summary = "Send a tensor through a logical channel";

  let arguments = (ins
    Index:$channelId,
    Index:$sourceCoreId,
    Index:$targetCoreId,
    SpatTensor:$input
  );

  let assemblyFormat = [{
    $input `channel` $channelId `from` $sourceCoreId `to` $targetCoreId
      attr-dict `:` type($input)
  }];
}

def SpatChannelReceiveOp : SpatOp<"channel_receive", []> {
  let summary = "Receive a tensor from a logical channel";

  let arguments = (ins
    Index:$channelId,
    Index:$sourceCoreId,
    Index:$targetCoreId
  );

  let results = (outs
    SpatTensor:$output
  );

  let assemblyFormat = [{
    `channel` $channelId `from` $sourceCoreId `to` $targetCoreId
      attr-dict `:` type($output)
  }];
}

def SpatHostStoreSyncOp : SpatOp<"host_store_sync", []> {
  let summary = "Store a tensor to host memory and signal its consumer";

  let arguments = (ins
    Index:$sourceCoreId,
    Index:$targetCoreId,
    Index:$hostOffset,
    Index:$eventRegister,
    SpatTensor:$input
  );

  let assemblyFormat = [{
    $input `from` $sourceCoreId `to` $targetCoreId
      `host_offset` $hostOffset `event` $eventRegister attr-dict `:` type($input)
  }];
}

def SpatHostWaitLoadOp : SpatOp<"host_wait_load", []> {
  let summary = "Wait for producers and load from host memory";

  let arguments = (ins
    Index:$sourceCoreId,
    Index:$targetCoreId,
    Index:$hostOffset,
    Index:$eventRegister,
    Index:$waitValue
  );

  let results = (outs
    SpatTensor:$output
  );

  let assemblyFormat = [{
    `from` $sourceCoreId `to` $targetCoreId
      `host_offset` $hostOffset `event` $eventRegister `count` $waitValue
      attr-dict `:` type($output)
  }];
}

def SpatSyncOp : SpatOp<"sync", []> {
  let summary = "Signal a synchronization register on another processor";

  let arguments = (ins
    Index:$targetCoreId,
    Index:$eventRegister
  );

  let assemblyFormat = [{
    $targetCoreId `event` $eventRegister attr-dict
  }];
}

def SpatWaitOp : SpatOp<"wait", []> {
  let summary = "Wait for a synchronization register value";

  let arguments = (ins
    Index:$eventRegister,
    Index:$waitValue
  );

  let assemblyFormat = [{
    $eventRegister `value` $waitValue attr-dict
  }];
}

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

def SpatVMMOp : SpatOp<"wvmm", []> {
  let summary = "Vector-matrix multiplication within a weighted compute operation";

  let arguments = (ins
    SpatTensor:$weight,
    SpatTensor:$input
  );

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;

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

def SpatVVDMulOp : SpatOp<"vvdmul", []> {
  let summary = "Dot product between two runtime vectors";

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

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;

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

def SpatVAddOp : SpatOp<"vadd", []> {
  let summary = "Element-wise addition between two tensors; rhs must match lhs or be 1x1";

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

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;

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

def SpatVSubOp : SpatOp<"vsub", []> {
  let summary = "Element-wise subtraction between two tensors; rhs must match lhs or be 1x1";

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

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;

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

def SpatVMulOp : SpatOp<"vmul", []> {
  let summary = "Element-wise multiplication between two tensors; rhs must match lhs or be 1x1";

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

  let results = (outs
    SpatTensor:$output
  );

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

def SpatVAvgOp : SpatOp<"vavg", []> {
  let summary = "Average all elements of the input tensor to a single scalar wrapped in a tensor";

  let arguments = (ins
    SpatTensor:$input
  );

  let results = (outs
    SpatTensor:$output
  );

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

def SpatSigmoidOp : SpatOp<"sigmoid", []> {
  let summary = "Element-wise sigmoid activation";

  let arguments = (ins
    SpatTensor:$input
  );

  let results = (outs
    SpatTensor:$output
  );

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

def SpatSoftmaxOp : SpatOp<"softmax", []> {
  let summary = "Softmax over the full input tensor slice";

  let arguments = (ins
    SpatTensor:$input
  );

  let results = (outs
    SpatTensor:$output
  );

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

def SpatReluOp : SpatOp<"relu", []> {
  let summary = "Element-wise ReLU activation";

  let arguments = (ins
    SpatTensor:$input
  );

  let results = (outs
    SpatTensor:$output
  );

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

def SpatVMaxOp : SpatOp<"vmax", []> {
  let summary = "Element-wise max between two tensors";

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

  let results = (outs
    SpatTensor:$output
  );

  let hasVerifier = 1;

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

#endif // SPATIAL_DIALECT_H