Avoid unnecessary output buffer allocation and initialization.

[bixia]

Summary

The sparse tensor code generator allocates memory for the output tensor. As
such, we only need to allocate a MemRefDescriptor to receive the output tensor
and do not need to allocate and initialize the storage for the tensor.

[aartbik]

Can you refine:

"The sparse tensor code generator allocates memory for the output tensor."

in the description test. The sparse codegen does not do bufferization for dense tensors, but uses the to_memref/to_tensor at the boundaries.
So the actual allocation comes from a later bufferization. Just to make sure details are right.

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

92–95

can you split this into two lines, and assign an intuitive name to the memref descriptor

something like

ref_out = rt.make ...
mem_out = ctypes.pointer( ...

[aartbik]

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

92–95

And actually add what you have in the description as comment

1. Allocate a MemRefDescriptor to receive the output tensor.
2. The buffer itself is allocated inside the MLIR code.

ref_out = ...
mem_out = ...

or something like that

[bixia]

Address review feedback.

[bixia]

Update tree.

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

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)
-  out = np.zeros((3, 2), 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)))
-  mem_out = ctypes.pointer(ctypes.pointer(rt.get_ranked_memref_descriptor(out)))
+  # 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);