[mlir][sparse][taco] Support f16.

[bixia]

Revert an unintentional change.

[aartbik]

Solid! Thank you.

[bixia]

Rebase.`

mlir/test/Integration/Dialect/SparseTensor/taco/test_tensor_types.py

 # RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
 
 import numpy as np
 import os
 import sys
 
 _SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(_SCRIPT_PATH)
 from tools import mlir_pytaco_api as pt
 
 compressed = pt.compressed
 dense = pt.dense
 
 passed = 0
-all_types = [pt.int8, pt.int16, pt.int32, pt.int64, pt.float32, pt.float64]
+all_types = [
+    pt.int8, pt.int16, pt.int32, pt.int64, pt.float16, pt.float32, pt.float64
+]
 for t in all_types:
   i, j = pt.get_index_vars(2)
   A = pt.tensor([2, 3], dtype=t)
   B = pt.tensor([2, 3], dtype=t)
   C = pt.tensor([2, 3], compressed, dtype=t)
   A.insert([0, 1], 10)
   A.insert([1, 2], 40)
   B.insert([0, 0], 20)
   B.insert([1, 2], 30)
   C[i, j] = A[i, j] + B[i, j]
 
   indices, values = C.get_coordinates_and_values()
   passed += isinstance(values[0], t.value)
   passed += np.array_equal(indices, [[0, 0], [0, 1], [1, 2]])
   passed += np.allclose(values, [20.0, 10.0, 70.0])
 
-# CHECK: Number of passed: 18
+# CHECK: Number of passed: 21
 print("Number of passed:", passed)

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

class Type(enum.Enum):
   """The data types supported by TACO.
 
   We use numpy data types to implement the enum data types.
   """
   INT8 = np.int8
   INT16 = np.int16
   INT32 = np.int32
   INT64 = np.int64
-  # numpy _ctype_from_dtype_scalar can't handle np.float16 yet.
+  FLOAT16 = np.float16
   FLOAT32 = np.float32
   FLOAT64 = np.float64
   COMPLEX64 = np.complex64
   COMPLEX128 = np.complex128
 
 
 # All floating point type enums.
-_FLOAT_TYPES = (Type.FLOAT32, Type.FLOAT64)
+_FLOAT_TYPES = (Type.FLOAT16, Type.FLOAT32, Type.FLOAT64)
 # All integral type enums.
 _INT_TYPES = (Type.INT8, Type.INT16, Type.INT32, Type.INT64)
 # All complex type enums.
 _COMPLEX_TYPES = (Type.COMPLEX64, Type.COMPLEX128)
 # Type alias for any numpy type used to implement the runtime support for the
 # enum data types.
-_AnyRuntimeType = Union[np.int8, np.int16, np.int32, np.int64, np.float32,
-                        np.float64, np.complex64, np.complex128]
+_AnyRuntimeType = Union[np.int8, np.int16, np.int32, np.int64, np.float16,
+                        np.float32, np.float64, np.complex64, np.complex128]
 
 
 @dataclasses.dataclass(frozen=True)
 class DType:
   """The data type class.
 
   We support the TACO API dtype class with an alias of this class.
 

 
 def _dtype_to_mlir_str(dtype: DType) -> str:
   """Returns the MLIR string for the given dtype."""
   dtype_to_str = {
       Type.INT16: "i8",
       Type.INT16: "i16",
       Type.INT32: "i32",
       Type.INT64: "i64",
+      Type.FLOAT16: "f16",
       Type.FLOAT32: "f32",
       Type.FLOAT64: "f64",
       Type.COMPLEX64: "complex<f32>",
       Type.COMPLEX128: "complex<f64>"
   }
   return dtype_to_str[dtype.kind]
 
 
 def _nptype_to_taco_type(ty: np.dtype) -> DType:
   """Returns the TACO type for the given numpy type."""
   nptype_to_dtype = {
       np.int8: Type.INT8,
       np.int16: Type.INT16,
       np.int32: Type.INT32,
       np.int64: Type.INT64,
+      np.float16: Type.FLOAT16,
       np.float32: Type.FLOAT32,
       np.float64: Type.FLOAT64,
       np.complex64: Type.COMPLEX64,
       np.complex128: Type.COMPLEX128
   }
   return DType(nptype_to_dtype[ty])
 
 
 def _mlir_type_from_taco_type(dtype: DType) -> ir.Type:
   """Returns the MLIR type corresponding to the given TACO type."""
   dtype_to_irtype = {
       Type.INT8: ir.IntegerType.get_signless(8),
       Type.INT16: ir.IntegerType.get_signless(16),
       Type.INT32: ir.IntegerType.get_signless(32),
       Type.INT64: ir.IntegerType.get_signless(64),
+      Type.FLOAT16: ir.F16Type.get(),
       Type.FLOAT32: ir.F32Type.get(),
       Type.FLOAT64: ir.F64Type.get(),
       Type.COMPLEX64: ir.ComplexType.get(ir.F32Type.get()),
       Type.COMPLEX128: ir.ComplexType.get(ir.F64Type.get())
   }
   return dtype_to_irtype[dtype.kind]
 
 def _ctype_pointer_from_array(array: np.ndarray) -> ctypes.pointer:

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_api.py

 index_expression = mlir_pytaco.IndexExpr
 access = mlir_pytaco.Access
 
 # Data type constants defined by PyTACO API.
 int8 = mlir_pytaco.DType(mlir_pytaco.Type.INT8)
 int16 = mlir_pytaco.DType(mlir_pytaco.Type.INT16)
 int32 = mlir_pytaco.DType(mlir_pytaco.Type.INT32)
 int64 = mlir_pytaco.DType(mlir_pytaco.Type.INT64)
+float16 = mlir_pytaco.DType(mlir_pytaco.Type.FLOAT16)
 float32 = mlir_pytaco.DType(mlir_pytaco.Type.FLOAT32)
 float64 = mlir_pytaco.DType(mlir_pytaco.Type.FLOAT64)
 complex64 = mlir_pytaco.DType(mlir_pytaco.Type.COMPLEX64)
 complex128 = mlir_pytaco.DType(mlir_pytaco.Type.COMPLEX128)
 
 # Storage format constants defined by the PyTACO API. In PyTACO, each storage
 # format constant has two aliasing names.
 compressed = mlir_pytaco.ModeFormat.COMPRESSED
 Compressed = mlir_pytaco.ModeFormat.COMPRESSED
 dense = mlir_pytaco.ModeFormat.DENSE
 Dense = mlir_pytaco.ModeFormat.DENSE

mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco_utils.py

try:
     support_types = [(np.int8, c_lib.convertToMLIRSparseTensorI8,
                       c_lib.convertFromMLIRSparseTensorI8),
                      (np.int16, c_lib.convertToMLIRSparseTensorI16,
                       c_lib.convertFromMLIRSparseTensorI16),
                      (np.int32, c_lib.convertToMLIRSparseTensorI32,
                       c_lib.convertFromMLIRSparseTensorI32),
                      (np.int64, c_lib.convertToMLIRSparseTensorI64,
                       c_lib.convertFromMLIRSparseTensorI64),
+                     (np.float16, c_lib.convertToMLIRSparseTensorF16,
+                      c_lib.convertFromMLIRSparseTensorF16),
                      (np.float32, c_lib.convertToMLIRSparseTensorF32,
                       c_lib.convertFromMLIRSparseTensorF32),
                      (np.float64, c_lib.convertToMLIRSparseTensorF64,
                       c_lib.convertFromMLIRSparseTensorF64),
                      (np.complex64, c_lib.convertToMLIRSparseTensorC32,
                       c_lib.convertFromMLIRSparseTensorC32),
                      (np.complex128, c_lib.convertToMLIRSparseTensorC64,
                       c_lib.convertFromMLIRSparseTensorC64)]