WIP [torch-mlir][sparse] replace temporary metadata/export with upstream PyTorch

python/torch_mlir/extras/fx_importer.py

@@ -264,64 +264,56 @@
         torch.ops.aten.sym_numel.default: torch.ops.aten.numel.default,
     }
 
-
-@dataclass(frozen=True)
-class SparsityMeta:
-    """
-    Class for keeping track of sparsity meta data.
-
-    NOTE: this will be fully replaced by
-          torch.fx.passes.shape_prop.SparseTensorMetadata
-    """
-
-    layout: torch.layout
-    batch_dim: int
-    sparse_dim: int
-    dense_dim: int
-    blocksize: Optional[Tuple[int, int]]
-    pos_dtype: torch.dtype
-    crd_dtype: torch.dtype
+SPARSE_LAYOUTS = [
+    torch.sparse_coo,
+    torch.sparse_csr,
+    torch.sparse_csc,
+    torch.sparse_bsr,
+    torch.sparse_bsc,
+]
 
 
-def sparsity_encoding(shape: torch.Size, sparsity: SparsityMeta) -> str:
-    """Returns sparse tensor encoding for the given sparse layout as string."""
-    assert sparsity is not None
+def sparsity_encoding(t: torch.Tensor) -> str:
+    """Returns sparse tensor encoding for the given sparse tensor as string."""
 
     # Sparse tensors have the form
     #   [ <batch_dimensions> , <sparse_dimensions>, <dense_dimensions> ]
     # which map directly to MLIR types.
