Julia 如何在避免溢出的同时将'abs2'应用于'CuArray{Complex{Int8}}}'

这个问题出现在julialang Slack helpdesk频道

---

---

如何计算

复杂{Int8}

数组的

abs2

，而元素不超过

Int8

的最大值，并且不增加CPU->GPU传输时间？我有两个预先分配的

CuArray

s的eltype

Complex{Int8}

和

Int32

。复杂数组的

abs2

结果应存储在

Int32

CuArray

由于julia的通用接口，答案基本上与CPU上的结果相同

using CUDA

let v = CUDA.rand(Complex{Int8}, 10)
    u = CUDA.zeros(Int32, length(v))
    u .= abs2.(Complex{Int32}.(v))
end

在这段代码中，我将

u

的每个元素转换为

Complex{Int32}

，然后将

abs2

应用于该元素，并将结果存储在

u

中。多亏了广播语义，上面所有的循环都融合在一起了，所以没有创建临时的

从本质上讲，这可以通过就地地图实现

let v = CUDA.rand(Complex{Int8}, 10)
    u = CUDA.zeros(Int32, length(v))
    map!(x -> abs2(Complex{Int32}(x)), u, v)
end

let v = CUDA.rand(Complex{Int8}, 10)
    u = CUDA.zeros(Int32, length(v))
    map!(x -> abs2(Complex{Int32}(x)), u, v)
end

10-element CuArray{Int32,1}:
 12682
  5821
 12676
 19061
  6025
   433
  8609
   205
  2125
 11169