| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | submodule (athena__diffstruc_extd) athena__diffstruc_extd_submodule_msgpass_kipf | ||
| 2 | !! Submodule containing implementations for extended diffstruc array operations | ||
| 3 | |||
| 4 | contains | ||
| 5 | |||
| 6 | !############################################################################### | ||
| 7 | − | module function kipf_propagate(vertex_features, adj_ia, adj_ja) result(c) | |
| 8 | !! Propagate values from one autodiff array to another | ||
| 9 | class(array_type), intent(in), target :: vertex_features | ||
| 10 | integer, dimension(:), intent(in) :: adj_ia | ||
| 11 | integer, dimension(:,:), intent(in) :: adj_ja | ||
| 12 | type(array_type), pointer :: c | ||
| 13 | |||
| 14 | integer :: v, w | ||
| 15 | real(real32) :: coeff | ||
| 16 | |||
| 17 | − | c => vertex_features%create_result() | |
| 18 | ! propagate 1D array by using shape to swap dimensions | ||
| 19 | − | do concurrent(v = 1:size(vertex_features%val,2)) | |
| 20 | − | c%val(:,v) = 0._real32 | |
| 21 | − | do w = adj_ia(v), adj_ia(v+1)-1 | |
| 22 | |||
| 23 | ! if( adj_ja(2,w) .eq. 0 )then | ||
| 24 | ! coeff = 1._real32 | ||
| 25 | ! !else | ||
| 26 | ! ! coeff = edge_weights(adj_ja(2,w)) | ||
| 27 | ! end if | ||
| 28 | !coeff = coeff * ( & | ||
| 29 | coeff = ( & | ||
| 30 | − | ( adj_ia(v+1) - adj_ia(v) ) * & | |
| 31 | − | ( adj_ia( adj_ja(1,w) + 1 ) - adj_ia( adj_ja(1,w) ) ) & | |
| 32 | − | ) ** ( -0.5_real32 ) | |
| 33 | |||
| 34 | − | c%val(:,v) = c%val(:,v) + coeff * [ vertex_features%val(:, adj_ja(1, w)) ] | |
| 35 | end do | ||
| 36 | end do | ||
| 37 | |||
| 38 | − | c%indices = adj_ia | |
| 39 | − | c%adj_ja = adj_ja | |
| 40 | − | c%get_partial_left => get_partial_kipf_propagate_left | |
| 41 | − | c%get_partial_left_val => get_partial_kipf_propagate_left_val | |
| 42 | − | if(vertex_features%requires_grad)then | |
| 43 | − | c%requires_grad = .true. | |
| 44 | − | c%is_forward = vertex_features%is_forward | |
| 45 | − | c%operation = 'kipf_propagate' | |
| 46 | − | c%left_operand => vertex_features | |
| 47 | − | c%owns_left_operand = vertex_features%is_temporary | |
| 48 | end if | ||
| 49 | − | end function kipf_propagate | |
| 50 | !------------------------------------------------------------------------------- | ||
| 51 | − | function get_partial_kipf_propagate_left(this, upstream_grad) result(output) | |
| 52 | class(array_type), intent(inout) :: this | ||
| 53 | type(array_type), intent(in) :: upstream_grad | ||
| 54 | type(array_type) :: output | ||
| 55 | |||
| 56 | output = reverse_kipf_propagate( upstream_grad, & | ||
| 57 | this%indices, this%adj_ja, & | ||
| 58 | − | num_features = [ & | |
| 59 | − | this%left_operand%shape(1), this%right_operand%shape(1) & | |
| 60 | ], & | ||
| 61 | num_elements = [ & | ||
| 62 | size(this%left_operand%val,2), size(this%right_operand%val,2) & | ||
| 63 | ] & | ||
| 64 | − | ) | |
| 65 | |||
| 66 | − | end function get_partial_kipf_propagate_left | |
| 67 | !------------------------------------------------------------------------------- | ||
| 68 | − | pure subroutine get_partial_kipf_propagate_left_val(this, upstream_grad, output) | |
| 69 | class(array_type), intent(in) :: this | ||
| 70 | real(real32), dimension(:,:), intent(in) :: upstream_grad | ||
| 71 | real(real32), dimension(:,:), intent(out) :: output | ||
| 72 | |||
| 73 | integer :: v, w, i | ||
| 74 | |||
| 75 | − | output = 0._real32 | |
| 76 | − | do concurrent(v=1:size(upstream_grad,2)) | |
| 77 | − | do w = this%indices(v), this%indices(v+1)-1 | |
| 78 | − | do concurrent(i = 1:size(upstream_grad,1)) | |
| 79 | − | output(i,this%adj_ja(1,w)) = output(i,this%adj_ja(1,w)) + & | |
| 80 | − | upstream_grad(i, v) | |
| 81 | end do | ||
| 82 | end do | ||
| 83 | end do | ||
| 84 | |||
| 85 | − | end subroutine get_partial_kipf_propagate_left_val | |
| 86 | !############################################################################### | ||
| 87 | |||
| 88 | |||
| 89 | !############################################################################### | ||
| 90 | − | function reverse_kipf_propagate( & | |
| 91 | − | a, adj_ia, adj_ja, num_features, num_elements & | |
| 92 | ) result(c) | ||
| 93 | !! Reverse propagate values from one autodiff array to another | ||
| 94 | class(array_type), intent(in), target :: a | ||
| 95 | integer, dimension(:), intent(in) :: adj_ia | ||
| 96 | integer, dimension(:,:), intent(in) :: adj_ja | ||
| 97 | integer, dimension(2), intent(in) :: num_features, num_elements | ||
| 98 | type(array_type), pointer :: c | ||
| 99 | |||
| 100 | integer :: v, w | ||
| 101 | |||
| 102 | c => a%create_result(array_shape=[ & | ||
| 103 | num_features(1), num_elements(1) & | ||
| 104 | − | ]) | |
| 105 | − | c%val = 0.0_real32 | |
| 106 | − | do concurrent(v=1:num_elements(1)) | |
| 107 | − | do w = adj_ia(v), adj_ia(v+1)-1 | |
| 108 | − | c%val(:,adj_ja(1,w)) = c%val(:,adj_ja(1,w)) + & | |
| 109 | − | [ a%val(:, v) ] | |
| 110 | end do | ||
| 111 | end do | ||
| 112 | |||
| 113 | − | c%indices = adj_ia | |
| 114 | − | c%adj_ja = adj_ja | |
| 115 | − | c%get_partial_left => get_partial_left_reverse_kipf_propagate | |
| 116 | − | c%get_partial_left_val => get_partial_left_reverse_kipf_propagate_val | |
| 117 | − | if(a%requires_grad)then | |
| 118 | − | c%requires_grad = .true. | |
| 119 | − | c%is_forward = a%is_forward | |
| 120 | − | c%operation = 'reverse_kipf_propagate' | |
| 121 | − | c%left_operand => a | |
| 122 | − | c%owns_left_operand = a%is_temporary | |
| 123 | end if | ||
| 124 | − | end function reverse_kipf_propagate | |
| 125 | !------------------------------------------------------------------------------- | ||
| 126 | − | function get_partial_left_reverse_kipf_propagate( & | |
| 127 | this, upstream_grad & | ||
| 128 | − | ) result(output) | |
| 129 | implicit none | ||
| 130 | class(array_type), intent(inout) :: this | ||
| 131 | type(array_type), intent(in) :: upstream_grad | ||
| 132 | type(array_type) :: output | ||
| 133 | |||
| 134 | output = kipf_propagate( upstream_grad, & | ||
| 135 | this%indices, this%adj_ja & | ||
| 136 | − | ) | |
| 137 | |||
| 138 | − | end function get_partial_left_reverse_kipf_propagate | |
| 139 | !------------------------------------------------------------------------------- | ||
| 140 | − | pure subroutine get_partial_left_reverse_kipf_propagate_val( & | |
| 141 | − | this, upstream_grad, output & | |
| 142 | ) | ||
| 143 | implicit none | ||
| 144 | class(array_type), intent(in) :: this | ||
| 145 | real(real32), dimension(:,:), intent(in) :: upstream_grad | ||
| 146 | real(real32), dimension(:,:), intent(out) :: output | ||
| 147 | |||
| 148 | integer :: v, w, i | ||
| 149 | − | output = 0._real32 | |
| 150 | − | do concurrent(v=1:size(upstream_grad,2)) | |
| 151 | − | do w = this%indices(v), this%indices(v+1)-1 | |
| 152 | − | do concurrent(i = 1:size(upstream_grad,1)) | |
| 153 | − | output(i,this%adj_ja(1,w)) = output(i,this%adj_ja(1,w)) + & | |
| 154 | − | upstream_grad(i, v) | |
| 155 | end do | ||
| 156 | end do | ||
| 157 | end do | ||
| 158 | |||
| 159 | − | end subroutine get_partial_left_reverse_kipf_propagate_val | |
| 160 | !############################################################################### | ||
| 161 | |||
| 162 | end submodule athena__diffstruc_extd_submodule_msgpass_kipf | ||
| 163 |