Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[mpact][compiler] more sparsity propagation rules with tests #39

Merged
merged 1 commit into from
Jun 17, 2024
Merged

[mpact][compiler] more sparsity propagation rules with tests #39

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 25 additions & 5 deletions python/mpact/mpactbackend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -359,6 +359,12 @@ def sparse_metadata(a: torch.Tensor) -> SparsityMeta:
raise RuntimeError(f"Unsupported sparse layout for {a}")


def sparse_arg(args, i):
if isinstance(args[i], torch.fx.node.Node):
return args[i].meta.get("sparsity", None)
return None


def sparse_export(
f: Callable, args: Tuple[Any, ...], kwargs: Optional[Dict[str, Any]] = None
) -> torch.export.ExportedProgram:
Expand Down Expand Up @@ -402,8 +408,22 @@ def sparse_export(
# TODO: use upstream _opname implementation when available
opname = node.target._schema.name.split("::")[1]
# Zero preserving elt-wise unary op.
if opname in {"abs", "neg", "relu", "sin", "mul"}:
node.meta["sparsity"] = node.args[0].meta.get("sparsity", None)
if opname in {"abs", "neg", "relu", "sin"}:
node.meta["sparsity"] = sparse_arg(node.args, 0)
# Some simplistic rules for preserving sparsity. Soon
# to be replaced by proper FX graph propagation.
elif opname in {"mul"}:
m0 = sparse_arg(node.args, 0)
m1 = sparse_arg(node.args, 1)
if m0 is not None:
node.meta["sparsity"] = m0
elif m1 is not None:
node.meta["sparsity"] = m1
elif opname in {"add", "mm"}:
m0 = sparse_arg(node.args, 0)
m1 = sparse_arg(node.args, 1)
if m0 is not None and m1 is not None:
node.meta["sparsity"] = m0
elif opname == "_to_sparse" or opname == "to_sparse":
dim = len(node.meta.get("val").shape)
node.meta["sparsity"] = SparsityMeta(
Expand All @@ -412,13 +432,13 @@ def sparse_export(
# TODO: Uncomment this to hack sparsity into the network.
# elif opname == "_to_dense" or opname == "to_dense":
# # hack (assumes we never really want the to_dense for now)
# node.meta["sparsity"] = node.args[0].meta.get("sparsity", None)
elif opname == "select" and node.args[0].meta.get("sparsity", None):
# node.meta["sparsity"] = sparse_arg(node.args, 0)
elif opname == "select" and sparse_arg(node.args, 0):
dim = len(node.meta.get("val").shape)
node.meta["sparsity"] = SparsityMeta(
torch.sparse_coo, 0, dim, 0, None, torch.int64, torch.int64
)
elif opname == "stack" and node.args[0][0].meta.get("sparsity", None):
elif opname == "stack" and sparse_arg(node.args[0], 0):
dim = len(node.meta.get("val").shape)
node.meta["sparsity"] = SparsityMeta(
torch.sparse_coo, 0, dim - 1, 1, None, torch.int64, torch.int64
Expand Down
89 changes: 89 additions & 0 deletions test/python/add.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
# RUN: %PYTHON %s | FileCheck %s

import torch
import numpy as np

from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run

from mpact.models.kernels import AddNet


def print_sparse(res):
print(res[0])
print(res[1])
print(res[2])


net = AddNet()

# Construct dense and sparse matrices.
X = torch.arange(0, 16, dtype=torch.float32).view(4, 4)
Y = torch.arange(16, 32, dtype=torch.float32).view(4, 4)
A = torch.tensor(
[
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 2.0],
[0.0, 0.0, 0.0, 0.0],
[3.0, 0.0, 0.0, 0.0],
],
dtype=torch.float32,
)
S = A.to_sparse_csr()

#
# CHECK: pytorch
# CHECK: tensor({{\[}}[16., 18., 20., 22.],
# CHECK: [24., 26., 28., 30.],
# CHECK: [32., 34., 36., 38.],
# CHECK: [40., 42., 44., 46.]{{\]}})
# CHECK: tensor({{\[}}[16., 18., 18., 19.],
# CHECK: [20., 21., 22., 25.],
# CHECK: [24., 25., 26., 27.],
# CHECK: [31., 29., 30., 31.]{{\]}})
# CHECK: tensor({{\[}}[ 0., 2., 2., 3.],
# CHECK: [ 4., 5., 6., 9.],
# CHECK: [ 8., 9., 10., 11.],
# CHECK: [15., 13., 14., 15.]{{\]}})
# CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]),
# CHECK: col_indices=tensor([1, 3, 0]),
# CHECK: values=tensor([2., 4., 6.]), size=(4, 4), nnz=3,
# CHECK: layout=torch.sparse_csr)
# CHECK: mpact
# CHECK: {{\[}}[16. 18. 20. 22.]
# CHECK: [24. 26. 28. 30.]
# CHECK: [32. 34. 36. 38.]
# CHECK: [40. 42. 44. 46.]{{\]}}
# CHECK: {{\[}}[16. 18. 18. 19.]
# CHECK: [20. 21. 22. 25.]
# CHECK: [24. 25. 26. 27.]
# CHECK: [31. 29. 30. 31.]{{\]}}
# CHECK: {{\[}}[ 0. 2. 2. 3.]
# CHECK: [ 4. 5. 6. 9.]
# CHECK: [ 8. 9. 10. 11.]
# CHECK: [15. 13. 14. 15.]{{\]}}
# CHECK: [0 1 2 2 3]
# CHECK: [1 3 0]
# CHECK: [2. 4. 6.]
#

# Run it with PyTorch.
print("pytorch")
res = net(X, Y)
print(res)
res = net(S, Y)
print(res)
res = net(X, S)
print(res)
res = net(S, S)
print(res)

# Run it with MPACT.
print("mpact")
res = mpact_jit(net, X, Y)
print(res)
res = mpact_jit(net, S, Y)
print(res)
res = mpact_jit(net, X, S)
print(res)
res = mpact_jit(net, S, S)
print_sparse(res)
89 changes: 89 additions & 0 deletions test/python/mm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,89 @@
# RUN: %PYTHON %s | FileCheck %s

import torch
import numpy as np

from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run

from mpact.models.kernels import MMNet


def print_sparse(res):
print(res[0])
print(res[1])
print(res[2])


net = MMNet()

# Construct dense and sparse matrices.
X = torch.arange(0, 16, dtype=torch.float32).view(4, 4)
Y = torch.arange(16, 32, dtype=torch.float32).view(4, 4)
A = torch.tensor(
[
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 2.0],
[0.0, 0.0, 0.0, 0.0],
[3.0, 0.0, 0.0, 0.0],
],
dtype=torch.float32,
)
S = A.to_sparse_csr()

#
# CHECK: pytorch
# CHECK: tensor({{\[}}[ 152., 158., 164., 170.],
# CHECK: [ 504., 526., 548., 570.],
# CHECK: [ 856., 894., 932., 970.],
# CHECK: [1208., 1262., 1316., 1370.]{{\]}})
# CHECK: tensor({{\[}}[20., 21., 22., 23.],
# CHECK: [56., 58., 60., 62.],
# CHECK: [ 0., 0., 0., 0.],
# CHECK: [48., 51., 54., 57.]{{\]}})
# CHECK: tensor({{\[}}[ 9., 0., 0., 2.],
# CHECK: [21., 4., 0., 10.],
# CHECK: [33., 8., 0., 18.],
# CHECK: [45., 12., 0., 26.]{{\]}})
# CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]),
# CHECK: col_indices=tensor([3, 0, 1]),
# CHECK: values=tensor([2., 6., 3.]), size=(4, 4), nnz=3,
# CHECK: layout=torch.sparse_csr)
# CHECK: mpact
# CHECK: {{\[}}[ 152. 158. 164. 170.]
# CHECK: [ 504. 526. 548. 570.]
# CHECK: [ 856. 894. 932. 970.]
# CHECK: [1208. 1262. 1316. 1370.]{{\]}}
# CHECK: {{\[}}[20. 21. 22. 23.]
# CHECK: [56. 58. 60. 62.]
# CHECK: [ 0. 0. 0. 0.]
# CHECK: [48. 51. 54. 57.]{{\]}}
# CHECK: {{\[}}[ 9. 0. 0. 2.]
# CHECK: [21. 4. 0. 10.]
# CHECK: [33. 8. 0. 18.]
# CHECK: [45. 12. 0. 26.]{{\]}}
# CHECK: [0 1 2 2 3]
# CHECK: [3 0 1]
# CHECK: [2. 6. 3.]
#

# Run it with PyTorch.
print("pytorch")
res = net(X, Y)
print(res)
res = net(S, Y)
print(res)
res = net(X, S)
print(res)
res = net(S, S)
print(res)

# Run it with MPACT.
print("mpact")
res = mpact_jit(net, X, Y)
print(res)
res = mpact_jit(net, S, Y)
print(res)
res = mpact_jit(net, X, S)
print(res)
res = mpact_jit(net, S, S)
print_sparse(res)
87 changes: 87 additions & 0 deletions test/python/mul.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,87 @@
# RUN: %PYTHON %s | FileCheck %s

import torch
import numpy as np

from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run

from mpact.models.kernels import MulNet


def print_sparse(res):
print(res[0])
print(res[1])
print(res[2])


net = MulNet()

# Construct dense and sparse matrices.
X = torch.arange(0, 16, dtype=torch.float32).view(4, 4)
Y = torch.arange(16, 32, dtype=torch.float32).view(4, 4)
A = torch.tensor(
[
[0.0, 1.0, 0.0, 0.0],
[0.0, 0.0, 0.0, 2.0],
[0.0, 0.0, 0.0, 0.0],
[3.0, 0.0, 0.0, 0.0],
],
dtype=torch.float32,
)
S = A.to_sparse_csr()

#
# CHECK: pytorch
# CHECK: tensor({{\[}}[ 0., 17., 36., 57.],
# CHECK: [ 80., 105., 132., 161.],
# CHECK: [192., 225., 260., 297.],
# CHECK: [336., 377., 420., 465.]{{\]}})
# CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]),
# CHECK: col_indices=tensor([1, 3, 0]),
# CHECK: values=tensor([17., 46., 84.]), size=(4, 4), nnz=3,
# CHECK: layout=torch.sparse_csr)
# CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]),
# CHECK: col_indices=tensor([1, 3, 0]),
# CHECK: values=tensor([ 1., 14., 36.]), size=(4, 4), nnz=3,
# CHECK: layout=torch.sparse_csr)
# CHECK: tensor(crow_indices=tensor([0, 1, 2, 2, 3]),
# CHECK: col_indices=tensor([1, 3, 0]),
# CHECK: values=tensor([1., 4., 9.]), size=(4, 4), nnz=3,
# CHECK: layout=torch.sparse_csr)
# CHECK: mpact
# CHECK: {{\[}}[ 0. 17. 36. 57.]
# CHECK: [ 80. 105. 132. 161.]
# CHECK: [192. 225. 260. 297.]
# CHECK: [336. 377. 420. 465.]{{\]}}
# CHECK: [0 1 2 2 3]
# CHECK: [1 3 0]
# CHECK: [17. 46. 84.]
# CHECK: [0 1 2 2 3]
# CHECK: [1 3 0]
# CHECK: [ 1. 14. 36.]
# CHECK: [0 1 2 2 3]
# CHECK: [1 3 0]
# CHECK: [1. 4. 9.]
#

# Run it with PyTorch.
print("pytorch")
res = net(X, Y)
print(res)
res = net(S, Y)
print(res)
res = net(X, S)
print(res)
res = net(S, S)
print(res)

# Run it with MPACT.
print("mpact")
res = mpact_jit(net, X, Y)
print(res)
res = mpact_jit(net, S, Y)
print_sparse(res)
res = mpact_jit(net, X, S)
print_sparse(res)
res = mpact_jit(net, S, S)
print_sparse(res)
13 changes: 9 additions & 4 deletions test/python/sddmm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,14 @@

from mpact.models.kernels import MMNet, SDDMMNet


def print_sparse(res):
print(res[0])
print(res[1])
print(res[2])
print(res[3])


mmnet = MMNet()
sddmmnet = SDDMMNet()

Expand Down Expand Up @@ -54,7 +62,4 @@
dense = mpact_jit(mmnet, A, B)
print(dense)
res = mpact_jit(sddmmnet, S, A, B)
print(res[0])
print(res[1])
print(res[2])
print(res[3])
print_sparse(res)
Loading