Быстрая тензор инициализация в Джулии

Question

Я хотел бы инициализировать 3d тензора (многомерный массив) со значениями «по диагонали Gaussian»

exp(-32*(u^2 + 16*(v^2 + w^2))) 

где u = 1/sqrt(3)*(x+y+z) и v,w являются любыми две координатами ортогональных u, дискретизирован на равномерной сетке на [-1,1]^3. Следующий код обеспечивает:

function gaussian3d(n) 
    q = qr(ones(3,1), thin=false)[1] 
    x = linspace(-1.,1., n) 
    p = Array(Float64,(n,n,n)) 
    square(x) = x*x 
    Base.@nloops 3 i p begin 
     @inbounds p[i_1,i_2,i_3] = 
      exp(
       -32*(
        square(q[1,1]*x[i_1] + q[2,1]*x[i_2] + q[3,1]*x[i_3]) 
        + 16*(
         square(q[1,2]*x[i_1] + q[2,2]*x[i_2] + q[3,2]*x[i_3]) + 
         square(q[1,3]*x[i_1] + q[2,3]*x[i_2] + q[3,3]*x[i_3]) 
        ) 
       ) 
      ) 
    end 
    return p 
end 

Это, однако, довольно медленно. Например, если я заменяю определяющую функцию exp(x*y*z), код работает быстрее на 50 раз. Кроме того, макрос @time сообщает ~ 20% GC-времени для вышеуказанного кода, который я не понимаю, откуда они. (Эти числовые значения были получены с n = 128.) Мои вопросы поэтому

• Как я могу ускорить этот кусок кода?
• Где скрыто выделение памяти, которое вызывает накладные расходы GC?

Answer 1

не зная ничего о 3D тензоров со значениями «по диагонали Gaussian», используя square комментарий от оригинального поста, «набрав» q (@code_warntype helps here: Большой скачок производительности), и далее Специализируя @nloops, это работает намного быстрее на платформах, на которых я его пробовал.

julia> square(x::Float64) = x * x 
square (generic function with 1 method) 

julia> function my_gaussian3d(n) 
      q::Array{Float64,2} = qr(ones(3,1), thin=false)[1] 
      x = linspace(-1.,1., n) 
      p = Array(Float64,(n,n,n)) 
      Base.@nloops 3 i p d->x_d=x[i_d] begin 
       @inbounds p[i_1,i_2,i_3] = 
        exp(
         -32*(
          square(q[1,1]*x_1 + q[2,1]*x_2 + q[3,1]*x_3) 
          + 16*(
           square(q[1,2]*x_1 + q[2,2]*x_2 + q[3,2]*x_3) + 
           square(q[1,3]*x_1 + q[2,3]*x_2 + q[3,3]*x_3) 
          ) 
         ) 
        ) 
      end 
      return p 
     end 
my_gaussian3d (generic function with 1 method) 

julia> @time gaussian3d(128); 
elapsed time: 3.952389337 seconds (1264 MB allocated, 4.50% gc time in 57 pauses with 0 full sweep) 

julia> @time gaussian3d(128); 
elapsed time: 3.527316699 seconds (1264 MB allocated, 4.42% gc time in 58 pauses with 0 full sweep) 

julia> @time my_gaussian3d(128); 
elapsed time: 0.285837566 seconds (16 MB allocated) 

julia> @time my_gaussian3d(128); 
elapsed time: 0.28476448 seconds (16 MB allocated, 1.22% gc time in 0 pauses with 0 full sweep) 

julia> my_gaussian3d(128) == gaussian3d(128) 
true