launchGemmAndBiasCublasLt Class — pytorch Architecture

Source Code

aten/src/ATen/native/cuda/Blas.cpp lines 260–300

template <typename scalar_t, typename res_scalar_t = scalar_t>
bool launchGemmAndBiasCublasLt(
    // args contains result which is modified
    cublasCommonArgs& args,
    const std::optional<Tensor>& self,
    const Scalar& alpha,
    Activation activation = Activation::None
) {
  // We apply bias in the epilogue only when it is 1D,
  // or when it can be squeezed to 1D.
  // self_ptr == nullptr implies ignore bias epilogue
  // and use standard gemm-like API.
  const auto* self_ptr = self.has_value() ? self.value().const_data_ptr<scalar_t>() : static_cast<const scalar_t*>(nullptr);


  const auto tuning_ctx = at::cuda::tunable::getTuningContext();
  if (tuning_ctx->IsTunableOpEnabled()) {
    // TODO: maybe also return some success state?
    launchTunableGemmAndBias<scalar_t>(
      args, alpha, self_ptr, activation_to_gemm_and_blas_arg(activation)
    );
    return true;
  }

  return at::cuda::blas::gemm_and_bias<scalar_t, res_scalar_t>(
    args.transa == 't',
    args.transb == 't',
    args.m,
    args.n,
    args.k,
    alpha.to<at::opmath_type<scalar_t>>(),
    args.mata->const_data_ptr<scalar_t>(),
    args.lda,
    args.matb->const_data_ptr<scalar_t>(),
    args.ldb,
    self_ptr,
    args.result->data_ptr<res_scalar_t>(),
    args.result_ld,
    activation_to_gemm_and_blas_arg(activation)
  );
}