[mpact][compiler] add normalization and scaling test (#33)
* [mpact][compiler] enable vectorization (SIMD)
* [mpact][compiler] add normalization and scaling test
Note, this still needs torch-mlir work to be fully functional
* lint scale
diff --git a/python/mpact/models/gcn.py b/python/mpact/models/gcn.py
index 8a015de..df9b198 100644
--- a/python/mpact/models/gcn.py
+++ b/python/mpact/models/gcn.py
@@ -17,3 +17,7 @@
# Add bias.
output = output + self.bias
return output
+
+
+def graphconv44():
+ return GraphConv(input_dim=4, output_dim=4)
diff --git a/python/mpact/models/kernels.py b/python/mpact/models/kernels.py
index f8dbf17..dce3355 100644
--- a/python/mpact/models/kernels.py
+++ b/python/mpact/models/kernels.py
@@ -29,3 +29,21 @@
class SDDMMNet(torch.nn.Module):
def forward(self, x, y, z):
return torch.mul(x, torch.mm(y, z))
+
+
+class FeatureScale(torch.nn.Module):
+ def forward(self, F):
+ sum_vector = torch.sum(F, dim=1)
+ reciprocal_vector = 1 / sum_vector
+ reciprocal_vector[reciprocal_vector == float("inf")] = 0
+ scaling_diagonal = torch.diag(reciprocal_vector).to_sparse()
+ return scaling_diagonal @ F
+
+
+class Normalization(torch.nn.Module):
+ def forward(self, A):
+ sum_vector = torch.sum(A, dim=1)
+ reciprocal_vector = 1 / sum_vector
+ reciprocal_vector[reciprocal_vector == float("inf")] = 0
+ scaling_diagonal = torch.diag(reciprocal_vector).to_sparse()
+ return scaling_diagonal @ A @ scaling_diagonal
diff --git a/test/python/sparse_gcn.py b/test/python/gcn.py
similarity index 96%
rename from test/python/sparse_gcn.py
rename to test/python/gcn.py
index a1109a2..d44f1ea 100644
--- a/test/python/sparse_gcn.py
+++ b/test/python/gcn.py
@@ -3,9 +3,10 @@
import torch
from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
-from mpact.models.gcn import GraphConv
-net = GraphConv(4, 4)
+from mpact.models.gcn import graphconv44
+
+net = graphconv44()
# Get random (but reproducible) matrices.
torch.manual_seed(0)
diff --git a/test/python/sparse_lif.py b/test/python/lif.py
similarity index 99%
rename from test/python/sparse_lif.py
rename to test/python/lif.py
index 38973c9..b7c4288 100644
--- a/test/python/sparse_lif.py
+++ b/test/python/lif.py
@@ -3,6 +3,7 @@
import torch
from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+
from mpact.models.lif import Block
net = Block()
diff --git a/test/python/norm.py b/test/python/norm.py
new file mode 100644
index 0000000..f589e46
--- /dev/null
+++ b/test/python/norm.py
@@ -0,0 +1,45 @@
+# 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 Normalization
+
+net = Normalization()
+
+# Construct adjacency matrix.
+V = 8
+edges = np.array([[0, 1], [0, 4], [1, 4], [3, 4], [4, 3]], dtype=np.int32)
+E = edges.shape[0]
+adj_mat = torch.sparse_coo_tensor(edges.T, torch.ones(E), (V, V), dtype=torch.int64)
+adj_mat = (
+ torch.eye(V) + adj_mat
+) # Add self-loops to the adjacency matrix (becomes dense)
+
+#
+# CHECK: pytorch
+# CHECK: tensor({{\[}}[0.1111, 0.1667, 0.0000, 0.0000, 0.1667, 0.0000, 0.0000, 0.0000],
+# CHECK: [0.0000, 0.2500, 0.0000, 0.0000, 0.2500, 0.0000, 0.0000, 0.0000],
+# CHECK: [0.0000, 0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
+# CHECK: [0.0000, 0.0000, 0.0000, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000],
+# CHECK: [0.0000, 0.0000, 0.0000, 0.2500, 0.2500, 0.0000, 0.0000, 0.0000],
+# CHECK: [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000],
+# CHECK: [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 0.0000],
+# CHECK: [0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000]{{\]}})
+# CHECK: mpact
+#
+
+# Run it with PyTorch.
+print("pytorch")
+res = net(adj_mat)
+print(res)
+
+# Run it with MPACT.
+#
+# TODO: make this work, crashes in TORCH-MLIR
+#
+print("mpact")
+# res = mpact_jit(net, adj_mat)
+# print(res)
diff --git a/test/python/scale.py b/test/python/scale.py
new file mode 100644
index 0000000..880fdf2
--- /dev/null
+++ b/test/python/scale.py
@@ -0,0 +1,39 @@
+# 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 FeatureScale
+
+net = FeatureScale()
+
+# Get random (but reproducible) matrices.
+torch.manual_seed(0)
+features = torch.rand(7, 7)
+
+#
+# CHECK: pytorch
+# CHECK: tensor({{\[}}[0.1702, 0.2634, 0.0303, 0.0453, 0.1054, 0.2174, 0.1680],
+# CHECK: [0.3064, 0.1557, 0.2161, 0.1192, 0.1373, 0.0076, 0.0577],
+# CHECK: [0.0856, 0.1510, 0.2031, 0.2329, 0.0469, 0.0822, 0.1984],
+# CHECK: [0.2207, 0.0957, 0.2108, 0.1011, 0.1333, 0.2297, 0.0087],
+# CHECK: [0.0774, 0.1561, 0.1275, 0.3896, 0.0735, 0.1128, 0.0630],
+# CHECK: [0.0093, 0.0611, 0.2731, 0.2124, 0.2180, 0.1546, 0.0716],
+# CHECK: [0.2026, 0.0115, 0.0481, 0.0839, 0.2826, 0.2749, 0.0964]{{\]}})
+# CHECK: mpact
+#
+
+# Run it with PyTorch.
+print("pytorch")
+res = net(features)
+print(res)
+
+# Run it with MPACT.
+#
+# TODO: make this work, crashes in TORCH-MLIR
+#
+print("mpact")
+# res = mpact_jit(net, features)
+# print(res)
diff --git a/test/python/spmv.py b/test/python/spmv.py
index 604fab7..4f52ea0 100644
--- a/test/python/spmv.py
+++ b/test/python/spmv.py
@@ -3,6 +3,7 @@
import torch
from mpact.mpactbackend import mpact_jit, mpact_jit_compile, mpact_jit_run
+
from mpact.models.kernels import MVNet
net = MVNet()
