get_partial_matmul_right_val Subroutine

subroutine get_partial_matmul_right_val(this, upstream_grad, output)

Compute gradient w.r.t. right operand for matmul in reverse mode.

For C = A * B (where B is the right operand): dL/dB = A^T * dL/dC (2D case) dL/dB = left_operand (x) upstream_grad (outer product case)

Arguments

Type IntentOptional Attributes Name
class(array_type), intent(in) :: this
real(kind=real32), intent(in), dimension(:,:) :: upstream_grad
real(kind=real32), intent(out), dimension(:,:) :: output

Source Code

  subroutine get_partial_matmul_right_val(this, upstream_grad, output)
#else
  pure subroutine get_partial_matmul_right_val(this, upstream_grad, output)
#endif
    !! Compute gradient w.r.t. right operand for matmul in reverse mode.
    !!
    !! For C = A * B (where B is the right operand):
    !!   dL/dB = A^T * dL/dC  (2D case)
    !!   dL/dB = left_operand (x) upstream_grad  (outer product case)
    implicit none
    class(array_type), intent(in) :: this
    real(real32), dimension(:,:), intent(in) :: upstream_grad
    real(real32), dimension(:,:), intent(out) :: output

    integer :: i, j, s, m, n, num_elements, num_batches, num_upstream

    num_batches = size(upstream_grad, 2)

    if(size(this%left_operand%shape).eq.2)then
       m = this%left_operand%shape(1)
       n = this%left_operand%shape(2)
       if(this%left_operand%is_sample_dependent)then
          ! Per-sample: weight matrix varies per sample.
          ! output(:,s) = W_s^T * upstream_grad(:,s)
          ! where W_s = reshape(left%val(:,s), [m, n])
          ! Mathematically: dL/dB_flat = A_s^T * dL/dC per sample.
          ! Uses intrinsic matmul — BLAS overhead not worthwhile for per-sample vectors.
          block
            real(real32), dimension(m, n) :: temp
            do s = 1, num_batches
               temp = reshape(this%left_operand%val(:,s), [m, n])
               output(:,s) = matmul(transpose(temp), upstream_grad(:,s))
            end do
          end block
       else
          ! Non-sample-dependent: single batch operation across all samples.
          ! output(n, S) = W^T(n, m) * upstream_grad(m, S)
          ! Mathematically: dL/dB = A^T * dL/dC for each sample, batched.
#ifdef USE_BLAS
          block
            real(real32), pointer :: W(:,:)
            ! Pointer view avoids reshape and transpose: left_operand%val(:,1) is
            ! column-major contiguous (m, n) data. SGEMM transposes in-place.
            W(1:m, 1:n) => this%left_operand%val(:,1)
            ! sgemm: C = alpha * A^T * B + beta * C
            ! A = W(m, n), transposed to W^T(n, m); B = upstream_grad(m, S)
            ! m_arg = n (rows of W^T and output)
            ! n_arg = num_batches (columns of upstream_grad and output)
            ! k = m (columns of W^T / rows of upstream_grad)
            ! lda = m (leading dim of W before transpose), ldb = m, ldc = n
            call sgemm('T', 'N', n, num_batches, m, &
                 1.0_real32, W, m, upstream_grad, m, &
                 0.0_real32, output, n)
          end block
#else
          block
            real(real32), dimension(n, m) :: temp_t
            temp_t = transpose(reshape(this%left_operand%val(:,1), [m, n]))
            output = matmul(temp_t, upstream_grad)
          end block
#endif
       end if
    else
       ! Outer product case: output(i + (j-1)*num_el, s) = left(i,s) * grad(j,s)
       num_elements = size(this%left_operand%val,1)
       num_upstream = size(upstream_grad, 1)
       if(this%left_operand%is_sample_dependent)then
          do concurrent(s = 1:num_batches, j = 1:num_upstream)
             output((j-1)*num_elements+1:j*num_elements, s) = &
                  upstream_grad(j,s) * this%left_operand%val(:,s)
          end do
       else
          do concurrent(s=1:num_batches, j=1:num_upstream)
             output((j-1)*num_elements+1:j*num_elements, s) = &
                  this%left_operand%val(:, 1) * upstream_grad(j, s)
          end do
       end if
    end if

  end subroutine get_partial_matmul_right_val