Class: Ignis::Tensor::Contraction

Inherits:
Object
  • Object
show all
Defined in:
lib/nvruby/tensor/contraction.rb

Overview

Tensor contraction operation using cuTENSOR Supports Einstein notation for expressing tensor contractions

Examples:

Matrix multiplication (ij,jk->ik)

contraction = Contraction.new("ij,jk->ik", a, b)
result = contraction.execute
contraction.destroy!

Instance Attribute Summary collapse

Class Method Summary collapse

Instance Method Summary collapse

Constructor Details

#initialize(expression, tensor_a, tensor_b, alpha: 1.0, beta: 0.0) ⇒ Contraction

Initialize tensor contraction

Parameters:

  • expression (String) —

    Einstein notation expression

  • tensor_a (NvArray) —

    First input tensor

  • tensor_b (NvArray) —

    Second input tensor

  • alpha (Float) (defaults to: 1.0) —

    Scaling factor

  • beta (Float) (defaults to: 0.0) —

    Scaling for output (for accumulation)



52
53
54
55
56
57
58
59
60
61
62
63
64
65
 
# File 'lib/nvruby/tensor/contraction.rb', line 52

def initialize(expression, tensor_a, tensor_b, alpha: 1.0, beta: 0.0)
  @expression = expression
  @tensor_a = tensor_a
  @tensor_b = tensor_b
  @alpha = alpha
  @beta = beta
  @parsed = self.class.parse_expression(expression)
  @planned = false
  @handle = nil
  @plan = nil

  validate_inputs!
  compute_output_shape!
end

Instance Attribute Details

#expression ⇒ String (readonly)

Returns Einstein notation expression.

Returns:

  • (String) —

    Einstein notation expression



16
17
18
 
# File 'lib/nvruby/tensor/contraction.rb', line 16

def expression
  @expression
end

#output_shape ⇒ Array<Integer> (readonly)

Returns Output shape.

Returns:

  • (Array<Integer>) —

    Output shape



25
26
27
 
# File 'lib/nvruby/tensor/contraction.rb', line 25

def output_shape
  @output_shape
end

#tensor_a ⇒ NvArray (readonly)

Returns First input tensor.

Returns:



19
20
21
 
# File 'lib/nvruby/tensor/contraction.rb', line 19

def tensor_a
  @tensor_a
end

#tensor_b ⇒ NvArray (readonly)

Returns Second input tensor.

Returns:

  • (NvArray) —

    Second input tensor



22
23
24
 
# File 'lib/nvruby/tensor/contraction.rb', line 22

def tensor_b
  @tensor_b
end

Class Method Details

.parse_expression(expression) ⇒ Hash

Parse Einstein notation expression

Parameters:

  • expression (String) —

    Expression like “ij,jk->ik”

Returns:

  • (Hash) —

    Parsed components



30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
 
# File 'lib/nvruby/tensor/contraction.rb', line 30

def self.parse_expression(expression)
  unless expression.include?(",") && expression.include?("->")
    raise ArgumentError, "Invalid Einstein notation: #{expression}. Expected format: 'ij,jk->ik'"
  end

  input_part, output_modes = expression.split("->")
  input_a_modes, input_b_modes = input_part.split(",")

  {
    input_a_modes: input_a_modes.chars,
    input_b_modes: input_b_modes.chars,
    output_modes: output_modes.chars,
    contracted_modes: (input_a_modes.chars & input_b_modes.chars) - output_modes.chars
  }
end

Instance Method Details

#destroy! ⇒ void

This method returns an undefined value.

Free cuTENSOR resources



91
92
93
94
95
 
# File 'lib/nvruby/tensor/contraction.rb', line 91

def destroy!
  @handle = nil
  @plan = nil
  @planned = false
end

#execute ⇒ NvArray

Execute the tensor contraction

Returns:



69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
 
# File 'lib/nvruby/tensor/contraction.rb', line 69

def execute
  # Ensure tensors are on device
  a = @tensor_a.on_device? ? @tensor_a : @tensor_a.to_device
  b = @tensor_b.on_device? ? @tensor_b : @tensor_b.to_device

  # Create output tensor
  output = NvArray.zeros(@output_shape, dtype: a.dtype, device: a.device_index)
  output = output.to_device unless output.on_device?

  # For now, use optimized path for common patterns
  # Full cuTENSOR integration will use the plan/execute pattern
  if matrix_multiply?
    execute_as_matmul(a, b, output)
  else
    execute_general_contraction(a, b, output)
  end

  output
end