Dual Views managed by Python

examples/bspmm.py

@@ -100,7 +100,7 @@ def main():
         backend = KokkosBackend.KokkosBackend(decompose_tensors=True, parallel_strategy=par, index_instance=instance, num_instances=len(parStrats))
         instance += 1
         module_kokkos = backend.compile(moduleText)
-        C_kokkos = module_kokkos.pte_local_bspmm(rowptrs, colinds, values, ((m, n, b), (m+1, nnz, nnz*b)), B)
+        C_kokkos = module_kokkos.pte_local_bspmm(rowptrs, colinds, values, ((m, n, b), (m+1, nnz, nnz*b)), B).asnumpy()
         # For debugging: print the CSRV formatted matrix
         # module_kokkos.print_csrv(rowptrs, colinds, values, ((m, n, b), (m+1, nnz, nnz*b)))
         if np.allclose(C_gold, C_kokkos):

examples/csrv_softmax.py

@@ -155,7 +155,7 @@ def main():
         # For debugging: print the CSRV formatted matrix
         #module_kokkos.print_csrv(A)
         result = module_kokkos.pte_softmax(A)
-        resultDense = module_kokkos.csrv_to_dense(result)
+        resultDense = module_kokkos.csrv_to_dense(result).asnumpy()
         print("Result (converted to dense): ")
         print(resultDense)
         if checkResult is None:

examples/gemm_no_alloc.py

@@ -32,9 +32,14 @@ def main():
     ckokkos = torch.zeros((m, n))
 
     backend = KokkosBackend.KokkosBackend(dump_mlir=False)
-    k_backend = backend.compile(module)
+    should_compile = True
+    if should_compile:
+        k_backend = backend.compile(module)
+    else:
+        import lapis_package.lapis_package as k_backend
 
     print("a*b from kokkos")
+    print(f"{type(a)=} {type(b)=} {type(ckokkos)=}")
     k_backend.forward(a, b, ckokkos)
     print(ckokkos)
 

examples/hitting_times.py

@@ -212,6 +212,7 @@ def main():
     (mom1, mom2) = module_kokkos.mht(A.indptr, A.indices, A.data, ((n, n), (n + 1, nnz, nnz)), mask, D, 0.99, 1e-10, 20)
     # Normalize 2nd moment
     mom2_norm = module_kokkos.normalize_mom2(mom1, mom2)
+    mom1, mom2_norm = mom1.asnumpy(), mom2_norm.asnumpy()
     print("1st moment:", mom1)
     print("2nd moment:", mom2_norm)
     mom1_gold = [9.9999999999999957e+01, 9.9999999999999957e+01, 1.9947875961498600e+01, 0, 1.8542105566249749e+01, 1.9909372734041774e+01, 9.9999999999999957e+01, 1.9970947033662114e+01, 2.0196661010469668e+01, 1.4624345260746775e+01, 1.8252877036565582e+01, 1.9213516388393590e+01]

examples/issue76.py

@@ -69,7 +69,7 @@ def main():
     module_kokkos.print_dcsr(A_dcsr)
     [result, rank, nnz] = module_kokkos.column_sums(A_dcsr)
     # Convert the sparse vector result to dense to check the output
-    result_dense = module_kokkos.sparse_vec_to_dense(result)
+    result_dense = module_kokkos.sparse_vec_to_dense(result).asnumpy()
     print("Results:        ", rank, nnz, result_dense)
     print("Correct result: ", 1, gold_nnz, gold)
 

examples/multiple_results.py

