InfiniOps
Operator Library for Accelerators
Loading...
Searching...
No Matches
infini::ops::CutlassScaledMm Class Referenceabstract

#include <cutlass_scaled_mm.h>

Inheritance diagram for infini::ops::CutlassScaledMm:
infini::ops::Operator< CutlassScaledMm > infini::ops::OperatorBase infini::ops::OperatorBase

Public Member Functions

 CutlassScaledMm (const Tensor a, const Tensor b, const Tensor scale_a, const Tensor scale_b, std::optional< Tensor > bias, const DataType out_dtype, Tensor out)
 
virtual void operator() (const Tensor a, const Tensor b, const Tensor scale_a, const Tensor scale_b, std::optional< Tensor > bias, const DataType out_dtype, Tensor out) const =0
 
- Public Member Functions inherited from infini::ops::Operator< CutlassScaledMm >
void operator() (const Handle &handle, const Args &... args)
 
void operator() (const Args &... args) const
 
void operator() (const Handle &handle, const Args &... args)
 
void operator() (const Args &... args) const
 
- Public Member Functions inherited from infini::ops::OperatorBase
virtual ~OperatorBase ()=default
 
virtual std::size_t workspace_size_in_bytes () const
 
void set_handle (const Handle &handle)
 
void set_config (const Config &config)
 
void set_stream (void *stream)
 
void set_workspace (void *workspace)
 
void set_workspace_size_in_bytes (std::size_t workspace_size_in_bytes)
 
virtual ~OperatorBase ()=default
 
virtual std::size_t workspace_size_in_bytes () const
 
void set_handle (const Handle &handle)
 
void set_config (const Config &config)
 
void set_stream (void *stream)
 
void set_workspace (void *workspace)
 
void set_workspace_size_in_bytes (std::size_t workspace_size_in_bytes)
 

Protected Attributes

Tensor::Size m_ {0}
 
Tensor::Size n_ {0}
 
Tensor::Size k_ {0}
 
Tensor::Stride lda_ {0}
 
Tensor::Stride ldb_ {0}
 
Tensor::Stride ldo_ {0}
 
Tensor::Size scale_a_size_ {0}
 
Tensor::Size scale_b_size_ {0}
 
DataType out_dtype_
 
- Protected Attributes inherited from infini::ops::OperatorBase
std::unique_ptr< Handle > handle_ptr_
 
std::unique_ptr< Config > config_ptr_
 
void * stream_ {nullptr}
 
void * workspace_ {nullptr}
 
std::size_t workspace_size_in_bytes_ {0}
 

Additional Inherited Members

- Static Public Member Functions inherited from infini::ops::Operator< CutlassScaledMm >
static void clear_cache ()
 
static void clear_cache ()
 
static std::unique_ptr< Operator > Make (const Config &config, const Tensor tensor, Args &&... args)
 
static std::unique_ptr< Operator > Make (const Tensor tensor, Args &&... args)
 
static std::unique_ptr< Operator > Make (const Config &config, const std::vector< Tensor > tensors, Args &&... args)
 
static std::unique_ptr< Operator > Make (const std::vector< Tensor > tensors, Args &&... args)
 
static std::unique_ptr< Operator > Make (const Config &config, const Tensor tensor, Args &&... args)
 
static std::unique_ptr< Operator > Make (const Tensor tensor, Args &&... args)
 
static std::unique_ptr< Operator > Make (const Config &config, const std::vector< Tensor > tensors, Args &&... args)
 
static std::unique_ptr< Operator > Make (const std::vector< Tensor > tensors, Args &&... args)
 
static void Call (const Handle &handle, const Config &config, const Args &... args)
 
static void Call (const Tensor tensor, const Args &... args)
 
static auto Call (const TensorLike &tensor, const Args &... args)
 
static void Call (const Handle &handle, const Config &config, const Args &... args)
 
static void Call (const Tensor tensor, const Args &... args)
 
static auto Call (const TensorLike &tensor, const Args &... args)
 
static std::vector< std::size_t > active_implementation_indices (Device::Type dev_type)
 
static std::vector< std::size_t > active_implementation_indices (Device::Type dev_type)
 
- Static Protected Attributes inherited from infini::ops::Operator< CutlassScaledMm >
static constexpr Device::Type device_type_
 
static constexpr std::size_t implementation_index_
 

Constructor & Destructor Documentation

◆ CutlassScaledMm()

infini::ops::CutlassScaledMm::CutlassScaledMm ( const Tensor  a,
const Tensor  b,
const Tensor  scale_a,
const Tensor  scale_b,
std::optional< Tensor >  bias,
const DataType  out_dtype,
Tensor  out 
)
inline

Member Function Documentation

◆ operator()()

virtual void infini::ops::CutlassScaledMm::operator() ( const Tensor  a,
const Tensor  b,
const Tensor  scale_a,
const Tensor  scale_b,
std::optional< Tensor >  bias,
const DataType  out_dtype,
Tensor  out 
) const
pure virtual

Member Data Documentation

◆ k_

Tensor::Size infini::ops::CutlassScaledMm::k_ {0}
protected

◆ lda_

Tensor::Stride infini::ops::CutlassScaledMm::lda_ {0}
protected

◆ ldb_

Tensor::Stride infini::ops::CutlassScaledMm::ldb_ {0}
protected

◆ ldo_

Tensor::Stride infini::ops::CutlassScaledMm::ldo_ {0}
protected

◆ m_

Tensor::Size infini::ops::CutlassScaledMm::m_ {0}
protected

◆ n_

Tensor::Size infini::ops::CutlassScaledMm::n_ {0}
protected

◆ out_dtype_

DataType infini::ops::CutlassScaledMm::out_dtype_
protected

◆ scale_a_size_

Tensor::Size infini::ops::CutlassScaledMm::scale_a_size_ {0}
protected

◆ scale_b_size_

Tensor::Size infini::ops::CutlassScaledMm::scale_b_size_ {0}
protected

The documentation for this class was generated from the following file: