UTs

ROCm · Feb 5, 2022 · bad1aa2 · bad1aa2
1 parent 0ccd72f
commit bad1aa2
Showing 1 changed file with 101 additions and 0 deletions.
diff --git a/orttraining/orttraining/test/training_ops/cpu/reduction/reduction_ops_test.cc b/orttraining/orttraining/test/training_ops/cpu/reduction/reduction_ops_test.cc
@@ -6,6 +6,7 @@
 #include <type_traits>
 #include "gtest/gtest.h"
 #include "test/providers/provider_test_utils.h"
+#include "test/common/cuda_op_test_utils.h"
 
 namespace onnxruntime {
 namespace test {
@@ -163,6 +164,106 @@ TEST_P(ReductionOpTest, ReduceAllL2HalfFloat) {
 }
 #endif
 
+#if defined(USE_CUDA) || defined(USE_ROCM)
+TEST_P(ReductionOpTest, ReduceAllL2_BFloat16_BFloat16) {
+#ifdef USE_CUDA
+  int min_cuda_architecture = 530;
+  if (!HasCudaEnvironment(min_cuda_architecture)) {
+    LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16";
+    return;
+  }
+#endif
+  OpTester test("ReduceAllL2", 1, onnxruntime::kMSDomain, true);
+  test.SetDeterminism(GetParam());
+
+  std::vector<float> data0 = {1.0f, 2.0f, 3.0f};
+  std::vector<BFloat16> data0_bf16 = FloatsToBFloat16s(data0);
+
+  std::vector<float> data1 = {-1.0f, -2.0f};
+  std::vector<BFloat16> data1_bf16 = FloatsToBFloat16s(data1);
+
+  std::vector<float> result = {4.358898943540674f};
+  std::vector<BFloat16> result_bf16 = FloatsToBFloat16s(result);
+
+  test.AddInput<BFloat16>("data0", {3}, data0_bf16);
+  test.AddInput<BFloat16>("data1", {2}, data1_bf16);
+
+  test.AddOutput<BFloat16>("reduced", {}, result_bf16);
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+#ifdef USE_CUDA
+  execution_providers.push_back(DefaultCudaExecutionProvider());
+#elif USE_ROCM
+  execution_providers.push_back(DefaultRocmExecutionProvider());
+#endif 
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+}
+
+TEST_P(ReductionOpTest, ReduceAllL2_BFloat16_Float) {
+#ifdef USE_CUDA
+  int min_cuda_architecture = 530;
+  if (!HasCudaEnvironment(min_cuda_architecture)) {
+    LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16";
+    return;
+  }
+#endif
+  OpTester test("ReduceAllL2", 1, onnxruntime::kMSDomain, true);
+  test.SetDeterminism(GetParam());
+
+  std::vector<float> data0 = {1.0f, 2.0f, 3.0f};
+  std::vector<BFloat16> data0_bf16 = FloatsToBFloat16s(data0);
+
+  std::vector<float> data1 = {-1.0f, -2.0f};
+  std::vector<BFloat16> data1_bf16 = FloatsToBFloat16s(data1);
+
+  std::vector<float> result = {4.358898943540674f};
+
+  test.AddInput<BFloat16>("data0", {3}, data0_bf16);
+  test.AddInput<BFloat16>("data1", {2}, data1_bf16);
+
+  test.AddOutput<float>("reduced", {}, result);
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+#ifdef USE_CUDA
+  execution_providers.push_back(DefaultCudaExecutionProvider());
+#elif USE_ROCM
+  execution_providers.push_back(DefaultRocmExecutionProvider());
+#endif 
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+}
+
+TEST_P(ReductionOpTest, ReduceAllL2_Float_BFloat16) {
+#ifdef USE_CUDA
+  int min_cuda_architecture = 530;
+  if (!HasCudaEnvironment(min_cuda_architecture)) {
+    LOGS_DEFAULT(WARNING) << "Hardware NOT support BFP16";
+    return;
+  }
+#endif
+  OpTester test("ReduceAllL2", 1, onnxruntime::kMSDomain, true);
+  test.SetDeterminism(GetParam());
+
+  std::vector<float> data0 = {1.0f, 2.0f, 3.0f};
+  std::vector<float> data1 = {-1.0f, -2.0f};
+
+  std::vector<float> result = {4.358898943540674f};
+  std::vector<BFloat16> result_bf16 = FloatsToBFloat16s(result);
+
+  test.AddInput<float>("data0", {3}, data0);
+  test.AddInput<float>("data1", {2}, data1);
+
+  test.AddOutput<BFloat16>("reduced", {}, result_bf16);
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+#ifdef USE_CUDA
+  execution_providers.push_back(DefaultCudaExecutionProvider());
+#elif USE_ROCM
+  execution_providers.push_back(DefaultRocmExecutionProvider());
+#endif 
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+}
+#endif
+
 void TestMultiTensorReduce(
     const int tensor_count,
     const int min_tensor_size,