@@ -7,10 +7,11 @@
 module {
   func.func @plus_norm(%x: tensor<?xf64>, %y: tensor<?xf64>) -> (tensor<?xf64>, f64) {
     %c0 = arith.constant 0 : index
+    %f0 = arith.constant 0.0 : f64
     %n = tensor.dim %y, %c0 : tensor<?xf64>
-    %alloc1 = tensor.empty (%n) : tensor<?xf64>
+    %alloc1 = tensor.splat %f0[%n] : tensor<?xf64>
     %x_plus_y = linalg.add ins(%x, %y : tensor<?xf64>, tensor<?xf64>) outs(%alloc1: tensor<?xf64>) -> tensor<?xf64>
-    %alloc2 = tensor.empty () : tensor<f64>
+    %alloc2 = tensor.splat %f0 : tensor<f64>
     %sumTensor = linalg.dot ins(%x_plus_y, %x_plus_y : tensor<?xf64>, tensor<?xf64>) outs(%alloc2: tensor<f64>) -> tensor<f64>
     %sum = tensor.extract %sumTensor[] : tensor<f64>
     %norm = math.sqrt %sum : f64
@@ -29,7 +30,8 @@ def main():
     print("x:", x)
     print("y:", y)
 
-    (x_plus_y, norm) = module_kokkos.plus_norm(x, y)
+    x_plus_y, norm = module_kokkos.plus_norm(x, y)
+    x_plus_y = x_plus_y.asnumpy()
 
     print("x + y:", x_plus_y)
     print("norm(x+y):", norm)

examples/pcg_solve.py

@@ -53,16 +53,16 @@
     %f1 = arith.constant 1.0 : f64
     %fm1 = arith.constant -1.0 : f64
     // Preallocate some intermediate tensors for dst-passing style
-    %buf0 = tensor.empty(%n) : tensor<?xf64>
-    %buf1 = tensor.empty(%n) : tensor<?xf64>
-    %buf2 = tensor.empty(%n) : tensor<?xf64>
+    %buf0 = tensor.splat %f0[%n] : tensor<?xf64>
+    %buf1 = tensor.splat %f0[%n] : tensor<?xf64>
+    %buf2 = tensor.splat %f0[%n] : tensor<?xf64>
     // Assume initial guess x0 = 0
     // Then r0 = b - A*x0 = b
     %r0 = linalg.copy ins(%b : tensor<?xf64>) outs(%buf0 : tensor<?xf64>) -> tensor<?xf64>
     %z0 = func.call @mult(%r0, %dinv, %buf1) : (tensor<?xf64>, tensor<?xf64>, tensor<?xf64>) -> tensor<?xf64>
     %p0 = linalg.copy ins(%z0 : tensor<?xf64>) outs(%buf2 : tensor<?xf64>) -> tensor<?xf64>
     %x0 = tensor.splat %f0[%n] : tensor<?xf64>
-    %Apbuf = tensor.empty(%n) : tensor<?xf64>
+    %Apbuf = tensor.splat %f0[%n] : tensor<?xf64>
     %rr0 = func.call @dot(%r0, %r0) : (tensor<?xf64>, tensor<?xf64>) -> f64
     %initres = math.sqrt %rr0 : f64
     %x, %p, %z, %r, %final_relres, %rz, %iters = scf.while (%xiter = %x0, %piter = %p0, %ziter = %z0, %riter = %r0, %rziter = %f0, %i = %c1) : (tensor<?xf64>, tensor<?xf64>, tensor<?xf64>, tensor<?xf64>, f64, index) -> (tensor<?xf64>, tensor<?xf64>, tensor<?xf64>, tensor<?xf64>, f64, f64, index)
@@ -118,11 +118,16 @@ def main():
     maxiter = 40
 
     backend = KokkosBackend.KokkosBackend(decompose_tensors=True)
-    module_kokkos = backend.compile(moduleText)
+    should_compile = True
+    if should_compile:
+        module_kokkos = backend.compile(moduleText)
+    else:
+        import lapis_package.lapis_package as module_kokkos
 
     print("x exact solution (first 10 elements):", xgold[:10])
 
     (x, numiter, relres) = module_kokkos.pcg(A.indptr, A.indices, A.data, ((n, n), (n + 1, nnz, nnz)), b, dinv, reltol, maxiter)
+    x = x.asnumpy()
     print("Ran CG for", numiter, "iterations and achieved relative residual norm", relres)
     print("x vector from LAPIS (first 10 elements):", x[:10])
 

examples/scalar_args.py

@@ -37,8 +37,12 @@ def main():
         xcorrect[i] += alpha
 
     backend = KokkosBackend.KokkosBackend(decompose_tensors=True)
-    module_kokkos = backend.compile(moduleText)
-    x = module_kokkos.f(A, alpha, k, False)
+    should_compile = True
+    if should_compile:
+        module_kokkos = backend.compile(moduleText)
+    else:
+        import lapis_package.lapis_package as module_kokkos
+    x = module_kokkos.f(A, alpha, k, False).asnumpy()
     print("Result 1:", x)
 
     if not np.allclose(x, xcorrect):
@@ -49,7 +53,7 @@ def main():
     for i in range(8):
         xcorrect[i] -= alpha
 
-    x = module_kokkos.f(A, alpha, k, True)
+    x = module_kokkos.f(A, alpha, k, True).asnumpy()
     print("Result 2:", x)
 
     if not np.allclose(x, xcorrect):

examples/sparse_axpy.py

@@ -85,9 +85,7 @@ def check_axpy(module_kokkos, v1_pos, v1_inds, v1_vals, v2_pos, v2_inds, v2_vals
         print("Result crd:", result_crd)
         print("Result val:", result_val)
     print("Test case result:")
-    # Alternate way to print sparse vector, using sparse_tensor.print op
-    #module_kokkos.print_sparse_vec(result)
-    result = module_kokkos.sparse_to_dense(result)
+    result = module_kokkos.sparse_to_dense(result).asnumpy()
     if correct_nnz != actual_nnz:
         print("Failed: result nonzero count incorrect")
         return False

examples/spmv_noalloc.py

@@ -23,7 +23,11 @@ def main():
     ykokkos = np.zeros((m), dtype=np.double)
 
     backend = KokkosBackend.KokkosBackend(decompose_tensors=True)
-    module_kokkos = backend.compile(moduleText)
+    should_compile = True
+    if should_compile:
+        module_kokkos = backend.compile(moduleText)
+    else:
+        import lapis_package.lapis_package as module_kokkos
 
     print("y = Ax from kokkos:")
     module_kokkos.spmv(rowptrs, colinds, values, ((m, n), (len(rowptrs), len(colinds), len(values))), x, ykokkos)