device_matrix_data and mtx_io

ginkgo-project · Oct 24, 2024 · eb75b06 · eb75b06
1 parent 013584d
commit eb75b06
Show file tree

Hide file tree

Showing 8 changed files with 60 additions and 30 deletions.
diff --git a/common/cuda_hip/base/device_matrix_data_kernels.cpp b/common/cuda_hip/base/device_matrix_data_kernels.cpp
@@ -12,6 +12,7 @@
 #include <thrust/sort.h>
 #include <thrust/tuple.h>
 
+#include "common/cuda_hip/base/math.hpp"
 #include "common/cuda_hip/base/thrust.hpp"
 #include "common/cuda_hip/base/types.hpp"
 
@@ -22,6 +23,15 @@ namespace GKO_DEVICE_NAMESPACE {
 namespace components {
 
 
+// __half != only in __device__
+// Although gko::is_nonzero is constexpr, it still shows calling __device__ in
+// __host__
+template <typename T>
+GKO_INLINE __device__ constexpr bool is_nonzero(T value)
+{
+    return value != zero<T>();
+}
+
 template <typename ValueType, typename IndexType>
 void remove_zeros(std::shared_ptr<const DefaultExecutor> exec,
                   array<ValueType>& values, array<IndexType>& row_idxs,
@@ -58,7 +68,7 @@ void remove_zeros(std::shared_ptr<const DefaultExecutor> exec,
     }
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_REMOVE_ZEROS_KERNEL);
 
 
@@ -102,7 +112,7 @@ void sum_duplicates(std::shared_ptr<const DefaultExecutor> exec, size_type,
     }
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_SUM_DUPLICATES_KERNEL);
 
 
@@ -117,7 +127,7 @@ void sort_row_major(std::shared_ptr<const DefaultExecutor> exec,
                         it + data.get_num_stored_elements(), vals);
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_SORT_ROW_MAJOR_KERNEL);
 
 

diff --git a/common/unified/base/device_matrix_data_kernels.cpp b/common/unified/base/device_matrix_data_kernels.cpp
@@ -30,7 +30,7 @@ void soa_to_aos(std::shared_ptr<const DefaultExecutor> exec,
         in.get_const_col_idxs(), in.get_const_values(), out);
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_SOA_TO_AOS_KERNEL);
 
 
@@ -50,7 +50,7 @@ void aos_to_soa(std::shared_ptr<const DefaultExecutor> exec,
         out.get_values());
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_AOS_TO_SOA_KERNEL);
 
 

diff --git a/core/base/device_matrix_data.cpp b/core/base/device_matrix_data.cpp
@@ -157,7 +157,8 @@ device_matrix_data<ValueType, IndexType>::empty_out()
 
 #define GKO_DECLARE_DEVICE_MATRIX_DATA(ValueType, IndexType) \
     struct device_matrix_data<ValueType, IndexType>
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_DEVICE_MATRIX_DATA);
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
+    GKO_DECLARE_DEVICE_MATRIX_DATA);
 
 
 }  // namespace gko
diff --git a/core/base/mtx_io.cpp b/core/base/mtx_io.cpp
@@ -14,6 +14,7 @@
 #include <type_traits>
 
 #include <ginkgo/core/base/exception_helpers.hpp>
+#include <ginkgo/core/base/half.hpp>
 #include <ginkgo/core/base/math.hpp>
 #include <ginkgo/core/base/utils.hpp>
 
@@ -757,19 +758,28 @@ static constexpr uint64 binary_format_magic()
 {
     constexpr auto is_int = std::is_same<IndexType, int32>::value;
     constexpr auto is_long = std::is_same<IndexType, int64>::value;
+    constexpr auto is_half = std::is_same<ValueType, half>::value;
     constexpr auto is_double = std::is_same<ValueType, double>::value;
     constexpr auto is_float = std::is_same<ValueType, float>::value;
     constexpr auto is_complex_double =
         std::is_same<ValueType, std::complex<double>>::value;
     constexpr auto is_complex_float =
         std::is_same<ValueType, std::complex<float>>::value;
+    constexpr auto is_complex_half =
+        std::is_same<ValueType, std::complex<half>>::value;
     static_assert(is_int || is_long, "invalid storage index type");
-    static_assert(
-        is_double || is_float || is_complex_double || is_complex_float,
-        "invalid storage value type");
+    static_assert(is_half || is_complex_half || is_double || is_float ||
+                      is_complex_double || is_complex_float,
+                  "invalid storage value type");
     constexpr auto index_bit = is_int ? 'I' : 'L';
     constexpr auto value_bit =
-        is_double ? 'D' : (is_float ? 'S' : (is_complex_double ? 'Z' : 'C'));
+        is_double
+            ? 'D'
+            : (is_float
+                   ? 'S'
+                   : (is_complex_double
+                          ? 'Z'
+                          : (is_complex_float ? 'C' : (is_half ? 'H' : 'X'))));
     constexpr uint64 shift = 256;
     constexpr uint64 type_bits = index_bit * shift + value_bit;
     return 'G' +
@@ -879,12 +889,16 @@ matrix_data<ValueType, IndexType> read_binary_raw(std::istream& is)
     }
     DECLARE_OVERLOAD(double, int32)
     DECLARE_OVERLOAD(float, int32)
+    DECLARE_OVERLOAD(half, int32)
     DECLARE_OVERLOAD(std::complex<double>, int32)
     DECLARE_OVERLOAD(std::complex<float>, int32)
+    DECLARE_OVERLOAD(std::complex<half>, int32)
     DECLARE_OVERLOAD(double, int64)
     DECLARE_OVERLOAD(float, int64)
