allow torch.bmm on nested_tensors of dim == 3 or (dim==4 and size(1)==1)

[puririshi98]

🚀 The feature, motivation and pitch

In the case of 4 dimensions, we should additionally check that the size(1)==1. This is because for us to use torch.bmm for segment_matmul, we split a tensor of shape (x, y, z) and we get a list of x tensors w/ shape (1,y, z). If I make a nested tensor out of these and put it into bmm it complains that it 'must be a 3D tensor' from these checks

example:

>>> torch.bmm(torch.nested.nested_tensor(list(torch.randn((10,5,5)).split(1))),torch.nested.nested_tensor(list(torch.ran
dn((10,5,5)).split(1))))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
RuntimeError: batch1 must be a 3D tensor

This is problematic because this is essential for our segment_matmul implementation we want the following behavior:

    r"""Performs dense-dense matrix multiplication according to segments along
    the first dimension of :obj:`inputs` as given by :obj:`ptr`, utilizing
    dedicated kernels that effectively parallelize over groups.

    .. code-block:: python

        inputs = torch.randn(8, 16)
        ptr = torch.tensor([0, 5, 8])
        other = torch.randn(2, 16, 32)

        out = pyg_lib.ops.segment_matmul(inputs, ptr, other)
        assert out.size() == (8, 32)
        assert out[0:5] == inputs[0:5] @ other[0]
        assert out[5:8] == inputs[5:8] @ other[1]

    Args:
        input (torch.Tensor): The left operand 2D matrix of shape
            :obj:`[N, K]`.
        ptr (torch.Tensor): Compressed vector of shape :obj:`[B + 1]`, holding
            the boundaries of segments.
        other (torch.Tensor): The right operand 3D tensor of shape
            :obj:`[B, K, M]`.

    Returns:
        torch.Tensor: The 2D output matrix of shape :obj:`[N, M]`.
    """

Here is the code I am using to allow torch>=1.14 to use torch.bmm instead of our custom kernel:

at::Tensor segment_matmul_kernel(const at::Tensor& input,
                                 const at::Tensor& ptr,
                                 const at::Tensor& other) {
  const auto size = pyg::utils::size_from_ptr(ptr).cpu();
  const auto sizes = at::IntArrayRef(size.data_ptr<int64_t>(), size.numel());
#if TORCH_VERSION_MINOR >= 14 or TORCH_VERSION_MAJOR > 1
  auto input_nested = at::_nested_tensor_from_tensor_list(
      input.contiguous().split_with_sizes(/*split_size=*/sizes, /*dim=*/0));
  auto other_nested = at::_nested_tensor_from_tensor_list(
                          other.contiguous().split(/*split_size=*/1, /*dim=*/0))
                          .squeeze(1);
  auto out_nested = at::native::bmm_nested_cuda(input_nested, other_nested);
  auto out = at::cat(out_nested.contiguous().unbind());
#else
  const auto out = input.new_empty({input.size(0), other.size(-1)});
  grouped_matmul_out_kernel(
      input.contiguous().split_with_sizes(/*split_size=*/sizes, /*dim=*/0),
      other.contiguous().split(/*split_size=*/1, /*dim=*/0),
      out.split_with_sizes(/*split_size=*/sizes, /*dim=*/0));
#endif

  return out;
}

In the case of torch>=1.14 we trigger a batch2 must be a 3D tensor error unless I call squeeze(1) on the output other.contiguous().split(/*split_size=*/1, /*dim=*/0)) . However this is a problem since when the nested tensor consists of many subtensors this slows things down dramatically compared to the torch < 1.14 pathway. This is because our custom grouped_matmul kernel functions even when other has shape (x, 1, y, z). This should be possible for the native implementation too and it would make it feasible for us to use. As you can see from this graph, even though our custom kernel is not faster than torch.bmm we end up with a much slower op using torch 1.14 vs torch < 1.14

Alternatives

I have tried implementing in python and cpp and in both cases using the new method is slower even though our kernel itself is not faster

Additional context

No response

cc @cpuhrsch @jbschlosser @bhosmer @drisspg @mikaylagawarecki

[cpuhrsch]

