The map(c, a.array) allocates an array, you can use the dot-broadcasting syntax instead of map c.(a.array), and then fuse everything in a single line, with the @. macro. The advantage is that it's...
The map(c, a.array) allocates an array, you can use the dot-broadcasting syntax instead of map c.(a.array), and then fuse everything in a single line, with the @. macro. The advantage is that it's clear from the syntax alone that all operations are element wise, so the compiler can easily fuse them in a single loop.
out = @. c(a.array) * c(b.array)
There's also no need to do a type assertion on the return type here, since you explicitly return a type already.
function mul(a::QArray{Int8}, b::QArray{Int8})
c = x -> convert(Int16, x)
out = @. c(a.array) * c(b.array)
QArray{Int16}(out, a.scale * b.scale)
end
Finally, you could make your function more generic:
function convert_and_mul(::Type{O}, a::QArray, b::QArray) where O <:Integer
c = x -> convert(O, x)
out = @. c(a.array) * c(b.array)
QArray{O}(out, a.scale * b.scale)
end
convert_and_mul(Int32, A, B)
Note that people would typically let the type promotion system to automatically deal with these kind of conversions (if you multiply an Int8 and an Int16 it returns an Int16, i.e. Base.promote_type(Int8, Int16) == Int16).
The
map(c, a.array)
allocates an array, you can use the dot-broadcasting syntax instead of mapc.(a.array)
, and then fuse everything in a single line, with the@.
macro. The advantage is that it's clear from the syntax alone that all operations are element wise, so the compiler can easily fuse them in a single loop.There's also no need to do a type assertion on the return type here, since you explicitly return a type already.
Finally, you could make your function more generic:
Note that people would typically let the type promotion system to automatically deal with these kind of conversions (if you multiply an Int8 and an Int16 it returns an Int16, i.e.
Base.promote_type(Int8, Int16) == Int16
).edit: fixed some bugs in my example