[Mlir-commits] [mlir] ae7ee65 - [mlir][taco] Add a utility to create an MLIR sparse tensor from a file.
Bixia Zheng
llvmlistbot at llvm.org
Tue Feb 1 15:43:59 PST 2022
Author: Bixia Zheng
Date: 2022-02-01T15:43:53-08:00
New Revision: ae7ee655a9f1386feed49d6eb5c902bfebb752ec
URL: https://github.com/llvm/llvm-project/commit/ae7ee655a9f1386feed49d6eb5c902bfebb752ec
DIFF: https://github.com/llvm/llvm-project/commit/ae7ee655a9f1386feed49d6eb5c902bfebb752ec.diff
LOG: [mlir][taco] Add a utility to create an MLIR sparse tensor from a file.
Move the functions that retrieve the supporting C library, compile an MLIR
module and build a JIT execution engine to mlir_pytaco_utils.
Add a function to create an MLIR sparse tensor from a file and return a pointer
to the MLIR sparse tensor as well as the shape of the sparse tensor.
Add unit tests.
Reviewed By: aartbik
Differential Revision: https://reviews.llvm.org/D118496
Added:
mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_utils.py
Modified:
mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py
Removed:
################################################################################
diff --git a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
index f3e865ba860a3..f74ae09b9087e 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
@@ -30,8 +30,6 @@
import threading
# Import MLIR related modules.
-from mlir import all_passes_registration # Register MLIR compiler passes.
-from mlir import execution_engine
from mlir import ir
from mlir import runtime
from mlir.dialects import arith
@@ -40,7 +38,6 @@
from mlir.dialects import std
from mlir.dialects import sparse_tensor
from mlir.dialects.linalg.opdsl import lang
-from mlir.passmanager import PassManager
from . import mlir_pytaco_utils as utils
@@ -51,13 +48,6 @@
# Bitwidths for pointers and indices.
_POINTER_BIT_WIDTH = 0
_INDEX_BIT_WIDTH = 0
-# The name for the environment variable that provides the full path for the
-# supporting library.
-_SUPPORTLIB_ENV_VAR = "SUPPORTLIB"
-# The default supporting library if the environment variable is not provided.
-_DEFAULT_SUPPORTLIB = "libmlir_c_runner_utils.so"
-# The JIT compiler optimization level.
-_OPT_LEVEL = 2
# The entry point to the JIT compiled program.
_ENTRY_NAME = "main"
@@ -134,33 +124,6 @@ def _mlir_type_from_taco_type(dtype: DType) -> ir.Type:
return dtype_to_irtype[dtype.kind]
-def _compile_mlir(module: ir.Module) -> ir.Module:
- """Compiles an MLIR module and returns the compiled module."""
- # TODO: Replace this with a pipeline implemented for
- # https://github.com/llvm/llvm-project/issues/51751.
- pipeline = (
- f"sparsification,"
- 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"arith-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")
- PassManager.parse(pipeline).run(module)
- return module
-
-
- at functools.lru_cache()
-def _get_support_lib_name() -> str:
- """Returns the string for the supporting C shared library."""
- return os.getenv(_SUPPORTLIB_ENV_VAR, _DEFAULT_SUPPORTLIB)
-
-
def _ctype_pointer_from_array(array: np.ndarray) -> ctypes.pointer:
"""Returns the ctype pointer for the given numpy array."""
return ctypes.pointer(
@@ -900,8 +863,7 @@ def ctype_pointer(self) -> ctypes.pointer:
shape = np.array(self._shape, np.int64)
indices = np.array(self._coords, np.int64)
values = np.array(self._values, self._dtype.value)
- ptr = utils.coo_tensor_to_sparse_tensor(_get_support_lib_name(), shape,
- values, indices)
+ ptr = utils.coo_tensor_to_sparse_tensor(shape, values, indices)
return ctypes.pointer(ctypes.cast(ptr, ctypes.c_void_p))
def get_coordinates_and_values(
@@ -1316,18 +1278,12 @@ def evaluate(
input_accesses = []
self._visit(_gather_input_accesses_index_vars, (input_accesses,))
- support_lib = _get_support_lib_name()
# Build and compile the module to produce the execution engine.
with ir.Context(), ir.Location.unknown():
module = ir.Module.create()
self._emit_assignment(module, dst, dst_indices, expr_to_info,
input_accesses)
- compiled_module = _compile_mlir(module)
-
- # We currently rely on an environment to pass in the full path of a
- # supporting library for the execution engine.
- engine = execution_engine.ExecutionEngine(
- compiled_module, opt_level=_OPT_LEVEL, shared_libs=[support_lib])
+ engine = utils.compile_and_build_engine(module)
# Gather the pointers for the input buffers.
input_pointers = [a.tensor.ctype_pointer() for a in input_accesses]
@@ -1351,7 +1307,6 @@ def evaluate(
# Check and return the sparse tensor output.
rank, nse, shape, values, indices = utils.sparse_tensor_to_coo_tensor(
- support_lib,
ctypes.cast(arg_pointers[-1][0], ctypes.c_void_p),
np.float64,
)
diff --git a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py
index 867a129e9a09b..2e4d5a3880dc3 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py
@@ -4,21 +4,47 @@
# This file contains the utilities to process sparse tensor outputs.
-from typing import Tuple
+from typing import Sequence, Tuple
import ctypes
import functools
import numpy as np
+import os
+
+# Import MLIR related modules.
+from mlir import all_passes_registration # Register MLIR compiler passes.
+from mlir import execution_engine
+from mlir import ir
+from mlir import runtime
+from mlir.dialects import sparse_tensor
+from mlir.passmanager import PassManager
+
+# The name for the environment variable that provides the full path for the
+# supporting library.
+_SUPPORTLIB_ENV_VAR = "SUPPORTLIB"
+# The default supporting library if the environment variable is not provided.
+_DEFAULT_SUPPORTLIB = "libmlir_c_runner_utils.so"
+
+# The JIT compiler optimization level.
+_OPT_LEVEL = 2
+# The entry point to the JIT compiled program.
+_ENTRY_NAME = "main"
@functools.lru_cache()
-def _get_c_shared_lib(lib_name: str) -> ctypes.CDLL:
- """Loads and returns the requested C shared library.
+def _get_support_lib_name() -> str:
+ """Gets the string name for the supporting C shared library."""
+ return os.getenv(_SUPPORTLIB_ENV_VAR, _DEFAULT_SUPPORTLIB)
- Args:
- lib_name: A string representing the C shared library.
+
+ at functools.lru_cache()
+def _get_c_shared_lib() -> ctypes.CDLL:
+ """Loads the supporting C shared library with the needed routines.
+
+ The name of the supporting C shared library is either provided by an
+ an environment variable or a default value.
Returns:
- The C shared library.
+ The supporting C shared library.
Raises:
OSError: If there is any problem in loading the shared library.
@@ -26,7 +52,7 @@ def _get_c_shared_lib(lib_name: str) -> ctypes.CDLL:
"""
# This raises OSError exception if there is any problem in loading the shared
# library.
- c_lib = ctypes.CDLL(lib_name)
+ c_lib = ctypes.CDLL(_get_support_lib_name())
try:
c_lib.convertToMLIRSparseTensor.restype = ctypes.c_void_p
@@ -44,14 +70,12 @@ def _get_c_shared_lib(lib_name: str) -> ctypes.CDLL:
def sparse_tensor_to_coo_tensor(
- lib_name: str,
sparse_tensor: ctypes.c_void_p,
dtype: np.dtype,
) -> Tuple[int, int, np.ndarray, np.ndarray, np.ndarray]:
"""Converts an MLIR sparse tensor to a COO-flavored format tensor.
Args:
- lib_name: A string for the supporting C shared library.
sparse_tensor: A ctypes.c_void_p to the MLIR sparse tensor descriptor.
dtype: The numpy data type for the tensor elements.
@@ -69,7 +93,7 @@ def sparse_tensor_to_coo_tensor(
OSError: If there is any problem in loading the shared library.
ValueError: If the shared library doesn't contain the needed routines.
"""
- c_lib = _get_c_shared_lib(lib_name)
+ c_lib = _get_c_shared_lib()
rank = ctypes.c_ulonglong(0)
nse = ctypes.c_ulonglong(0)
@@ -84,16 +108,14 @@ def sparse_tensor_to_coo_tensor(
shape = np.ctypeslib.as_array(shape, shape=[rank.value])
values = np.ctypeslib.as_array(values, shape=[nse.value])
indices = np.ctypeslib.as_array(indices, shape=[nse.value, rank.value])
- return rank, nse, shape, values, indices
+ return rank.value, nse.value, shape, values, indices
-def coo_tensor_to_sparse_tensor(lib_name: str, np_shape: np.ndarray,
- np_values: np.ndarray,
+def coo_tensor_to_sparse_tensor(np_shape: np.ndarray, np_values: np.ndarray,
np_indices: np.ndarray) -> int:
"""Converts a COO-flavored format sparse tensor to an MLIR sparse tensor.
Args:
- lib_name: A string for the supporting C shared library.
np_shape: A 1D numpy array of integers, for the shape of the tensor.
np_values: A 1D numpy array, for the non-zero values in the tensor.
np_indices: A 2D numpy array of integers, representing the indices for the
@@ -115,7 +137,136 @@ def coo_tensor_to_sparse_tensor(lib_name: str, np_shape: np.ndarray,
ctypes.POINTER(np.ctypeslib.as_ctypes_type(np_values.dtype)))
indices = np_indices.ctypes.data_as(ctypes.POINTER(ctypes.c_ulonglong))
- c_lib = _get_c_shared_lib(lib_name)
+ c_lib = _get_c_shared_lib()
ptr = c_lib.convertToMLIRSparseTensor(rank, nse, shape, values, indices)
assert ptr is not None, "Problem with calling convertToMLIRSparseTensor"
return ptr
+
+
+def compile_and_build_engine(
+ module: ir.Module) -> execution_engine.ExecutionEngine:
+ """Compiles an MLIR module and builds a JIT execution engine.
+
+ Args:
+ module: The MLIR module.
+
+ Returns:
+ A JIT execution engine for the MLIR module.
+
+ """
+ pipeline = (
+ f"sparsification,"
+ 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"arith-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")
+ PassManager.parse(pipeline).run(module)
+ return execution_engine.ExecutionEngine(
+ module, opt_level=_OPT_LEVEL, shared_libs=[_get_support_lib_name()])
+
+
+class _SparseTensorDescriptor(ctypes.Structure):
+ """A C structure for an MLIR sparse tensor."""
+ _fields_ = [
+ # A pointer for the MLIR sparse tensor storage.
+ ("storage", ctypes.POINTER(ctypes.c_ulonglong)),
+ # An MLIR MemRef descriptor for the shape of the sparse tensor.
+ ("shape", runtime.make_nd_memref_descriptor(1, ctypes.c_ulonglong)),
+ ]
+
+
+def _output_one_dim(dim: int, rank: int, shape: str) -> str:
+ """Produces the MLIR text code to output the size for the given dimension."""
+ return f"""
+ %c{dim} = arith.constant {dim} : index
+ %d{dim} = tensor.dim %t, %c{dim} : tensor<{shape}xf64, #enc>
+ memref.store %d{dim}, %b[%c{dim}] : memref<{rank}xindex>
+"""
+
+
+# TODO: With better support from MLIR, we may improve the current implementation
+# by doing the following:
+# (1) Use Python code to generate the kernel instead of doing MLIR text code
+# stitching.
+# (2) Use scf.for instead of an unrolled loop to write out the dimension sizes
+# when tensor.dim supports non-constant dimension value.
+def _get_create_sparse_tensor_kernel(
+ sparsity_codes: Sequence[sparse_tensor.DimLevelType]) -> str:
+ """Creates an MLIR text kernel to contruct a sparse tensor from a file.
+
+ The kernel returns a _SparseTensorDescriptor structure.
+ """
+ rank = len(sparsity_codes)
+
+ # Use ? to represent a dimension in the dynamic shape string representation.
+ shape = "x".join(map(lambda d: "?", range(rank)))
+
+ # Convert the encoded sparsity values to a string representation.
+ sparsity = ", ".join(
+ map(lambda s: '"compressed"' if s.value else '"dense"', sparsity_codes))
+
+ # Get the MLIR text code to write the dimension sizes to the output buffer.
+ output_dims = "\n".join(
+ map(lambda d: _output_one_dim(d, rank, shape), range(rank)))
+
+ # Return the MLIR text kernel.
+ return f"""
+!Ptr = type !llvm.ptr<i8>
+#enc = #sparse_tensor.encoding<{{
+ dimLevelType = [ {sparsity} ]
+}}>
+func @{_ENTRY_NAME}(%filename: !Ptr) -> (tensor<{shape}xf64, #enc>, memref<{rank}xindex>)
+attributes {{ llvm.emit_c_interface }} {{
+ %t = sparse_tensor.new %filename : !Ptr to tensor<{shape}xf64, #enc>
+ %b = memref.alloc() : memref<{rank}xindex>
+ {output_dims}
+ return %t, %b : tensor<{shape}xf64, #enc>, memref<{rank}xindex>
+}}"""
+
+
+def create_sparse_tensor(
+ filename: str, sparsity: Sequence[sparse_tensor.DimLevelType]
+) -> Tuple[ctypes.c_void_p, np.ndarray]:
+ """Creates an MLIR sparse tensor from the input file.
+
+ Args:
+ filename: A string for the name of the file that contains the tensor data in
+ a COO-flavored format.
+ sparsity: A sequence of DimLevelType values, one for each dimension of the
+ tensor.
+
+ Returns:
+ A Tuple containing the following values:
+ storage: A ctypes.c_void_p for the MLIR sparse tensor storage.
+ shape: A 1D numpy array of integers, for the shape of the tensor.
+
+ Raises:
+ OSError: If there is any problem in loading the supporting C shared library.
+ ValueError: If the shared library doesn't contain the needed routine.
+ """
+ with ir.Context() as ctx, ir.Location.unknown():
+ module = _get_create_sparse_tensor_kernel(sparsity)
+ module = ir.Module.parse(module)
+ engine = compile_and_build_engine(module)
+
+ # A sparse tensor descriptor to receive the kernel result.
+ c_tensor_desc = _SparseTensorDescriptor()
+ # Convert the filename to a byte stream.
+ c_filename = ctypes.c_char_p(bytes(filename, "utf-8"))
+
+ arg_pointers = [
+ ctypes.byref(ctypes.pointer(c_tensor_desc)),
+ ctypes.byref(c_filename)
+ ]
+
+ # Invoke the execution engine to run the module and return the result.
+ engine.invoke(_ENTRY_NAME, *arg_pointers)
+ shape = runtime.ranked_memref_to_numpy(ctypes.pointer(c_tensor_desc.shape))
+ return c_tensor_desc.storage, shape
diff --git a/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_utils.py b/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_utils.py
new file mode 100644
index 0000000000000..273b913b3a205
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_utils.py
@@ -0,0 +1,121 @@
+# RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
+
+from typing import Sequence
+import dataclasses
+import numpy as np
+import os
+import sys
+import tempfile
+
+from mlir.dialects import sparse_tensor
+
+_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
+sys.path.append(_SCRIPT_PATH)
+from tools import mlir_pytaco
+from tools import mlir_pytaco_utils as pytaco_utils
+
+# Define the aliases to shorten the code.
+_COMPRESSED = mlir_pytaco.ModeFormat.COMPRESSED
+_DENSE = mlir_pytaco.ModeFormat.DENSE
+
+
+def _to_string(s: Sequence[int]) -> str:
+ """Converts a sequence of integer to a space separated value string."""
+ return " ".join(map(lambda e: str(e), s))
+
+
+def _add_one(s: Sequence[int]) -> Sequence[int]:
+ """Adds one to each element in the sequence of integer."""
+ return [i + 1 for i in s]
+
+
+ at dataclasses.dataclass(frozen=True)
+class _SparseTensorCOO:
+ """Values for a COO-flavored format sparse tensor.
+
+ Attributes:
+ rank: An integer rank for the tensor.
+ nse: An integer for the number of non-zero values.
+ shape: A sequence of integer for the dimension size.
+ values: A sequence of float for the non-zero values of the tensor.
+ indices: A sequence of coordinate, each coordinate is a sequence of integer.
+ """
+ rank: int
+ nse: int
+ shape: Sequence[int]
+ values: Sequence[float]
+ indices: Sequence[Sequence[int]]
+
+
+def _coo_values_to_tns_format(t: _SparseTensorCOO) -> str:
+ """Converts a sparse tensor COO-flavored values to TNS text format."""
+ # The coo_value_str contains one line for each (coordinate value) pair.
+ # Indices are 1-based in TNS text format but 0-based in MLIR.
+ coo_value_str = "\n".join(
+ map(lambda i: _to_string(_add_one(t.indices[i])) + " " + str(t.values[i]),
+ range(t.nse)))
+
+ # Returns the TNS text format representation for the tensor.
+ return f"""{t.rank} {t.nse}
+{_to_string(t.shape)}
+{coo_value_str}
+"""
+
+
+def _implement_read_tns_test(
+ t: _SparseTensorCOO,
+ sparsity_codes: Sequence[sparse_tensor.DimLevelType]) -> int:
+ tns_data = _coo_values_to_tns_format(t)
+
+ # Write sparse tensor data to a file.
+ with tempfile.TemporaryDirectory() as test_dir:
+ file_name = os.path.join(test_dir, "data.tns")
+ with open(file_name, "w") as file:
+ file.write(tns_data)
+
+ # Read the data from the file and construct an MLIR sparse tensor.
+ sparse_tensor, o_shape = pytaco_utils.create_sparse_tensor(
+ file_name, sparsity_codes)
+
+ passed = 0
+
+ # Verify the output shape for the tensor.
+ if np.allclose(o_shape, t.shape):
+ passed += 1
+
+ # Use the output MLIR sparse tensor pointer to retrieve the COO-flavored
+ # values and verify the values.
+ o_rank, o_nse, o_shape, o_values, o_indices = (
+ pytaco_utils.sparse_tensor_to_coo_tensor(sparse_tensor, np.float64))
+ if o_rank == t.rank and o_nse == t.nse and np.allclose(
+ o_shape, t.shape) and np.allclose(o_values, t.values) and np.allclose(
+ o_indices, t.indices):
+ passed += 1
+
+ return passed
+
+
+# A 2D sparse tensor data in COO-flavored format.
+_rank = 2
+_nse = 3
+_shape = [4, 5]
+_values = [3.0, 2.0, 4.0]
+_indices = [[0, 4], [1, 0], [3, 1]]
+
+_t = _SparseTensorCOO(_rank, _nse, _shape, _values, _indices)
+_s = [_COMPRESSED, _COMPRESSED]
+# CHECK: PASSED 2D: 2
+print("PASSED 2D: ", _implement_read_tns_test(_t, _s))
+
+
+# A 3D sparse tensor data in COO-flavored format.
+_rank = 3
+_nse = 3
+_shape = [2, 5, 4]
+_values = [3.0, 2.0, 4.0]
+_indices = [[0, 4, 3], [1, 3, 0], [1, 3, 1]]
+
+_t = _SparseTensorCOO(_rank, _nse, _shape, _values, _indices)
+_s = [_DENSE, _COMPRESSED, _COMPRESSED]
+# CHECK: PASSED 3D: 2
+print("PASSED 3D: ", _implement_read_tns_test(_t, _s))
More information about the Mlir-commits
mailing list