+    DECLARE_OVERLOAD(half, int64)
     DECLARE_OVERLOAD(std::complex<double>, int64)
     DECLARE_OVERLOAD(std::complex<float>, int64)
+    DECLARE_OVERLOAD(std::complex<half>, int64)
 #undef DECLARE_OVERLOAD
     else
     {
@@ -970,11 +984,14 @@ void write_raw(std::ostream& os, const matrix_data<ValueType, IndexType>& data,
                           const matrix_data<ValueType, IndexType>& data)
 #define GKO_DECLARE_READ_GENERIC_RAW(ValueType, IndexType) \
     matrix_data<ValueType, IndexType> read_generic_raw(std::istream& is)
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_READ_RAW);
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_WRITE_RAW);
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_READ_BINARY_RAW);
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_WRITE_BINARY_RAW);
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_READ_GENERIC_RAW);
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(GKO_DECLARE_READ_RAW);
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(GKO_DECLARE_WRITE_RAW);
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
+    GKO_DECLARE_READ_BINARY_RAW);
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
+    GKO_DECLARE_WRITE_BINARY_RAW);
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
+    GKO_DECLARE_READ_GENERIC_RAW);
 
 
 }  // namespace gko
diff --git a/core/device_hooks/common_kernels.inc.cpp b/core/device_hooks/common_kernels.inc.cpp
@@ -241,14 +241,16 @@ GKO_STUB_TEMPLATE_TYPE_WITH_HALF(GKO_DECLARE_FILL_SEQ_ARRAY_KERNEL);
 GKO_STUB_TEMPLATE_TYPE_WITH_HALF(GKO_DECLARE_REDUCE_ADD_ARRAY_KERNEL);
 GKO_STUB_VALUE_TYPE_WITH_HALF(GKO_DECLARE_INPLACE_ABSOLUTE_ARRAY_KERNEL);
 GKO_STUB_VALUE_TYPE_WITH_HALF(GKO_DECLARE_OUTPLACE_ABSOLUTE_ARRAY_KERNEL);
-GKO_STUB_VALUE_AND_INDEX_TYPE(
+GKO_STUB_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_REMOVE_ZEROS_KERNEL);
-GKO_STUB_VALUE_AND_INDEX_TYPE(
+GKO_STUB_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_SUM_DUPLICATES_KERNEL);
-GKO_STUB_VALUE_AND_INDEX_TYPE(
+GKO_STUB_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_SORT_ROW_MAJOR_KERNEL);
-GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_DEVICE_MATRIX_DATA_AOS_TO_SOA_KERNEL);
-GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_DEVICE_MATRIX_DATA_SOA_TO_AOS_KERNEL);
+GKO_STUB_VALUE_AND_INDEX_TYPE_WITH_HALF(
+    GKO_DECLARE_DEVICE_MATRIX_DATA_AOS_TO_SOA_KERNEL);
+GKO_STUB_VALUE_AND_INDEX_TYPE_WITH_HALF(
+    GKO_DECLARE_DEVICE_MATRIX_DATA_SOA_TO_AOS_KERNEL);
 
 template <typename IndexType, typename RowPtrType>
 GKO_DECLARE_CONVERT_PTRS_TO_IDXS(IndexType, RowPtrType)

diff --git a/dpcpp/base/device_matrix_data_kernels.dp.cpp b/dpcpp/base/device_matrix_data_kernels.dp.cpp
@@ -49,7 +49,7 @@ void remove_zeros(std::shared_ptr<const DefaultExecutor> exec,
     }
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_REMOVE_ZEROS_KERNEL);
 
 
@@ -112,7 +112,7 @@ void sort_row_major(std::shared_ptr<const DefaultExecutor> exec,
               });
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_SORT_ROW_MAJOR_KERNEL);
 
 

diff --git a/omp/base/device_matrix_data_kernels.cpp b/omp/base/device_matrix_data_kernels.cpp
@@ -69,7 +69,7 @@ void remove_zeros(std::shared_ptr<const DefaultExecutor> exec,
     }
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_REMOVE_ZEROS_KERNEL);
 
 
@@ -127,7 +127,7 @@ void sum_duplicates(std::shared_ptr<const DefaultExecutor> exec,
     }
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_SUM_DUPLICATES_KERNEL);
 
 
@@ -142,7 +142,7 @@ void sort_row_major(std::shared_ptr<const DefaultExecutor> exec,
     aos_to_soa(exec, tmp, data);
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_SORT_ROW_MAJOR_KERNEL);
 
 

diff --git a/reference/base/device_matrix_data_kernels.cpp b/reference/base/device_matrix_data_kernels.cpp
@@ -29,7 +29,7 @@ void soa_to_aos(std::shared_ptr<const DefaultExecutor> exec,
     }
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_SOA_TO_AOS_KERNEL);
 
 
@@ -46,7 +46,7 @@ void aos_to_soa(std::shared_ptr<const DefaultExecutor> exec,
     }
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_AOS_TO_SOA_KERNEL);
 
 
@@ -78,7 +78,7 @@ void remove_zeros(std::shared_ptr<const DefaultExecutor> exec,
     }
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_REMOVE_ZEROS_KERNEL);
 
 
@@ -127,7 +127,7 @@ void sum_duplicates(std::shared_ptr<const DefaultExecutor> exec, size_type,
     }
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_SUM_DUPLICATES_KERNEL);
 
 
@@ -142,7 +142,7 @@ void sort_row_major(std::shared_ptr<const DefaultExecutor> exec,
     aos_to_soa(exec, tmp, data);
 }
 
-GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE_WITH_HALF(
     GKO_DECLARE_DEVICE_MATRIX_DATA_SORT_ROW_MAJOR_KERNEL);