Hey @puririshi98 - we just landed a shortcut into an optimized bmm via torch.matmul (see #86856). Maybe matmul is what you're looking for here?

[puririshi98]

RuntimeError: matmul: both inputs must have the same rank

I get the above issue if I don't call squeeze(1) on other_nested

at::Tensor segment_matmul_kernel(const at::Tensor& input,
                                 const at::Tensor& ptr,
                                 const at::Tensor& other) {
  const auto size = pyg::utils::size_from_ptr(ptr).cpu();
  // TODO (matthias) Allow for other types than `int64_t`.
  const auto sizes = at::IntArrayRef(size.data_ptr<int64_t>(), size.numel());
#if TORCH_VERSION_MINOR >= 14 or TORCH_VERSION_MAJOR > 1
  auto input_nested = at::_nested_tensor_from_tensor_list(
      input.contiguous().split_with_sizes(/*split_size=*/sizes, /*dim=*/0));
  auto other_nested = at::_nested_tensor_from_tensor_list(
                          other.contiguous().split(/*split_size=*/1, /*dim=*/0));
  auto out_nested = at::matmul(input_nested, other_nested);
  auto out = at::cat(out_nested.contiguous().unbind());
#else
  const auto out = input.new_empty({input.size(0), other.size(-1)});
  // TODO (matthias) Better handle non-contiguous memory layouts.
  grouped_matmul_out_kernel(
      input.contiguous().split_with_sizes(/*split_size=*/sizes, /*dim=*/0),
      other.contiguous().split(/*split_size=*/1, /*dim=*/0),
      out.split_with_sizes(/*split_size=*/sizes, /*dim=*/0));
#endif

  return out;
}

but if I squeeze other_tensor then it works:

x.py::test_segment_matmul_autograd[cuda:0] [W NestedTensorImpl.cpp:177] Warning: The PyTorch API of nested tensors is in prototype stage and will change in the near future. (function operator())
PASSED
x.py::test_grouped_matmul_autograd[cuda:0] PASSED

at::matmul or at::native:bmm_nested_cuda both behave the same for me when given nested tensor inputs, need a solution that works without me needing to squeeze other_nested

As can be seen by the graph, the torch>=1.14 pathway is much slower even though the only difference is that squeeze, which our custom kernel does not require.

[cpuhrsch]

@puririshi98 - can you share some typical input shapes to matmul for your use case (without unsqueeze etc.) or some self-contained, short code in a gist for a microbenchmark that can act as a proxy for development.

[puririshi98]

https://github.com/puririshi98/rgcn_pyg_lib_forward_bench/blob/main/segmatmul_bench.py
this is the code I use to make the graph shown above:

git clone -b nested_tensors --recurse-submodules https://github.com/puririshi98/pyg-lib.git; cd pyg-lib; pip uninstall -y torch; pip3 install --pre torch --extra-index-url https://download.pytorch.org/whl/nightly/cu117; git checkout nested_tensors; git pull origin nested_tensors; python3 -m pip uninstall -y pyg-lib; python3 -m pip install -e .; git clone https://github.com/puririshi98/rgcn_pyg_lib_forward_bench.git; cd rgcn_pyg_lib_forward_bench; python3 segmatmul_bench.py; cd /workspace/pyg-lib; pip3 uninstall -y torch; pip3 install torch --extra-index-url https://download.pytorch.org/whl/cu117;
pip uninstall -y pyg-lib; python3 -m pip install -e .; cd rgcn_pyg_lib_forward_bench; python3 segmatmul_bench.py;

This is the code for the segment_matmul op we are benchmarking:

at::Tensor segment_matmul_kernel(const at::Tensor& input,
                                 const at::Tensor& ptr,
                                 const at::Tensor& other) {
  const auto size = pyg::utils::size_from_ptr(ptr).cpu();
  // TODO (matthias) Allow for other types than `int64_t`.
  const auto sizes = at::IntArrayRef(size.data_ptr<int64_t>(), size.numel());
#if TORCH_VERSION_MINOR >= 14 or TORCH_VERSION_MAJOR > 1
  auto input_nested = at::_nested_tensor_from_tensor_list(
      input.contiguous().split_with_sizes(/*split_size=*/sizes, /*dim=*/0));
  auto other_nested = at::_nested_tensor_from_tensor_list(
                          other.contiguous().split(/*split_size=*/1, /*dim=*/0)).squeeze(1);
  auto out_nested = at::matmul(input_nested, other_nested);
  auto out = at::cat(out_nested.contiguous().unbind());
#else
  const auto out = input.new_empty({input.size(0), other.size(-1)});
  grouped_matmul_out_kernel(
      input.contiguous().split_with_sizes(/*split_size=*/sizes, /*dim=*/0),
      other.contiguous().split(/*split_size=*/1, /*dim=*/0),
      out.split_with_sizes(/*split_size=*/sizes, /*dim=*/0));
#endif

  return out;
}

I would like the code in the >= 1.14 path way to be as fast as the <1.14 path way, and the only way I see this as being possible is if the bmm or matmul in >= 1.14 can handle not squeezing 'other'. Without squeezing the 'other' nested tensor will have 4 dimensions, (x, 1, y, z)) which currently is blocked by the dim==3 check, matmul or bmm should support tensors with a more flexible check: (dim == 3) or (dim == 4 and size(1) == 1) since torch.split gives tensors with that shape

[cpuhrsch]

@puririshi98

