| # 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 SqSum |
| |
| net = SqSum() |
| |
| # 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) |
| |
| # |
| # CHECK: pytorch |
| # CHECK: tensor(5) |
| # CHECK: mpact |
| # CHECK: 5 |
| |
| # Run it with PyTorch. |
| print("pytorch") |
| res = net(adj_mat) |
| print(res) |
| |
| # Run it with MPACT. |
| print("mpact") |
| # TODO: make this work, expose `sparse-emit-strategy=sparse-iterator` to |
| # mini-pipeline. |
| # res = mpact_jit(net, adj_mat, use_sp_it=True) |
| res = mpact_jit(net, adj_mat) |
| print(res) |
