[mlir][sparse][python] migrate more code from boilerplate into proper numpy land

[aartbik]

Summary

The boilerplate was setting up some arrays for testing. To fully illustrate
python - MLIR potential, however, this data should also come from numpy land.

[bixia]

mlir/test/python/dialects/sparse_tensor/test_SpMM.py

58–59

Can we document that the first argument of main is first converted to sparse tensor with the given attr before calling the generated kernel?

86–91

go/pystyle#naming
The local variable naming convention is local_variable_name.
So, maybe we can use np_a, np_b, np_c, np_out here?

93–99

Similar to the above, maybe we want to use c_a_memref_ptr etc?

110

Now that we can see the input tensors here, can we use numpy to compute the expected result instead?

[aartbik]

mlir/test/python/dialects/sparse_tensor/test_SpMM.py

110

Ah yes, good idea!

[aartbik]

good suggestions!

mlir/test/python/dialects/sparse_tensor/test_SpMM.py

def spMxM(*args):
       return matmul_dsl(args[0], args[1], outs=[args[2]])
 
   return module
 
 
 def boilerplate(attr: st.EncodingAttr):
   """Returns boilerplate main method.
 
   This method sets up a boilerplate main method that calls the generated
   sparse kernel. For convenience, this part is purely done as string input.
   """
   return f"""
-func @main(%c: tensor<3x2xf64>) -> tensor<3x2xf64>
+func @main(%ad: tensor<3x4xf64>, %b: tensor<4x2xf64>, %c: tensor<3x2xf64>) -> tensor<3x2xf64>
   attributes {{ llvm.emit_c_interface }} {{
-  %0 = constant dense<[ [ 1.1,  0.0,  0.0,  1.4 ],
-                        [ 0.0,  0.0,  0.0,  0.0 ],
-                        [ 0.0,  0.0,  3.3,  0.0 ]]> : tensor<3x4xf64>
-  %a = sparse_tensor.convert %0 : tensor<3x4xf64> to tensor<3x4xf64, {attr}>
-  %b = constant dense<[ [ 1.0,  2.0 ],
-                        [ 4.0,  3.0 ],
-                        [ 5.0,  6.0 ],
-                        [ 8.0,  7.0 ]]> : tensor<4x2xf64>
-  %1 = call @spMxM(%a, %b, %c) : (tensor<3x4xf64, {attr}>,
+  %a = sparse_tensor.convert %ad : tensor<3x4xf64> to tensor<3x4xf64, {attr}>
+  %0 = call @spMxM(%a, %b, %c) : (tensor<3x4xf64, {attr}>,
                                   tensor<4x2xf64>,
                                   tensor<3x2xf64>) -> tensor<3x2xf64>
-  return %1 : tensor<3x2xf64>
+  return %0 : tensor<3x2xf64>
 }}
 """
 
 
 def build_compile_and_run_SpMM(attr: st.EncodingAttr, support_lib: str,
                                compiler):
   # Build.
   module = build_SpMM(attr)
   func = str(module.operation.regions[0].blocks[0].operations[0].operation)
   module = ir.Module.parse(func + boilerplate(attr))
+
   # Compile.
   compiler(module)
   engine = execution_engine.ExecutionEngine(
       module, opt_level=0, shared_libs=[support_lib])
-  # Set up numpy input, invoke the kernel, and get numpy output.
-  # Built-in bufferization uses in-out buffers.
-  # TODO: replace with inplace comprehensive bufferization.
-  Cin = np.zeros((3, 2), np.double)
-  Cout = np.zeros((3, 2), np.double)
-  Cin_memref_ptr = ctypes.pointer(
-      ctypes.pointer(rt.get_ranked_memref_descriptor(Cin)))
+
+  # Set up numpy input and buffer for output.
+  Ca = np.array(
+      [[1.1, 0.0, 0.0, 1.4], [0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 3.3, 0.0]],
+      np.float64)
+  Cb = np.array([[1.0, 2.0], [4.0, 3.0], [5.0, 6.0], [8.0, 7.0]], np.float64)
+  Cc = np.zeros((3, 2), np.float64)
+  Cout = np.zeros((3, 2), np.float64)
+
+  Ca_memref_ptr = ctypes.pointer(
+      ctypes.pointer(rt.get_ranked_memref_descriptor(Ca)))
+  Cb_memref_ptr = ctypes.pointer(
+      ctypes.pointer(rt.get_ranked_memref_descriptor(Cb)))
+  Cc_memref_ptr = ctypes.pointer(
+      ctypes.pointer(rt.get_ranked_memref_descriptor(Cc)))
   Cout_memref_ptr = ctypes.pointer(
       ctypes.pointer(rt.get_ranked_memref_descriptor(Cout)))
-  engine.invoke('main', Cout_memref_ptr, Cin_memref_ptr)
-  Cresult = rt.ranked_memref_to_numpy(Cout_memref_ptr[0])
+
+  # Invoke the kernel and get numpy output.
+  # Built-in bufferization uses in-out buffers.
+  # TODO: replace with inplace comprehensive bufferization.
+  engine.invoke('main', Cout_memref_ptr, Ca_memref_ptr, Cb_memref_ptr,
+                Cc_memref_ptr)
 
   # Sanity check on computed result.
-  expected = [[12.3, 12.0], [0.0, 0.0], [16.5, 19.8]]
+  Cresult = rt.ranked_memref_to_numpy(Cout_memref_ptr[0])
+  expected = np.array([[12.3, 12.0], [0.0, 0.0], [16.5, 19.8]], np.float64)
   if np.allclose(Cresult, expected):
     pass
   else:
     quit(f'FAILURE')
 
 
 class SparseCompiler:
   """Sparse compiler passes."""
class SparseCompiler:
   def __call__(self, module: ir.Module):
     passmanager.PassManager.parse(self.pipeline).run(module)
 
 
 # CHECK-LABEL: TEST: testSpMM
 # CHECK: Passed 72 tests
 @run
 def testSpMM():
+  # Obtain path to runtime support library.
   support_lib = os.getenv('SUPPORT_LIB')
+  assert os.path.exists(support_lib), f'{support_lib} does not exist'
+
   with ir.Context() as ctx, ir.Location.unknown():
     count = 0
     # Fixed compiler optimization strategy.
     # TODO: explore state space here too
     par = 0
     vec = 0
     vl = 1
     e = False