-    batch_dim, sparse_dim, dense_dim = (
-        sparsity.batch_dim,
-        sparsity.sparse_dim,
-        sparsity.dense_dim,
+    dim, batch_dim, sparse_dim, dense_dim = (
+        t.ndim,
+        t.ndim - t.sparse_dim() - t.dense_dim(),
+        t.sparse_dim(),
+        t.dense_dim(),
     )
-    dim = batch_dim + sparse_dim + dense_dim
-    assert dim == len(shape)
-    blocksize = sparsity.blocksize
-
     dims = ",".join(f"d{d}" for d in range(dim))
 
-    if sparsity.layout is torch.sparse_coo:
-        assert sparse_dim >= 2 and blocksize is None
+    if t.layout is torch.sparse_coo:
+        assert sparse_dim >= 2
         trail_dim = batch_dim + sparse_dim - 1
         coords = ",".join(
             f"d{d}:singleton(nonunique,soa)" for d in range(batch_dim + 1, trail_dim)
         )
         sep = "," if sparse_dim > 2 else ""
         lvls = f"d{batch_dim}:compressed(nonunique),{coords}{sep}d{trail_dim}:singleton(soa)"
-    elif sparsity.layout is torch.sparse_csr:
-        assert sparse_dim == 2 and blocksize is None
+        idx_dtype = t._indices().dtype  # supports uncoalesced COO tensors
+    elif t.layout is torch.sparse_csr:
+        assert sparse_dim == 2
         lvls = f"d{batch_dim}:dense,d{batch_dim+1}:compressed"
-    elif sparsity.layout is torch.sparse_csc:
-        assert sparse_dim == 2 and blocksize is None
+        idx_dtype = t.col_indices().dtype
+    elif t.layout is torch.sparse_csc:
+        assert sparse_dim == 2
         lvls = f"d{batch_dim+1}:dense,d{batch_dim}:compressed"
+        idx_dtype = t.row_indices().dtype
     else:
-        assert sparse_dim == 2 and blocksize is not None
-        if sparsity.layout is torch.sparse_bsr:
+        assert sparse_dim == 2
+        blocksize = t.values().shape[batch_dim + 1 : batch_dim + 3]
+        if t.layout is torch.sparse_bsr:
             i, j = batch_dim, batch_dim + 1
+            idx_dtype = t.col_indices().dtype
         else:
-            assert sparsity.layout is torch.sparse_bsc
+            assert t.layout is torch.sparse_bsc
             j, i = batch_dim, batch_dim + 1
+            idx_dtype = t.row_indices().dtype
         m, n = blocksize
         lvls = (
             f"d{i} floordiv {m}:dense,d{j} floordiv {n}:compressed,"
@@ -336,8 +328,7 @@ def sparsity_encoding(shape: torch.Size, sparsity: SparsityMeta) -> str:
         dense = ",".join(f"d{d}:dense" for d in range(batch_dim + sparse_dim, dim))
         lvls = f"{lvls},{dense}"
 
-    posw = torch.iinfo(sparsity.pos_dtype).bits
-    crdw = torch.iinfo(sparsity.crd_dtype).bits
+    posw = crdw = torch.iinfo(idx_dtype).bits
     return f"#sparse_tensor.encoding<{{map=({dims})->({lvls}),posWidth={posw},crdWidth={crdw}}}>"
 
 
@@ -904,16 +895,15 @@ def get_vtensor_type(
         shape: torch.Size,
         dtype: torch.dtype,
         *,
-        sparsity: Optional[SparsityMeta] = None,
+        val: Optional[torch.Tensor] = None,
         mutable: bool = False,
     ):
         """Return IrType for !torch.vtensor with the given shape and dtype"""
         stem = "torch.tensor" if mutable else "torch.vtensor"
         shape_asm = self.format_asm_shape(shape)
         mlir_dtype = str(self.dtype_to_type(dtype))
-        if sparsity is not None:
-            encoding = sparsity_encoding(shape, sparsity)
-            assert encoding is not None
+        if val is not None and val.layout in SPARSE_LAYOUTS:
+            encoding = sparsity_encoding(val)
             return IrType.parse(
                 f"!{stem}<[{shape_asm}],{str(mlir_dtype)},{encoding}>",
                 context=self._c,
@@ -926,21 +916,19 @@ def node_val_to_type(self, node: torch_fx.Node, *, mutable: bool = False) -> IrT
         try:
             tensor_meta = node.meta.get("tensor_meta")
             val = node.meta.get("val")
-            sparsity = node.meta.get("sparsity", None)
         except KeyError as e:
             raise RuntimeError(
                 f"FIXME: Illegal access to torch.fx.Node.meta: {e} ({node.meta.keys()} : {node.meta})"
             )
         return self.value_info_to_type(
-            val, tensor_meta=tensor_meta, sparsity=sparsity, mutable=mutable
+            val, tensor_meta=tensor_meta, mutable=mutable
         )
 
     def value_info_to_type(
         self,
         val,
         *,
         tensor_meta: Optional[TensorMetadata] = None,
-        sparsity=None,
         mutable: bool = False,
     ):
         if tensor_meta is not None:
@@ -953,14 +941,14 @@ def value_info_to_type(
                 )
             else:
                 return self.tensor_metadata_to_type(
-                    tensor_meta, sparsity=sparsity, mutable=mutable
+                    tensor_meta, val=val, mutable=mutable
                 )
         elif val is not None:
             # some nodes with symbolic inputs pass a 'val' attribute rather than
             # tensor_meta
             if isinstance(val, TorchFakeTensor):
                 return self.get_vtensor_type(
-                    val.size(), val.dtype, sparsity=sparsity, mutable=mutable
+                    val.size(), val.dtype, val=val, mutable=mutable
                 )
 
         # Note that None is a valid scalar here, so it is important that this
@@ -978,18 +966,18 @@ def tensor_metadata_to_type(
         self,
         tm: TensorMetadata,
         *,
-        sparsity: Optional[SparsityMeta] = None,
+        val: Optional[torch.Tensor] = None,
         mutable: bool = False,
     ) -> IrType:
         tm_shape = tuple(
             item.node if is_symbolic(item) else item for item in list(tm.shape)
         )
 
-        key = (tm_shape, tm.dtype, sparsity, mutable)
+        key = (tm_shape, tm.dtype, val, mutable)
         t = self._tensor_metadata_cache.get(key)
         if t is None:
             t = self.get_vtensor_type(
-                tm.shape, tm.dtype, sparsity=sparsity, mutable=mutable
+                tm.shape, tm.dtype, val=val, mutable=mutable
             )
             self._tensor_metadata_cache[key] = t
         return t

test/python/fx_importer/sparse_test.py

@@ -8,12 +8,10 @@
 from typing import Any, Callable, Optional, Tuple, Dict
 
 import torch
-import torch.export
 import torch.nn as nn
 import numpy as np
 
 from torch_mlir.extras.fx_importer import FxImporter
-from torch_mlir.extras.fx_importer import SparsityMeta
 from torch_mlir import ir
 from torch_mlir.dialects import torch as torch_d
 from torch_mlir.compiler_utils import run_pipeline_with_repro_report
@@ -22,114 +20,12 @@
 )
 
 
-# All sparse layouts currently supported in torch.sparse.
-SPARSE_LAYOUTS = [
-    torch.sparse_coo,
-    torch.sparse_csr,
-    torch.sparse_csc,
-    torch.sparse_bsr,
-    torch.sparse_bsc,
-]
-
-
-def sparse_metadata(a: torch.Tensor) -> SparsityMeta:
-    """
-    Returns a meta data tuple for the given sparse tensor.
-
-    NOTE: this will be fully replaced by fx graph SparseTensorMetadata
-    """
-    sparse_dim = a.sparse_dim()
-    dense_dim = a.dense_dim()
-    batch_dim = a.ndim - dense_dim - sparse_dim
-    blocksize = None
-    if a.layout is torch.sparse_coo:
-        return SparsityMeta(
-            a.layout,
-            batch_dim,
-            sparse_dim,
-            dense_dim,
-            blocksize,
-            a._indices().dtype,
-            a._indices().dtype,
-        )
-    elif a.layout is torch.sparse_csr or a.layout is torch.sparse_bsr:
-        if a.layout is torch.sparse_bsr:
-            blocksize = a.values().shape[batch_dim + 1 : batch_dim + 3]
-        return SparsityMeta(
-            a.layout,
-            batch_dim,
-            sparse_dim,
-            dense_dim,
-            blocksize,
-            a.crow_indices().dtype,
-            a.col_indices().dtype,
-        )
-    elif a.layout is torch.sparse_csc or a.layout is torch.sparse_bsc:
-        if a.layout is torch.sparse_bsc:
-            blocksize = a.values().shape[batch_dim + 1 : batch_dim + 3]
-        return SparsityMeta(
-            a.layout,
-            batch_dim,
-            sparse_dim,
-            dense_dim,
-            blocksize,
-            a.ccol_indices().dtype,
-            a.row_indices().dtype,
-        )
-    else:
-        raise RuntimeError(f"Unsupported sparse layout for {a}")
-
-
-def sparse_export(
-    f: Callable, args: Tuple[Any, ...], kwargs: Optional[Dict[str, Any]] = None
-) -> torch.export.ExportedProgram:
-    """
-    This is a ***temporary*** wrapper around `torch.export.export`
-    that eventually should be removed and simply replaced by the
-    standard API for exporting traced graphs.
-
-    But until issue
-
-      https://github.com/pytorch/pytorch/pull/117907
-
-    is addressed, this wrapper provides support for the sparse
-    tensor types by first converting all operands to dense tensors,
-    building the traced graph as for the dense case, and then
-    annotation sparse parameters with their actual sparse layout
-    attributes. This temporary solution accelerates testing
-    torch-mlir with PyTorch sparse tensors until the issue is
-    resolved.
-    """
-    # Convert all arguments to dense.
-    dargs = tuple(a.to_dense() if a.layout in SPARSE_LAYOUTS else a for a in args)
-    mask = [a.layout in SPARSE_LAYOUTS for a in args]
-    # Build the regular FX traced graph with only dense arguments
-    # (the current version would crash otherwise, see issue above).
-    prog = torch.export.export(f, dargs, kwargs)
-    # Annotate sparse arguments in the graph. Note that we currently
-    # only account for sparsity defined by the user inputs to the model.
-    # TODO: support sparsity in model parameters (weights, biases)
-    # TODO: propagate sparsity into the layers
-    specs = prog.graph_signature.input_specs
-    alen = len(specs)
-    k = 0
-    for i, node in enumerate(prog.graph.nodes):
-        if i >= alen:
-            break
-        spec = specs[i]
-        if spec.kind is torch.export.graph_signature.InputKind.USER_INPUT:
-            if mask[k]:
-                node.meta["sparsity"] = sparse_metadata(args[k])
-            k = k + 1
-    return prog
-
-
 def export_and_import(f, *args, **kwargs):
     """This method implements Stella's importer, stripped down to essentials."""
     context = ir.Context()
     torch_d.register_dialect(context)
     fx_importer = FxImporter(context=context)
-    prog = sparse_export(f, args, kwargs)
+    prog = torch.export.export(f, args, kwargs)
     fx_importer.import_frozen_program(prog)
     return fx_importer.module
 
@@ -346,7 +242,6 @@ def forward(self, x, y):
     print("torch.mlir")
     print(res2)
 
-
 @run
 # CHECK-LABEL: test_sparse_eltwise
 # CHECK:       #[[$BCSR:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 : batch, d1 : dense, d2 : compressed), posWidth = 64, crdWidth = 64 }>
@@ -367,23 +262,22 @@ def forward(self, x, y):
 # CHECK:        col_indices=tensor([0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1,
 # CHECK:                            2, 3, 0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3]),
 # CHECK:        values=tensor({{\[}}[ -1.,  -2.],
-# CHECK:                            [ -3.,  -4.],
 #                                   ...
 # CHECK:                            [-63., -64.]{{\]}}), size=(8, 4, 2), nnz=32,
-# CHECK:                       layout=torch.sparse_csr)
+# CHECK:                      layout=torch.sparse_csr)
 # CHECK:        torch.mlir
 # CHECK:        {{\[\[}}[ -1.  -2.]
 # CHECK:                [ -3.  -4.]
 #                       ...
 # CHECK:                [-61. -62.]
 # CHECK:                [-63. -64.]{{\]\]}}
-#
 # CHECK:        torch.mlir.batch
 # CHECK:        {{\[\[}}[ -1.  -2.]
 # CHECK:                [ -3.  -4.]
 #                       ...
 # CHECK:                [-61. -62.]
 # CHECK:                [-63. -64.]{{\]\]}}
+#
 def test_sparse_eltwise():
     class EltNet(torch.nn.Module):
         def __init__(self):
@@ -425,12 +319,23 @@ def forward(self, x):
 # CHECK-LABEL: test_sparse_coo3
 # CHECK:       #[[$COO3:.*]] = #sparse_tensor.encoding<{ map = (d0, d1, d2) -> (d0 : compressed(nonunique), d1 : singleton(nonunique, soa), d2 : singleton(soa)), posWidth = 64, crdWidth = 64 }>
 # CHECK:       func.func @main(
-# CHECK-SAME:    %[[A:.*]]: !torch.vtensor<[10,20,30],f64,#sparse>) -> !torch.vtensor<[10,20,30],f64> {
-# CHECK:         %[[R:.*]] = torch.aten.relu %[[A]] : !torch.vtensor<[10,20,30],f64,#sparse> -> !torch.vtensor<[10,20,30],f64>
-# CHECK:         return %[[R]] : !torch.vtensor<[10,20,30],f64>
+# CHECK-SAME:    %[[A:.*]]: !torch.vtensor<[10,20,30],f64,#sparse>) -> !torch.vtensor<[10,20,30],f64,#sparse> {
+# CHECK:         %[[R:.*]] = torch.aten.relu %[[A]] : !torch.vtensor<[10,20,30],f64,#sparse> -> !torch.vtensor<[10,20,30],f64,#sparse>
+# CHECK:         return %[[R]] : !torch.vtensor<[10,20,30],f64,#sparse>
 # CHECK:       }
 #
-# TODO: make sure sparsity propagates through relu into the output and test actual JIT output
+# CHECK: torch.sparse
+# CHECK: tensor(indices=tensor([[ 0,  1,  1,  4,  9,  9],
+# CHECK:                        [ 0,  1,  1,  5, 19, 19],
+# CHECK:                        [ 0,  1,  3,  6, 28, 29]]),
+# CHECK: values=tensor([   0.,    0.,    1.,    2.,    3., 1000.]),
+# CHECK: size=(10, 20, 30), nnz=6, dtype=torch.float64, layout=torch.sparse_coo)
+# CHECK: torch.mlir
+# CHECK: tensor(indices=tensor([[ 0,  1,  1,  4,  9,  9],
+# CHECK:                        [ 0,  1,  1,  5, 19, 19],
+# CHECK:                        [ 0,  1,  3,  6, 28, 29]]),
+# CHECK: values=tensor([   0.,    0.,    1.,    2.,    3., 1000.]),
+# CHECK: size=(10, 20, 30), nnz=6, dtype=torch.float64, layout=torch.sparse_coo)
 #
 def test_sparse_coo3():
     class COO3Net(torch.nn.Module):
@@ -450,3 +355,11 @@ def forward(self, x):
 
     m = export_and_import(net, sparse_input)
     print(m)
+
+    # Run it with PyTorch torch.sparse and with TORCH-MLIR sparse_jit.
+    res1 = net(sparse_input)
+    res2 = sparse_jit(net, sparse_input)
+    print("torch.sparse")
+    print(res1)
+    print("torch.mlir")
+    print(res2)