[mlir][sparse] python driven test for SDDMM

[aartbik]

Summary

explores various sparsity combinations of
the SDMM kernel and verifies that the computed
result is the same for all cases

[aartbik]

removed debug code

[aartbik]

rebased main

[aartbik]

rebased and made expected pre-computed

mlir/test/Integration/Dialect/SparseTensor/python/test_SDDMM.py

 from mlir import passmanager
 
 from mlir.dialects import sparse_tensor as st
 from mlir.dialects import builtin
 from mlir.dialects.linalg.opdsl import lang as dsl
 
 
 @dsl.linalg_structured_op
-def matmul_dsl(
+def sddmm_dsl(
     A=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.K),
     B=dsl.TensorDef(dsl.T, dsl.S.K, dsl.S.N),
+    S=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.N),
     C=dsl.TensorDef(dsl.T, dsl.S.M, dsl.S.N, output=True)):
-  C[dsl.D.m, dsl.D.n] += A[dsl.D.m, dsl.D.k] * B[dsl.D.k, dsl.D.n]
+  C[dsl.D.m,
+    dsl.D.n] += S[dsl.D.m, dsl.D.n] * A[dsl.D.m, dsl.D.k] * B[dsl.D.k, dsl.D.n]
 
 
-def build_SpMM(attr: st.EncodingAttr):
-  """Build SpMM kernel.
+def build_SDDMM(attr: st.EncodingAttr):
+  """Build SDDMM kernel.
 
   This method generates a linalg op with for matrix multiplication using
   just the Python API. Effectively, a generic linalg op is constructed
-  that computes C(i,j) += A(i,k) * B(k,j) for annotated matrix A.
+  that computes C(i,j) += S(i,j) SUM_k A(i,k) B(k,j) for sparse S.
   """
   module = ir.Module.create()
   f64 = ir.F64Type.get()
-  a = ir.RankedTensorType.get([3, 4], f64, attr)
-  b = ir.RankedTensorType.get([4, 2], f64)
-  c = ir.RankedTensorType.get([3, 2], f64)
-  arguments = [a, b, c]
+  a = ir.RankedTensorType.get([8, 8], f64)
+  b = ir.RankedTensorType.get([8, 8], f64)
+  c = ir.RankedTensorType.get([8, 8], f64)
+  s = ir.RankedTensorType.get([8, 8], f64, attr)
+  arguments = [a, b, s, c]
   with ir.InsertionPoint(module.body):
 
     @builtin.FuncOp.from_py_func(*arguments)
-    def spMxM(*args):
-      return matmul_dsl(args[0], args[1], outs=[args[2]])
+    def sddmm(*args):
+      return sddmm_dsl(args[0], args[1], args[2], outs=[args[3]])
 
   return module
 
 
 def boilerplate(attr: st.EncodingAttr):
-  """Returns boilerplate main method.
-
-  This method sets up a boilerplate main method that takes three tensors
-  (a, b, c), converts the first tensor a into s sparse tensor, and then
-  calls the sparse kernel for matrix multiplication. For convenience,
-  this part is purely done as string input.
-  """
+  """Returns boilerplate code for main driver."""
   return f"""
