/*************************************************************************
 * Copyright (C) [2023-2024] by Cambricon, Inc.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
 * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
 * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
 * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
 * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
 * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 *************************************************************************/

#ifndef CSRC_OPS_OP_DESCRIPTOR_BATCHMATMUL_DESCRIPTOR_H_
#define CSRC_OPS_OP_DESCRIPTOR_BATCHMATMUL_DESCRIPTOR_H_

#include <functional>
#include <memory>
#include "base_descriptor.h"
#include "common/utils.h"

namespace tmo {
namespace op_desc {

class TMO_HIDDEN BatchMatMulDescImpl {
 public:
  BatchMatMulDescImpl();
  // using clear function to avoid compilation errors.
  ~BatchMatMulDescImpl() { clear(); }
  DELETE_COPY_ASSIGN_CONSTRUCT(BatchMatMulDescImpl);

  cnnlMatMulDescriptor_t bmm_desc_;
  cnnlMatMulHeuristicResult_t heuristic_result_;
  cnnlMatMulAlgo_t algo_;

 private:
  inline void clear() {
    CNNL_CHECK_FATAL(cnnlDestroyMatMulDescriptor(this->bmm_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyMatMulHeuristicResult(this->heuristic_result_));
    CNNL_CHECK_FATAL(cnnlDestroyMatMulAlgo(this->algo_));
  }
};

class TMO_EXPORT BatchMatMulDesc {
 public:
  using impl_type_ =
      std::unique_ptr<BatchMatMulDescImpl, std::function<void(BatchMatMulDescImpl *)>>;
  BatchMatMulDesc(cnnlMatMulHeuristicResult_t &heuristic_result,
                  cnnlMatMulAlgo_t &algo,
                  bool use_beta = false,
                  bool trans_a = false,
                  bool trans_b = true);

  CLASS_CAST_TYPE_OPERATOR_DEFINE(cnnlMatMulDescriptor_t, impl_->bmm_desc_)
  CLASS_CAST_TYPE_OPERATOR_DEFINE(cnnlMatMulHeuristicResult_t, impl_->heuristic_result_)
  CLASS_CAST_TYPE_OPERATOR_DEFINE(cnnlMatMulAlgo_t, impl_->algo_)

 private:
  impl_type_ impl_;
};

}  // namespace op_desc
}  // namespace tmo

#endif  // CSRC_OPS_OP_DESCRIPTOR_BATCHMATMUL_DESCRIPTOR_H_
