#ifndef SPATIAL_TO_PIM
#define SPATIAL_TO_PIM

#ifndef OP_BASE
include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Tensor/IR/TensorOps.td"
include "src/Dialect/ONNX/ONNX.td"
include "src/Accelerators/PIM/Dialect/Spatial/Spatial.td"
include "src/Accelerators/PIM/Dialect/Pim/Pim.td"
#endif // OP_BASE

def onnxToPimTranspose : Pat<
  (ONNXTransposeOp:$srcOpRes $data, $perms),
  (PimTransposeOp $data, $perms,
    (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
>;

def spatToPimVMM : Pat<
  (SpatWeightedVMMOp:$srcOpRes $weightIndex, $vector),
  (PimVMMOp $weightIndex, $vector,
    (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
>;

def spatToPimMVM : Pat<
  (SpatWeightedMVMOp:$srcOpRes $weightIndex, $vector),
  (PimMVMOp $weightIndex, $vector,
    (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
>;

def spatToPimVVAdd : Pat<
  (SpatVAddOp:$srcOpRes $a, $b),
  (PimVVAddOp $a, $b,
    (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
>;

def spatToPimVVMul : Pat<
  (SpatVMulOp:$srcOpRes $a, $b),
  (PimVVMulOp $a, $b,
    (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
>;

def spatToPimVVMax : Pat<
  (SpatVMaxOp:$srcOpRes $a, $b),
  (PimVVMaxOp $a, $b,
    (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
>;

def spatToPimVRelu : Pat<
  (SpatReluOp:$srcOpRes $input),
  (PimVReluOp $input,
    (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
>;

#endif // SPATIAL_TO_PIM
