addr bug for complex and inefficient implementation

[ngimel]

addr implementation here
https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/LinearAlgebra.cpp#L147-L159
does explicit multiplications by alpha and beta even though by default alpha and beta are 1, so no multiplication is necessary. This degrades performance and memory usage.
In addition, line 55 does invalid cast, so invocation with complex beta results in an error:

a=torch.randn(3, dtype=torch.complex64)                                                                                                                                                                
b=torch.randn(3, dtype=torch.complex64)                                                                                                                                                                
c=torch.randn(3,3, dtype=torch.complex64)   
c.addr(a,b, beta=0+0.1j)

RuntimeError: value cannot be converted to type double without overflow: (0,0.1)

related: #44364
cc @vishwakftw @jianyuh @nikitaved @pearu @mruberry @heitorschueroff
cc @ivannz

[RockingJavaBean]

@mruberry I would like to help fix the performance and memory issues, as well as the bug when inputs are complex numbers.

Please kindly confirm which option we should adopt.

• implement torch.addr as a composite op and prevent unnecessary multiplications when alpha or beta is the default value
• implement torch.addr based on BLAS libraries.

CC @ngimel @IvanKobzarev

[ivannz]

In my opinion, the implementation of torch.addr(input, vec1, vec2, beta=1, alpha=1, out=None) as a composite op by design incurs an impractical allocation of the product vec1 * vec2^T.

As a standalone outer product, produced by torch.outer, this is fine, because the caller presumably knows and agrees that the output is a rank-one matrix. However as an inplace operation, all of which are aimed at conserving memory, it is tacitly assumed that the implementation does not require extra storage beyond that of the arguments (since it is reasonable to expect it to be implemented as out[i,j] = out[i,j] * beta + vec1[i] * vec2[j] * alpha in an i-j loop).

Hence, I think, torch.outer(vec1, vec2, out=None) could be implemented as torch.addr(input, vec1, vec2, beta=0, alpha=1, out=None), wherein input is completely ignored if beta is exactly zero. This way, I believe, both the maintainability and versatility objectives would be feasible: torch.outer would always be a wrapper around torch.addr, while torch.addr, being the core procedure, would be responsible for dispatching to the correct dtype and device pathway.

[ngimel]

I agree that current addr implementation allocates unnecessary intermediate buffer, however, I'm not convinced that bringing back blas/cublas bindings for it is good idea.

1. they don't support bfloat16 and half datatypes.
2. they are cumbersome to use for complex
3. they are cumbersome to use with ROCm
4. as shown in the benchmarks in original port, outer is significantly faster implemented as a pointwise multiply, and addr suffers from needlessly inefficient implementation due to extra scalar multiplies, as mentioned in this issue.

addr can be efficiently implemented not relying on cublas, by using TensorIterator-based kernels (cpu_kernel_vec and gpu_kernel) To do it, one needs to create TensorIterator with correct inputs and declare cpu/gpu stubs specifying the correct pointwise operation for the kernels, you can look at how many pointwise operations are implemented to get the idea.

[RockingJavaBean]

@ngimel Thanks for the kind and detailed suggestions, I will start to implement torch.addr using TensorIterator-based kernels.