Instead of

    inputs = torch.randn((num_nodes_per_type * num_types, n_feats)).cuda()
    ptr = torch.tensor(list(range(0, (num_types + 1) * num_nodes_per_type, num_nodes_per_type))).cuda()
    other = torch.randn((num_types, n_feats, out_feats), requires_grad=True).cuda()
    
    for i in range(60):
        if i==10:
            since=time.time()
        pyg_lib.ops.segment_matmul(inputs, ptr, other)
    times.append((time.time()-since)/50.0)
    print("Avg time for", num_types, '=', times[-1])

we could look into enabling (I'm making assumptions about what ptr is)

    inputs = torch.randn((num_nodes_per_type * num_types, n_feats), device="cuda")
    ptr = list(range(0, (num_types + 1) * num_nodes_per_type, num_nodes_per_type))
    nt = torch.nested.nested_tensor([inputs[ptr[i]:ptr[i+1]] for i in range(len(num_types + 1))])
    other = torch.randn((num_types, n_feats, out_feats), requires_grad=True, device="cuda")

    result = torch.matmul(nt, other) # instead of segment_matmul

What do you think about that?

[puririshi98]

is it fine for ptr to be a cuda tensor or atleast a cpu tensor or would this cause issues?

[cpuhrsch]

@puririshi98 - It's fine if it's a CPU Tensor. It really just comes down to generating that list of Tensors for the nested_tensor constructor.

[puririshi98]

okay thats fine then

[puririshi98]

this is the code I will use with your new suggested changes:

at::Tensor segment_matmul_kernel(const at::Tensor& input,
                                 const at::Tensor& ptr,
                                 const at::Tensor& other) {
  const auto size = pyg::utils::size_from_ptr(ptr).cpu();
  // TODO (matthias) Allow for other types than `int64_t`.
  const auto sizes = at::IntArrayRef(size.data_ptr<int64_t>(), size.numel());
#if TORCH_VERSION_MINOR >= 14 or TORCH_VERSION_MAJOR > 1
  auto input_nested = at::_nested_tensor_from_tensor_list(
      input.contiguous().split_with_sizes(/*split_size=*/sizes, /*dim=*/0));
  auto out_nested = at::matmul(input_nested, other.contiguous());
  auto out = at::cat(out_nested.contiguous().unbind());
#else
  const auto out = input.new_empty({input.size(0), other.size(-1)});
  // TODO (matthias) Better handle non-contiguous memory layouts.
  grouped_matmul_out_kernel(
      input.contiguous().split_with_sizes(/*split_size=*/sizes, /*dim=*/0),
      other.contiguous().split(/*split_size=*/1, /*dim=*/0),
      out.split_with_sizes(/*split_size=*/sizes, /*dim=*/0));
#endif

  return out;
}

I am assuming the output of the matmul will be a nested tensor, is this correct?

[cpuhrsch]

Yep the output is a nestedtensor

[cpuhrsch]

@puririshi98 - Your main performance issues are from doing the construction and then unbinding the resulting. Construction and then unbind is very expensive. If you time it without the construction and unbind it'll be much faster and potentially competitive for your use case. This runs on our latest nightly. Note that I'm constructing a NestedTensor from other via unbind + nested_tensor.

import torch
import time
import os
os.environ['NVIDIA_TF32_OVERRIDE'] = '0'
times = []
num_nodes_per_type = 10000
n_feats = 128
out_feats = 64

def benchmark_torch_function(iters, f, *args, **kwargs):
    f(*args, **kwargs) # warmup
    torch.cuda.synchronize()
    start_event = torch.cuda.Event(enable_timing=True)
    end_event = torch.cuda.Event(enable_timing=True)
    start_event.record()
    for i in range(iters):
        f(*args, **kwargs)
    end_event.record()
    torch.cuda.synchronize()
    return start_event.elapsed_time(end_event)


for num_types in [4, 8, 16, 32, 64, 128, 256, 512, 1024]:
    inputs = torch.randn((num_nodes_per_type * num_types, n_feats)).cuda()
    ptr = list(range(0, (num_types + 1) * num_nodes_per_type, num_nodes_per_type))
    nt = torch.nested.nested_tensor([inputs[ptr[i]:ptr[i+1]] for i in range(len(ptr) - 1)], device='cuda')
    other = torch.randn((num_types, n_feats, out_feats), requires_grad=True).cuda()
    nt_other = torch.nested.nested_tensor(list(other.unbind()), device='cuda')
    time = benchmark_torch_function(1, torch.bmm, nt, nt_other)
    print(f"num_types: {num_types} time: {time}")

On a meta note, you could use NestedTensor as your internal representation and thus avoid this costly conversion. You'd also get various other operators as a side effect. However, NestedTensor still is in its prototype phase so there will be some breakages. From a technical perspective however, does that sound doable / interesting?