-func @main(%ad: tensor<3x4xf64>, %b: tensor<4x2xf64>, %c: tensor<3x2xf64>) -> tensor<3x2xf64>
-  attributes {{ llvm.emit_c_interface }} {{
-  %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 %0 : tensor<3x2xf64>
+func @main(%a: tensor<8x8xf64>,
+           %b: tensor<8x8xf64>,
+           %c: tensor<8x8xf64>) -> tensor<8x8xf64> attributes {{ llvm.emit_c_interface }} {{
+  %t = arith.constant sparse<[[0,0], [0,2], [4,1]], [1.0, 2.0, 3.0]> : tensor<8x8xf64>
+  %s = sparse_tensor.convert %t : tensor<8x8xf64> to tensor<8x8xf64, {attr}>
+  %0 = call @sddmm(%a, %b, %s, %c) : (tensor<8x8xf64>,
+                                      tensor<8x8xf64>,
+                                      tensor<8x8xf64, {attr}>,
+                                      tensor<8x8xf64>) -> tensor<8x8xf64>
+  return %0 : tensor<8x8xf64>
 }}
 """
 
 
-def build_compile_and_run_SpMM(attr: st.EncodingAttr, support_lib: str,
-                               compiler):
+def build_compile_and_run_SDDMMM(attr: st.EncodingAttr, opt: str,
+                                 support_lib: str, compiler):
   # Build.
-  module = build_SpMM(attr)
+  module = build_SDDMM(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 and buffer for output.
-  a = 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)
-  b = np.array([[1.0, 2.0], [4.0, 3.0], [5.0, 6.0], [8.0, 7.0]], np.float64)
-  c = np.zeros((3, 2), np.float64)
+  a = np.array([[1.1, 2.1, 3.1, 4.1, 5.1, 6.1, 7.1, 8.1],
+                [1.2, 2.2, 3.2, 4.2, 5.2, 6.2, 7.2, 8.2],
+                [1.3, 2.3, 3.3, 4.3, 5.3, 6.3, 7.3, 8.3],
+                [1.4, 2.4, 3.4, 4.4, 5.4, 6.4, 7.4, 8.4],
+                [1.5, 2.5, 3.5, 4.5, 5.5, 6.5, 7.5, 8.5],
+                [1.6, 2.6, 3.6, 4.6, 5.6, 6.6, 7.6, 8.6],
+                [1.7, 2.7, 3.7, 4.7, 5.7, 6.7, 7.7, 8.7],
+                [1.8, 2.8, 3.8, 4.8, 5.8, 6.8, 7.8, 8.8]], np.float64)
+  b = np.ones((8, 8), np.float64)
+  c = np.zeros((8, 8), np.float64)
 
   mem_a = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(a)))
   mem_b = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(b)))
   mem_c = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(c)))
+
   # Allocate a MemRefDescriptor to receive the output tensor.
   # The buffer itself is allocated inside the MLIR code generation.
   ref_out = rt.make_nd_memref_descriptor(2, ctypes.c_double)()
   mem_out = ctypes.pointer(ctypes.pointer(ref_out))
 
   # Invoke the kernel and get numpy output.
   # Built-in bufferization uses in-out buffers.
   # TODO: replace with inplace comprehensive bufferization.
   engine.invoke('main', mem_out, mem_a, mem_b, mem_c)
 
-  # Sanity check on computed result.
-  expected = np.matmul(a, b);
+  # Sanity check on computed result. Only a few elements
+  # are sampled from the full dense matrix multiplication.
+  full_matmul = np.matmul(a, b)
+  expected = np.zeros((8, 8), np.float64)
+  expected[0, 0] = 1.0 * full_matmul[0, 0]
+  expected[0, 2] = 2.0 * full_matmul[0, 2]
+  expected[4, 1] = 3.0 * full_matmul[4, 1]
   c = rt.ranked_memref_to_numpy(mem_out[0])
   if np.allclose(c, expected):
     pass
   else:
     quit(f'FAILURE')
 
 
 class SparseCompiler:
   """Sparse compiler passes."""
 
   def __init__(self, options: str):
     pipeline = (
-        f'builtin.func(linalg-generalize-named-ops,linalg-fuse-elementwise-ops),'
         f'sparsification{{{options}}},'
         f'sparse-tensor-conversion,'
         f'builtin.func(linalg-bufferize,convert-linalg-to-loops,convert-vector-to-scf),'
         f'convert-scf-to-std,'
         f'func-bufferize,'
         f'tensor-constant-bufferize,'
         f'builtin.func(tensor-bufferize,std-bufferize,finalizing-bufferize),'
         f'convert-vector-to-llvm{{reassociate-fp-reductions=1 enable-index-optimizations=1}},'
         f'lower-affine,'
         f'convert-memref-to-llvm,'
         f'convert-std-to-llvm,'
         f'reconcile-unrealized-casts')
     self.pipeline = pipeline
 
   def __call__(self, module: ir.Module):
     passmanager.PassManager.parse(self.pipeline).run(module)
 
 
 def main():
   support_lib = os.getenv('SUPPORT_LIB')
   assert support_lib is not None, 'SUPPORT_LIB is undefined'
   if not os.path.exists(support_lib):
-    raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), support_lib)
+    raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT),
+                            support_lib)
 
-  # CHECK-LABEL: TEST: testSpMM
-  print('\nTEST: testSpMM')
+  # CHECK-LABEL: TEST: testSDDMMM
+  print('\nTEST: testSDDMMM')
   with ir.Context() as ctx, ir.Location.unknown():
     count = 0
-    # Loop over various ways to compile and annotate the SpMM kernel with
+    # Loop over various ways to compile and annotate the SDDMM kernel with
     # a *single* sparse tensor. Note that we deliberate do not exhaustively
     # search the full state space to reduce runtime of the test. It is
     # straightforward to adapt the code below to explore more combinations.
-    par = 0
-    vec = 0
-    vl = 1
-    e = False
-    opt = (f'parallelization-strategy={par} '
-           f'vectorization-strategy={vec} '
-           f'vl={vl} enable-simd-index32={e}')
     levels = [[st.DimLevelType.dense, st.DimLevelType.dense],
               [st.DimLevelType.dense, st.DimLevelType.compressed],
               [st.DimLevelType.compressed, st.DimLevelType.dense],
               [st.DimLevelType.compressed, st.DimLevelType.compressed]]
     orderings = [
         ir.AffineMap.get_permutation([0, 1]),
         ir.AffineMap.get_permutation([1, 0])
     ]
-    bitwidths = [0]
     for level in levels:
       for ordering in orderings:
-        for pwidth in bitwidths:
-          for iwidth in bitwidths:
+        for pwidth in [32]:
+          for iwidth in [32]:
+            for par in [0]:
+              for vec in [0, 1]:
+                for e in [True]:
+                  vl = 1 if vec == 0 else 16
                   attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
+                  opt = (f'parallelization-strategy={par} '
+                         f'vectorization-strategy={vec} '
+                         f'vl={vl} enable-simd-index32={e}')
                   compiler = SparseCompiler(options=opt)
-            build_compile_and_run_SpMM(attr, support_lib, compiler)
+                  build_compile_and_run_SDDMMM(attr, opt, support_lib, compiler)
                   count = count + 1
-    # CHECK: Passed 8 tests
-    print('Passed ', count, 'tests')
+  # CHECK: Passed 16 tests
 
+
 if __name__ == '__main__':
   main()

mlir/test/Integration/Dialect/SparseTensor/python/test_SpMM.py

for level in levels:
           for iwidth in bitwidths:
             attr = st.EncodingAttr.get(level, ordering, pwidth, iwidth)
             compiler = SparseCompiler(options=opt)
             build_compile_and_run_SpMM(attr, support_lib, compiler)
             count = count + 1
     # CHECK: Passed 8 tests
     print('Passed ', count, 'tests')
 
+
 if __name__ == '__main__':
   main()