Add tests for OCP FP8 vector_type storage

test/data_type/test_custom_type.cpp

@@ -872,3 +872,153 @@ TEST(Complex_half, TestAsTypeReshape)
                   test_vec.at(num_elem * i + 1));
     });
 }
+
+#if CK_USE_OCP_FP8
+
+TEST(FP8OCP, TestSize)
+{
+    static_assert(std::is_same_v<f8_t, ck::f8_ocp_t>, "OCP FP8 is not enabled");
+    ASSERT_EQ(sizeof(f8_t), sizeof(ck::fp8_storage_t));
+    ASSERT_EQ(sizeof(vector_type<f8_t, 2>), sizeof(vector_type<ck::fp8_storage_t, 2>));
+    ASSERT_EQ(sizeof(vector_type<f8_t, 4>), sizeof(vector_type<ck::fp8_storage_t, 4>));
+    ASSERT_EQ(sizeof(vector_type<f8_t, 8>), sizeof(vector_type<ck::fp8_storage_t, 8>));
+    ASSERT_EQ(sizeof(vector_type<f8_t, 16>), sizeof(vector_type<ck::fp8_storage_t, 16>));
+    ASSERT_EQ(sizeof(vector_type<f8_t, 32>), sizeof(vector_type<ck::fp8_storage_t, 32>));
+    ASSERT_EQ(sizeof(vector_type<f8_t, 64>), sizeof(vector_type<ck::fp8_storage_t, 64>));
+}
+
+TEST(FP8OCP, TestAsType)
+{
+    static_assert(std::is_same_v<f8_t, ck::f8_ocp_t>, "OCP FP8 is not enabled");
+
+    // test size
+    std::array<float, 8> test_vec = {-4, -2, -0.5, -0.25, 1.0 / 8.0, 1, 1.5, 16};
+    constexpr int size            = test_vec.size();
+
+    // reference vector
+    vector_type<f8_t, size> right_vec;
+
+    // check default CTOR
+    ck::static_for<0, size, 1>{}(
+        [&](auto i) { ASSERT_EQ(right_vec.template AsType<f8_t>()(Number<i>{}), f8_t{0}); });
+
+    // assign test values to the vector
+    ck::static_for<0, size, 1>{}([&](auto i) {
+        right_vec.template AsType<f8_t>()(Number<i>{}) = ck::type_convert<f8_t>(test_vec.at(i));
+    });
+
+    // copy the vector
+    vector_type<f8_t, size> left_vec{right_vec};
+
+    // check if values were copied correctly
+    ck::static_for<0, size, 1>{}([&](auto i) {
+        ASSERT_EQ(left_vec.template AsType<f8_t>()(Number<i>{}),
+                  ck::type_convert<f8_t>(test_vec.at(i)));
+    });
+}
+
+TEST(FP8OCP, TestAsTypeReshape)
+{
+    static_assert(std::is_same_v<f8_t, ck::f8_ocp_t>, "OCP FP8 is not enabled");
+
+    // test size
+    std::array<float, 8> test_vec = {-8, -0.5, -0.25, 1.0 / 8.0, 1 / 256, 1, 1.5, 16};
+    constexpr int size            = test_vec.size();
+
+    // reference vector
+    vector_type<f8_t, size> right_vec;
+
+    // check default CTOR
+    ck::static_for<0, size, 1>{}(
+        [&](auto i) { ASSERT_EQ(right_vec.template AsType<f8_t>()(Number<i>{}), f8_t{0}); });
+
+    // assign test values to the vector
+    ck::static_for<0, size, 1>{}([&](auto i) {
+        right_vec.template AsType<f8_t>()(Number<i>{}) = ck::type_convert<f8_t>(test_vec.at(i));
+    });
+
+    // copy the first half of a vector
+    vector_type<f8_t, size / 2> left_vec{
+        right_vec.template AsType<vector_type<f8_t, size / 2>::type>()(Number<0>{})};
+
+    // check if values were copied correctly
+    ck::static_for<0, size / 2, 1>{}([&](auto i) {
+        ASSERT_EQ(left_vec.template AsType<f8_t>()(Number<i>{}),
+                  ck::type_convert<f8_t>(test_vec.at(i)));
+    });
+}
+
+TEST(BF8OCP, TestSize)
+{
+    static_assert(std::is_same_v<bf8_t, ck::bf8_ocp_t>, "OCP BF8 is not enabled");
+    ASSERT_EQ(sizeof(bf8_t), sizeof(ck::fp8_storage_t));
+    ASSERT_EQ(sizeof(vector_type<bf8_t, 2>), sizeof(vector_type<ck::fp8_storage_t, 2>));
+    ASSERT_EQ(sizeof(vector_type<bf8_t, 4>), sizeof(vector_type<ck::fp8_storage_t, 4>));
+    ASSERT_EQ(sizeof(vector_type<bf8_t, 8>), sizeof(vector_type<ck::fp8_storage_t, 8>));
+    ASSERT_EQ(sizeof(vector_type<bf8_t, 16>), sizeof(vector_type<ck::fp8_storage_t, 16>));
+    ASSERT_EQ(sizeof(vector_type<bf8_t, 32>), sizeof(vector_type<ck::fp8_storage_t, 32>));
+    ASSERT_EQ(sizeof(vector_type<bf8_t, 64>), sizeof(vector_type<ck::fp8_storage_t, 64>));
+}
+
+TEST(BF8OCP, TestAsType)
+{
+    static_assert(std::is_same_v<bf8_t, ck::bf8_ocp_t>, "OCP BF8 is not enabled");
+
+    // test size
+    std::array<float, 8> test_vec = {-4, -2, -0.5, -0.25, 1.0 / 8.0, 1, 1.5, 16};
+    constexpr int size            = test_vec.size();
+
+    // reference vector
+    vector_type<bf8_t, size> right_vec;
+
+    // check default CTOR
+    ck::static_for<0, size, 1>{}(
+        [&](auto i) { ASSERT_EQ(right_vec.template AsType<bf8_t>()(Number<i>{}), bf8_t{0}); });
+
+    // assign test values to the vector
+    ck::static_for<0, size, 1>{}([&](auto i) {
+        right_vec.template AsType<bf8_t>()(Number<i>{}) = ck::type_convert<bf8_t>(test_vec.at(i));
+    });
+
+    // copy the vector
+    vector_type<bf8_t, size> left_vec{right_vec};
+
+    // check if values were copied correctly
+    ck::static_for<0, size, 1>{}([&](auto i) {
+        ASSERT_EQ(left_vec.template AsType<bf8_t>()(Number<i>{}),
+                  ck::type_convert<bf8_t>(test_vec.at(i)));
+    });
+}
+
+TEST(BF8OCP, TestAsTypeReshape)
+{
+    static_assert(std::is_same_v<bf8_t, ck::bf8_ocp_t>, "OCP BF8 is not enabled");
+
+    // test size
+    std::array<float, 8> test_vec = {-8, -0.5, -0.25, 1.0 / 8.0, 1 / 256, 1, 1.5, 16};
+    constexpr int size            = test_vec.size();
+
+    // reference vector
+    vector_type<bf8_t, size> right_vec;
+
+    // check default CTOR
+    ck::static_for<0, size, 1>{}(
+        [&](auto i) { ASSERT_EQ(right_vec.template AsType<bf8_t>()(Number<i>{}), bf8_t{0}); });
+
+    // assign test values to the vector
+    ck::static_for<0, size, 1>{}([&](auto i) {
+        right_vec.template AsType<bf8_t>()(Number<i>{}) = ck::type_convert<bf8_t>(test_vec.at(i));
+    });
+
+    // copy the first half of a vector
+    vector_type<bf8_t, size / 2> left_vec{
+        right_vec.template AsType<vector_type<bf8_t, size / 2>::type>()(Number<0>{})};
+
+    // check if values were copied correctly
+    ck::static_for<0, size / 2, 1>{}([&](auto i) {
+        ASSERT_EQ(left_vec.template AsType<bf8_t>()(Number<i>{}),
+                  ck::type_convert<bf8_t>(test_vec.at(i)));
+    });
+}
+
+#endif