[BUG] Reduce operations give partly wrong results with 3D tensors.

[AtomicVar]

Describe the bug
I tried to use matx::sum on 3D tensors, which are common in Deep Learning (batch_size, sequence_len, embedding_size). The results are partly wrong. I don't know if it's my mis-using or it's not supported yet.

To Reproduce
Steps to reproduce the behavior:

1. Create a 3D tensor mat with shape {2, 3, 2}:

auto vec = matx::make_tensor<float, 1>({12});
vec.SetVals({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
auto mat = vec.View({2, 3, 2});
// now mat is
// [
//   [[1, 2],
//    [3, 4],
//    [5, 6]],
//   [[7, 8],
//    [9, 10],
//    [11, 12]]
// ]

2. Perform sum operation on each row:

auto mat_sum = matx::make_tensor<float, 1>({6});
matx::sum(mat_sum, mat);
mat_sum.Print();

3. The result is:

000000: 3.0000 
000001: 7.0000 
000002: 11.0000 
000003: 33554528.0000 
000004: 0.0000 
000005: 0.0000

Expected behavior
The Results should be:

000000: 3.0000 
000001: 7.0000 
000002: 11.0000 
000003: 15.0000 
000004: 19.0000 
000005: 23.0000

System details (please complete the following information):

• OS: Ubuntu 20.04
• CUDA version: CUDA 11.4
• g++ version: 9.3.0

[AtomicVar]

BTW, it works great with 1D and 2D tensors :)

[cliffburdick]

thanks for the report @ZJUGuoShuai. we recently moved most reductions to cub, so I will check for a regression

[cliffburdick]

Hi @ZJUGuoShuai, it looks like the order isn't quite what we expect, and we should have errored. The way our reductions work is the output tensor must have the same inner ranks as the input tensor. In other words,

{2,3,2} -> {2} // valid
{2,3,2} -> {6} // not valid

By going from a rank 3 to a rank one you are asking it to sum over the outer 2 dimensions, and the final output would have two values that are the sum of the two 3x2 matrices:

auto mat = vec.View({2, 3, 2});
auto mat_sum = matx::make_tensor<float, 1>({2});
matx::sum(mat_sum, mat);

000000: 21.0000
000001: 57.0000

It looks like instead you want to just sum over 6 rows of 2 columns each, which would be:

  auto vec = matx::make_tensor<float, 1>({12});
vec.SetVals({1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12});
auto mat = vec.View({6,2});
auto mat_sum = matx::make_tensor<float, 1>({6});
matx::sum(mat_sum, mat);
mat_sum.Print();

000000: 3.0000
000001: 7.0000
000002: 11.0000
000003: 15.0000
000004: 19.0000
000005: 23.0000

I also had to make a small change since we did have a regression when using CUB for row-wise reductions, so that's back to take a slower path for now and will be fixed next week. You should be able to get the correct answer now with the latest code and the example above.

[AtomicVar]

@cliffburdick Thanks, now I know how to use Reduce operations:

the output tensor must have the same inner ranks as the input tensor