/*************************************************************************
 * 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_MATMUL_DESCRIPTOR_H_
#define CSRC_OPS_OP_DESCRIPTOR_MATMUL_DESCRIPTOR_H_

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

namespace tmo {
namespace op_desc {

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

  cnnlMatMulExDescriptor_t matmul_desc_;
  cnnlActivationDescriptor_t act_desc_;

 private:
  inline void clear() {
    CNNL_CHECK_FATAL(cnnlDestroyMatMulExDescriptor(this->matmul_desc_));
    CNNL_CHECK_FATAL(cnnlDestroyActivationDescriptor(this->act_desc_));
  }
};

class TMO_EXPORT MatMulDesc {
 public:
  using impl_type_ = std::unique_ptr<MatMulDescImpl, std::function<void(MatMulDescImpl *)>>;
  MatMulDesc(const cnnlTensorDescriptor_t &bias_tensor,
             const void *bias_data,
             size_t &workspace_size,
             int32_t lda,
             int32_t ldb,
             int32_t ldc,
             const std::string &act_name = "none",
             bool use_beta = false,
             bool trans_a = false,
             bool trans_b = true,
             bool fast_act = true,
             bool approximate = true);

  CLASS_CAST_TYPE_OPERATOR_DEFINE(cnnlMatMulExDescriptor_t, impl_->matmul_desc_)
  CLASS_CAST_TYPE_OPERATOR_DEFINE(cnnlActivationDescriptor_t, impl_->act_desc_)

 private:
  size_t max_workspace_size_ = 128 * 1024 * 1024;
  impl_type_ impl_;
};

}  // namespace op_desc
}  // namespace tmo

#endif  // CSRC_OPS_OP_DESCRIPTOR_MATMUL_DESCRIPTOR_H_
