Connecting PyTorch sparse tensors with MLIR

aartbik: <p>Greetings! First time poster, here and a relative newcomer to PyTorch (but already loving it!). I am Aart Bik, and I am the tech lead of the <a href="https://developers.google.com/mlir-sparsifier" rel="noopener nofollow ugc">MLIR Sparsifier team at Google</a>. You can find more details about this team following the link, but in a nutshell, we think that sparsity should be a property, viz a type, and not a tedious implementation concern, and that a “sparse compiler”, aka sparsifier, should deal with all low level details of exploiting sparsity. These ideas have been fully worked out in MLIR for a CPU pipeline, with some GPU acceleration ramping up as well.</p>
<p>As such, we are very interested connecting a machine learning “front end” with our MLIR sparsifier “backend” (please forgive the very broad way of using those terms), and <em>torch.sparse</em> seems to fit our philosophy quite well, keeping the operator semantics (more or less) separate from the actual sparsity of their operands.</p>
<p>We have started to explore whether we can connect the <em>torch.sparse</em> data types with the MLIR sparse tensor types (the latter providing a superset of the former, since we use a <a href="https://mlir.llvm.org/docs/Dialects/SparseTensorOps/#sparsetensorencodingattr" rel="noopener nofollow ugc">TACO-flavored way</a> of defining sparse tensor types). For this, <em>torch-mlir</em> seems to make the most sense for us.</p>
<p>One of the tasks that probably requires minor changes inside the PyTorch framework itself is propagating the sparsity layout metadata somehow a bit further than currently seems to be the case, so that we can further propagate the layout while building the tensor IR in the torch dialect tensors of <em>torch-mlir</em>.</p>
<p>In any case, I just wanted to introduce myself and put this project on your radar screen. Please expect a lot of questions to follow. Also, your initial ideas are of course welcome already as well!</p>
<p>Nice “emeeting” you all!</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/1/0132ce8bcfc011fa05ce63452d54db80f4b8aa1a.gif" alt="mlir_sparsifier_name" data-base62-sha1="aBkuKYfRIVoYHYOomL8WhQfhMe" width="345" height="194" class="animated"></p>

aartbik: <p>And to make the discussion a bit more concrete, here is my first question.<br>
Given a very simple example:</p>
<pre><code class="lang-auto">class BikNet(torch.nn.Module):

  def __init__(self):
    super(BikNet, self).__init__()
    return

  def forward(self, x):
    return x.sum()
</code></pre>
<p>Then, building the traced graph for dense input works fine</p>
<pre><code class="lang-auto">dense_input = torch.ones(64, 64)
prog = torch.export.export(biknet, args=(dense_input,))
</code></pre>
<p>which yields</p>
<pre><code class="lang-auto">ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, l_x_: "f32[64, 64]"):
            # File: biknet.py:27, code: return x.sum()
            sum_1: "f32[]" = torch.ops.aten.sum.default(l_x_);  l_x_ = None
            return (sum_1,)
            
Graph signature: ExportGraphSignature(input_specs=[InputSpec(kind=&lt;InputKind.USER_INPUT: 1&gt;, arg=TensorArgument(name='l_x_'), target=None)], output_specs=[OutputSpec(kind=&lt;OutputKind.USER_OUTPUT: 1&gt;, arg=TensorArgument(name='sum_1'), target=None)])
Range constraints: {}
</code></pre>
<p>However, building the traced graph fails for sparse input</p>
<pre><code class="lang-auto">sparse_input = dense_input.to_sparse_csr()
prog =  torch.export.export(biknet, args=(sparse_input,))
</code></pre>
<p>throws an exception in builder.py when trying to clone.</p>
<pre><code class="lang-auto">torch._dynamo.exc.InternalTorchDynamoError: Sparse CSR tensors do not have strides
</code></pre>
<p>It would already go a long way if the traced graph builder would actually yield the same IR as shown above, but with the tensor input marked with the <em>torch.sparse_csr</em> layout. From there, we would be able to propagate the sparse type all the way down to MLIR.</p>
<p>Note that there is the subtlety that at runtime, sparse tensor input parameters of course have a 1:N relation with the actual passed in array arguments for positions, indices, and values (and same for a potential sparse output). But that is something we will have to deal with later (and we have done that in the past for e.g. Sparse JAX). For now, just getting the traced graph with sparse types would be really helpful.</p>
<p>Any comments on whether that is something that can be easily done? Or perhaps this is not the right direction, then I would love to hear why not as well.</p>

aartbik: <p>Note that I also opened <a href="https://github.com/pytorch/pytorch/issues/117188" rel="noopener nofollow ugc">feature request 117188</a> for this.</p>
<p>I hacked a bit around and was able to get a traced graph with the information I am (for now) looking for. In order to do this I had to</p>
<ol>
<li>add an extra layout field to <code>torch.fx.passes.shape_prop.TensorMetadata</code></li>
<li>add an extra layout field to <code>torch._subclasses.fake_tensor.FakeTensor</code></li>
<li>(obviously) modify printing to include those fields</li>
<li>hack my way around to propagate the sparse layout (very dirty solution)</li>
</ol>
<p>With that, I am able to generate the following traced graph.</p>
<pre><code class="lang-auto">ExportedProgram:
    class GraphModule(torch.nn.Module):
        def forward(self, l_x_: "f32[64, 64]:torch.sparse_csr"):    # ADDED!
            # File: biknet.py:27, code: return x.sum()
            sum_1: "f32[]" = torch.ops.aten.sum.default(l_x_);  l_x_ = None
            return (sum_1,)
           
Graph signature: ExportGraphSignature(
  input_specs=[
      InputSpec(
           kind=&lt;InputKind.USER_INPUT: 1&gt;,
           arg=TensorArgument(name='l_x_'),
           target=None,
           layout=torch.sparse_csr)       # ADDED!
  ],
  output_specs=[
     OutputSpec(
         kind=&lt;OutputKind.USER_OUTPUT: 1&gt;,
         arg=TensorArgument(name='sum_1'),
         target=None)
 ])
Range constraints: {}
</code></pre>
<p>With this information, I can prototype propagating this further in <code>torch-mlir</code> and then connecting this with the MLIR Sparsifier.</p>
<p>Stay tuned! I will report back here how that goes.</p>

smth: <p><a class="mention" href="/u/aartbik">@aartbik</a> I think you should post in the developer mailing list <a href="https://dev-discuss.pytorch.org/">https://dev-discuss.pytorch.org/</a> as your topic is really related to talking to the sparse devs and proposing making changes.</p>
<p><a href="http://discuss.pytorch.org">discuss.pytorch.org</a> is the user / support forum.</p>

aartbik: <p>Apologies for starting in the wrong forum and thanks for pointing me to the right place for this discussion. I hope to continue this discussion with <a href="https://dev-discuss.pytorch.org/t/connecting-pytorch-sparse-tensors-with-mlir/1803" rel="noopener nofollow ugc">a new posting in that forum</a>.</p>