#pragma once

#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Dialect.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/RegionKindInterface.h"
#include "mlir/IR/Types.h"
#include "mlir/Interfaces/ParallelCombiningOpInterface.h"

#include "llvm/ADT/DenseSet.h"
#include "llvm/ADT/SetVector.h"

#include <map>
#include <optional>
#include <string>
#include <tuple>

/// Include the auto-generated header files containing the declarations
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialDialect.hpp.inc"

#define GET_TYPEDEF_CLASSES
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialTypes.hpp.inc"

#define GET_OP_CLASSES
#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp.inc"

namespace onnx_mlir {
namespace spatial {

bool isGraphComputeLike(mlir::Operation* op);
bool isGraphBatchComputeLike(mlir::Operation* op);
bool isScheduledComputeLike(mlir::Operation* op);
bool isScheduledBatchComputeLike(mlir::Operation* op);
bool isAnySpatialComputeLike(mlir::Operation* op);
bool isAnySpatialComputeBatchLike(mlir::Operation* op);

using SpatCompute = SpatGraphCompute;
using SpatComputeBatch = SpatGraphComputeBatch;

} // namespace spatial
} // namespace onnx_mlir