Module: Ignis::LinAlg::Epilog

Defined in:
lib/nvruby/linalg/epilog.rb

Overview

Epilog operations for fused GEMM

Constant Summary collapse

TYPES =

Available epilog types

{
  default: CuBLASBindings::CUBLASLT_EPILOGUE_DEFAULT,
  relu: CuBLASBindings::CUBLASLT_EPILOGUE_RELU,
  relu_aux: CuBLASBindings::CUBLASLT_EPILOGUE_RELU_AUX,
  bias: CuBLASBindings::CUBLASLT_EPILOGUE_BIAS,
  relu_bias: CuBLASBindings::CUBLASLT_EPILOGUE_RELU_BIAS,
  relu_aux_bias: CuBLASBindings::CUBLASLT_EPILOGUE_RELU_AUX_BIAS,
  drelu: CuBLASBindings::CUBLASLT_EPILOGUE_DRELU,
  drelu_bgrad: CuBLASBindings::CUBLASLT_EPILOGUE_DRELU_BGRAD,
  gelu: CuBLASBindings::CUBLASLT_EPILOGUE_GELU,
  gelu_aux: CuBLASBindings::CUBLASLT_EPILOGUE_GELU_AUX,
  gelu_bias: CuBLASBindings::CUBLASLT_EPILOGUE_GELU_BIAS,
  gelu_aux_bias: CuBLASBindings::CUBLASLT_EPILOGUE_GELU_AUX_BIAS,
  dgelu: CuBLASBindings::CUBLASLT_EPILOGUE_DGELU,
  dgelu_bgrad: CuBLASBindings::CUBLASLT_EPILOGUE_DGELU_BGRAD,
  bgrada: CuBLASBindings::CUBLASLT_EPILOGUE_BGRADA,
  bgradb: CuBLASBindings::CUBLASLT_EPILOGUE_BGRADB
}.freeze

Class Method Summary collapse

Class Method Details

.available ⇒ Array<Symbol>

List available epilog types

Returns:

  • (Array<Symbol>) 


40
41
42
 
# File 'lib/nvruby/linalg/epilog.rb', line 40

def available
  TYPES.keys
end

.backward?(type) ⇒ Boolean

Check if epilog type is for backward pass

Parameters:

  • type (Symbol) —

    Epilog type

Returns:

  • (Boolean) 


61
62
63
 
# File 'lib/nvruby/linalg/epilog.rb', line 61

def backward?(type)
  %i[drelu drelu_bgrad dgelu dgelu_bgrad bgrada bgradb].include?(type)
end

.get(type) ⇒ Integer

Get epilog type constant

Parameters:

  • type (Symbol) —

    Epilog type

Returns:

  • (Integer) —

    cuBLASLt epilog constant

Raises:

  • (ArgumentError) 


31
32
33
34
35
36
 
# File 'lib/nvruby/linalg/epilog.rb', line 31

def get(type)
  constant = TYPES[type]
  raise ArgumentError, "Unknown epilog type: #{type}" unless constant

  constant
end

.supports_aux?(type) ⇒ Boolean

Check if epilog type supports auxiliary output

Parameters:

  • type (Symbol) —

    Epilog type

Returns:

  • (Boolean) 


54
55
56
 
# File 'lib/nvruby/linalg/epilog.rb', line 54

def supports_aux?(type)
  %i[relu_aux relu_aux_bias gelu_aux gelu_aux_bias].include?(type)
end

.supports_bias?(type) ⇒ Boolean

Check if epilog type supports bias

Parameters:

  • type (Symbol) —

    Epilog type

Returns:

  • (Boolean) 


47
48
49
 
# File 'lib/nvruby/linalg/epilog.rb', line 47

def supports_bias?(type)
  %i[bias relu_bias relu_aux_bias gelu_bias gelu_aux_bias].include?(type)
end