Inconsistent DirectSolver interface.

[anadodik]

Hi everyone! I've noticed a small inconsistency in the sparse solver interface, namely when it comes to using it with a 2D tensor where the last dimension is 1.

Given a 3x3 lhs sparse matrix and a 3x2 rhs dense matrix:

a = torch.sparse_coo_tensor([[0, 1, 2], [0, 1, 2]], [1.0, 1.0, 1.0], device="cuda").to_sparse_csr()
b = torch.randn([3, 2], device="cuda")

we can set up a solver for with:

nvsparse.DirectSolver(a, b.mT.contiguous().mT)

This works fine. However, for a 3x1 rhs:

b = torch.randn([3,1], device="cuda")

setting up the solver in the same way breaks with:

TypeError("The RHS tensor([[1.],\n        [1.],\n        [1.]], device='cuda:0') must be a matrix or vector with col-major layout, and for implicitly-batched RHS (N-D >= 3), each matrix sample must have col-major layout (the second dimension from the end must have unit stride.")

It would be great not to have to special-case this.

Best,
Ana

[yangcal]

Thanks for reporting. After some investigation, I believe the inconsistency is coming from the array package, not from nvmath. Using your snippet as an example, the follow two code paths should both create an F ordered matrix with shape (3, 1) :

# Path 1
b = torch.randn([3, 1], device="cuda").mT.contiguous().mT
print(b.shape, b.strides())

# Path 2
b = torch.randn([1, 3], device="cuda").mT
print(b.shape, b.strides())

I've observed that for the first path, the returned b strides is (1, 1) while the second path gives strides of (1, 3).

In this particular case, the second dimension is 1 so it doesn't matter whether the stride for the second dimension is specified as 1 or 3. Pedantically it should be 3. In fact I've also observed such inconsistency in other packages like numpy & cupy. One example is that:

shape = (3, 1)
b0 = np.asarray(np.ones(shape), order="F")
b1 = np.ones(shape, order="F")

print(b0.strides, b1.strides) # strides (8, 8) & strides (8, 24), it's different from torch due to a scale by itemsize

Going back to your case, you may consider one of the following as WAR:

1. provide b as a vector when the column size is 1
2. Make sure the array has the pedantically correct strides, like we did with Path 2 above. Path 2 also helps avoid a copy process.
3. If you're comfortable editing the source code, you may update the following block to something like below:

def check_dense_tensor_layout(shape: Sequence[int], strides: Sequence[int], *, explicitly_batched=None):
    """
    Check that the dense vector or matrix (each sample matrix in a N-D tensor) is
    in col-major format.
    """

    num_dimensions = len(shape)
    assert explicitly_batched is not None and len(strides) == num_dimensions, "Internal Error."

    # For explicitly batched matrices, each sample must be a matrix or vector.
    if explicitly_batched and num_dimensions > 2:
        return False

    first = int(num_dimensions > 1)

    # Only column-major matrices are currently supported.
    is_col_major = strides[-1 - first] == 1

    # Check that the LD doesn't lead to overlapping memory.
    # For batched matrices, LD has to be equal to the first matrix dimension.
    comparison_op_for_ld = operator.eq if num_dimensions > 2 else operator.ge

    def compare_ld(strides, shape):
        if shape[-1] == 1 and strides[-1] == 1:
            # NOTE: a special case for F-order matrix with shape (..., m, 1)
            # Depending on how the array is created, the strides property
            # provided by the package can become either:
            #    (..., 1, 1) or (..., 1, m)
            return True
        return comparison_op_for_ld(strides[-1], shape[-2])

    if num_dimensions > 1:
        return is_col_major and compare_ld(strides, shape)

    return is_col_major

I'm reporting back to the team and we will decide whether to include this fix in our next release

[anadodik]

Hi, thanks for getting back to me! I’ve managed to find a work-around at the time, but figured you might be interested in addressing this on your end.

From what you’re saying, the memory layout is actually fine, it’s just an inconsistency with the stride computation internal to PyTorch? My 2c is that it would make sense to patch this at the library level; most people will end up having to work around this on the user side and there seems to be no real performance penalty from supporting it on your end, from what I understand.

Thank you for your consideration!

[yangcal]

the memory layout is actually fine, it’s just an inconsistency with the stride computation internal to PyTorch?

Yes the memory layout is fine. The stride for the dimension = 1 is moot as there was no need to access data along that dimension. I think for frameworks, when they get a transposed view from the original C ordered array with strides (1, 1), they just permute the strides so it ended up being (1, 1) too.

it would make sense to patch this at the library level

Yes. The affected scope may go beyond the DirectSolver class and we will discuss what to include in our next release.