/*************************************************************************
 * 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.
 *************************************************************************/

#include "batchmatmul_descriptor.h"

namespace tmo {
namespace op_desc {

namespace {
static OpDescPool<BatchMatMulDescImpl> bmm_instance;
}

BatchMatMulDesc::BatchMatMulDesc(cnnlMatMulHeuristicResult_t &heuristic_result,
                                 cnnlMatMulAlgo_t &algo,
                                 bool use_beta,
                                 bool trans_a,
                                 bool trans_b) {
  impl_ = bmm_instance.get();
  auto bmm_desc_ = impl_->bmm_desc_;
  heuristic_result = impl_->heuristic_result_;
  algo = impl_->algo_;
  int32_t matmul_trans_a = int(trans_a);
  int32_t matmul_trans_b = int(trans_b);
  int32_t matmul_use_tf32 = 0;
  int32_t matmul_use_beta = int(use_beta);
  CNNL_CHECK_FATAL(cnnlSetMatMulDescAttr(bmm_desc_, CNNL_MATMUL_DESC_TRANSA, &(matmul_trans_a),
                                         sizeof(int32_t)));
  CNNL_CHECK_FATAL(cnnlSetMatMulDescAttr(bmm_desc_, CNNL_MATMUL_DESC_TRANSB, &(matmul_trans_b),
                                         sizeof(int32_t)));
  CNNL_CHECK_FATAL(cnnlSetMatMulDescAttr(bmm_desc_, CNNL_MATMUL_ALLOW_TF32, &(matmul_use_tf32),
                                         sizeof(int32_t)));
  CNNL_CHECK_FATAL(
      cnnlSetMatMulDescAttr(bmm_desc_, CNNL_MATMUL_USE_BETA, &(matmul_use_beta), sizeof(int32_t)));
}

BatchMatMulDescImpl::BatchMatMulDescImpl() {
  CNNL_CHECK_FATAL(cnnlCreateMatMulDescriptor(&this->bmm_desc_));
  CNNL_CHECK_FATAL(cnnlCreateMatMulHeuristicResult(&this->heuristic_result_));
  CNNL_CHECK_FATAL(cnnlCreateMatMulAlgo(&this->algo_));
}

}  // namespace op_desc
}  // namespace tmo
