#pragma once

#include "mlir/IR/BuiltinTypes.h"

#include "llvm/ADT/SmallVector.h"

#include <cstdint>

namespace onnx_mlir {

struct ContractionProblem {
  llvm::SmallVector<int64_t> lhsBatchShape;
  llvm::SmallVector<int64_t> rhsBatchShape;
  llvm::SmallVector<int64_t> outputBatchShape;
  int64_t lhsBatch = 1;
  int64_t rhsBatch = 1;
  int64_t batch = 1;
  int64_t m = 0;
  int64_t k = 0;
  int64_t n = 0;
  mlir::Type lhsElementType;
  mlir::Type rhsElementType;
  mlir::Type resultElementType;
};

} // namespace onnx_mlir