diff --git a/.github/workflows/ut.yml b/.github/workflows/ut.yml index 2fbe2733..4e60fde7 100644 --- a/.github/workflows/ut.yml +++ b/.github/workflows/ut.yml @@ -72,11 +72,26 @@ jobs: cd build ARK_UT_CORRECTNESS=1 ARK_ROOT=$PWD ctest -L correctness --stop-on-failure --verbose --schedule-random + - name: Install Python Dependencies + if: github.event_name != 'schedule' + run: | + python3 -m pip install -r requirements.txt + + - name: Run Python UT + if: github.event_name != 'schedule' + run: | + cd build + PYTHONPATH=$PWD/python ARK_ROOT=$PWD python3 -m pytest \ + --cov=python/ark \ + --cov-report lcov:py_coverage.info \ + --verbose \ + ../python/unittest/ + - name: C++ Coverage if: github.event_name != 'schedule' run: | cd build - lcov --ignore-errors negative --capture --directory . --output-file cpp_coverage.info + lcov --ignore-errors negative,mismatch --capture --directory . --output-file cpp_coverage.info lcov --ignore-errors unused --remove cpp_coverage.info \ '/usr/*' \ '/tmp/*' \ @@ -90,21 +105,6 @@ jobs: --output-file cpp_coverage.info lcov --list cpp_coverage.info - - name: Install Python Dependencies - if: github.event_name != 'schedule' - run: | - python3 -m pip install -r requirements.txt - - - name: Run Python UT - if: github.event_name != 'schedule' - run: | - cd build - PYTHONPATH=$PWD/python ARK_ROOT=$PWD python3 -m pytest \ - --cov=python/ark \ - --cov-report lcov:py_coverage.info \ - --verbose \ - ../python/unittest/ - - name: Report Coverage if: github.event_name != 'schedule' env: diff --git a/ark/CMakeLists.txt b/ark/CMakeLists.txt index 225bcfe4..97b32399 100644 --- a/ark/CMakeLists.txt +++ b/ark/CMakeLists.txt @@ -79,21 +79,10 @@ if(ARK_BUILD_TESTS) endforeach() # Tests that call op_test() and require ARK_UT_CORRECTNESS=1 for real validation. - # Keep in sync: grep -l 'op_test(' ark/ops/*_test.cpp + # Remaining C++ correctness tests (most ops migrated to python/unittest/ops/). set(CORRECTNESS_TESTS ops_all_reduce_test - ops_arithmetic_test - ops_cast_test ops_copy_test - ops_embedding_test - ops_layernorm_test - ops_math_test - ops_matmul_test - ops_reduce_test - ops_rope_test - ops_scalar_test - ops_softmax_test - ops_transpose_test ) set_tests_properties(${CORRECTNESS_TESTS} PROPERTIES LABELS "correctness") endif() diff --git a/ark/ops/ops_arithmetic_test.cpp b/ark/ops/ops_arithmetic_test.cpp deleted file mode 100644 index 6a878c66..00000000 --- a/ark/ops/ops_arithmetic_test.cpp +++ /dev/null @@ -1,451 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -#include "ops_test_common.hpp" - -template -void baseline_add(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, int) { - T *out = static_cast(outputs[0]); - T *t0 = static_cast(inputs[0]); - T *t1 = static_cast(inputs[1]); - - // NumPy-style broadcasted addition - ark::Dims osh = output_shapes[0].dims4(); - ark::Dims ish0 = input_shapes[0].dims4(); - ark::Dims ish1 = input_shapes[1].dims4(); - for (ark::DimType n = 0; n < osh[0]; ++n) { - for (ark::DimType c = 0; c < osh[1]; ++c) { - for (ark::DimType h = 0; h < osh[2]; ++h) { - for (ark::DimType w = 0; w < osh[3]; ++w) { - out[w + h * osh[3] + c * osh[2] * osh[3] + - n * osh[1] * osh[2] * osh[3]] = - t0[(w % ish0[3]) + (h % ish0[2]) * ish0[3] + - (c % ish0[1]) * ish0[2] * ish0[3] + - (n % ish0[0]) * ish0[1] * ish0[2] * ish0[3]] + - t1[(w % ish1[3]) + (h % ish1[2]) * ish1[3] + - (c % ish1[1]) * ish1[2] * ish1[3] + - (n % ish1[0]) * ish1[1] * ish1[2] * ish1[3]]; - } - } - } - } -}; - -template -void baseline_sub(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, int) { - T *out = static_cast(outputs[0]); - T *t0 = static_cast(inputs[0]); - T *t1 = static_cast(inputs[1]); - - // NumPy-style broadcasted addition - ark::Dims osh = output_shapes[0].dims4(); - ark::Dims ish0 = input_shapes[0].dims4(); - ark::Dims ish1 = input_shapes[1].dims4(); - for (ark::DimType n = 0; n < osh[0]; ++n) { - for (ark::DimType c = 0; c < osh[1]; ++c) { - for (ark::DimType h = 0; h < osh[2]; ++h) { - for (ark::DimType w = 0; w < osh[3]; ++w) { - out[w + h * osh[3] + c * osh[2] * osh[3] + - n * osh[1] * osh[2] * osh[3]] = - t0[(w % ish0[3]) + (h % ish0[2]) * ish0[3] + - (c % ish0[1]) * ish0[2] * ish0[3] + - (n % ish0[0]) * ish0[1] * ish0[2] * ish0[3]] - - t1[(w % ish1[3]) + (h % ish1[2]) * ish1[3] + - (c % ish1[1]) * ish1[2] * ish1[3] + - (n % ish1[0]) * ish1[1] * ish1[2] * ish1[3]]; - } - } - } - } -}; - -template -void baseline_mul(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, int) { - T *out = static_cast(outputs[0]); - T *t0 = static_cast(inputs[0]); - T *t1 = static_cast(inputs[1]); - - // NumPy-style broadcasted multiplication - ark::Dims osh = output_shapes[0].dims4(); - ark::Dims ish0 = input_shapes[0].dims4(); - ark::Dims ish1 = input_shapes[1].dims4(); - for (ark::DimType n = 0; n < osh[0]; ++n) { - for (ark::DimType c = 0; c < osh[1]; ++c) { - for (ark::DimType h = 0; h < osh[2]; ++h) { - for (ark::DimType w = 0; w < osh[3]; ++w) { - out[w + h * osh[3] + c * osh[2] * osh[3] + - n * osh[1] * osh[2] * osh[3]] = - t0[(w % ish0[3]) + (h % ish0[2]) * ish0[3] + - (c % ish0[1]) * ish0[2] * ish0[3] + - (n % ish0[0]) * ish0[1] * ish0[2] * ish0[3]] * - t1[(w % ish1[3]) + (h % ish1[2]) * ish1[3] + - (c % ish1[1]) * ish1[2] * ish1[3] + - (n % ish1[0]) * ish1[1] * ish1[2] * ish1[3]]; - } - } - } - } -}; - -template -void baseline_div(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, int) { - T *out = static_cast(outputs[0]); - T *t0 = static_cast(inputs[0]); - T *t1 = static_cast(inputs[1]); - - // NumPy-style broadcasted division - ark::Dims osh = output_shapes[0].dims4(); - ark::Dims ish0 = input_shapes[0].dims4(); - ark::Dims ish1 = input_shapes[1].dims4(); - for (ark::DimType n = 0; n < osh[0]; ++n) { - for (ark::DimType c = 0; c < osh[1]; ++c) { - for (ark::DimType h = 0; h < osh[2]; ++h) { - for (ark::DimType w = 0; w < osh[3]; ++w) { - out[w + h * osh[3] + c * osh[2] * osh[3] + - n * osh[1] * osh[2] * osh[3]] = - t0[(w % ish0[3]) + (h % ish0[2]) * ish0[3] + - (c % ish0[1]) * ish0[2] * ish0[3] + - (n % ish0[0]) * ish0[1] * ish0[2] * ish0[3]] / - t1[(w % ish1[3]) + (h % ish1[2]) * ish1[3] + - (c % ish1[1]) * ish1[2] * ish1[3] + - (n % ish1[0]) * ish1[1] * ish1[2] * ish1[3]]; - } - } - } - } -}; - -ark::unittest::State test_add_fp32() { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP32); - ark::Tensor t1 = m.tensor({8192}, ark::FP32); - ark::Tensor out = m.add(t0, t1); - - auto result = - ark::op_test("add_fp32", m, {t0, t1}, {out}, baseline_add); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_add_fp16() { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP16); - ark::Tensor t1 = m.tensor({8192}, ark::FP16); - ark::Tensor out = m.add(t0, t1); - - auto result = - ark::op_test("add_fp16", m, {t0, t1}, {out}, baseline_add); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_add_bf16() { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::BF16); - ark::Tensor t1 = m.tensor({8192}, ark::BF16); - ark::Tensor out = m.add(t0, t1); - - auto result = ark::op_test("add_bf16", m, {t0, t1}, {out}, - baseline_add); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_add_overwrite() { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP16); - ark::Tensor t1 = m.tensor({8192}, ark::FP16); - ark::Tensor out = m.add(t0, t1, t1); - - auto result = ark::op_test("add_overwrite", m, {t0, t1}, {out}, - baseline_add); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_add_broadcast() { - { - ark::Model m; - ark::Tensor t0 = m.tensor({4, 1024}, ark::FP16); - ark::Tensor t1 = m.tensor({1, 1024}, ark::FP16); - ark::Tensor out = m.add(t0, t1); - - auto result = ark::op_test("add_broadcast", m, {t0, t1}, {out}, - baseline_add); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({4, 64}, ark::FP16); - ark::Tensor t1 = m.tensor({4, 1}, ark::FP16, {4, 2}); - ark::Tensor out = m.add(t0, t1); - - auto result = ark::op_test("add_broadcast", m, {t0, t1}, {out}, - baseline_add); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({3, 1, 1024}, ark::FP16); - ark::Tensor t1 = m.tensor({1, 4, 1}, ark::FP16, {1, 4, 2}); - ark::Tensor out = m.add(t0, t1); - - auto result = ark::op_test("add_broadcast", m, {t0, t1}, {out}, - baseline_add); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_add_offset() { - { - ark::Model m; - ark::Tensor t0 = m.tensor({2, 64}, ark::FP16, {4, 128}, {2, 64}); - ark::Tensor t1 = m.tensor({2, 64}, ark::FP16); - ark::Tensor out = m.add(t0, t1); - - auto result = ark::op_test("add_offset", m, {t0, t1}, {out}, - baseline_add); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_add_invalid() { - { - ark::Model m; - ark::Tensor t0 = m.tensor({1024}, ark::FP16); - ark::Tensor t1 = m.tensor({1024}, ark::FP32); - UNITTEST_THROW(m.add(t0, t1), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP16); - ark::Tensor t1 = m.tensor({8192}, ark::FP16); - ark::Tensor out = m.tensor({8192}, ark::FP32); - UNITTEST_THROW(m.add(t0, t1, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP16); - ark::Tensor t1 = m.tensor({8192}, ark::FP16); - ark::Tensor out = m.tensor({1024}, ark::FP16); - UNITTEST_THROW(m.add(t0, t1, out), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_sub_fp32() { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP32); - ark::Tensor t1 = m.tensor({8192}, ark::FP32); - ark::Tensor out = m.sub(t0, t1); - - auto result = - ark::op_test("sub_fp32", m, {t0, t1}, {out}, baseline_sub); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_sub_invalid() { - { - ark::Model m; - ark::Tensor t0 = m.tensor({1024}, ark::FP16); - ark::Tensor t1 = m.tensor({1024}, ark::FP32); - UNITTEST_THROW(m.sub(t0, t1), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP16); - ark::Tensor t1 = m.tensor({8192}, ark::FP16); - ark::Tensor out = m.tensor({8192}, ark::FP32); - UNITTEST_THROW(m.sub(t0, t1, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP16); - ark::Tensor t1 = m.tensor({8192}, ark::FP16); - ark::Tensor out = m.tensor({1024}, ark::FP16); - UNITTEST_THROW(m.sub(t0, t1, out), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_mul_fp32() { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP32); - ark::Tensor t1 = m.tensor({8192}, ark::FP32); - ark::Tensor out = m.mul(t0, t1); - - auto result = - ark::op_test("mul_fp32", m, {t0, t1}, {out}, baseline_mul); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_mul_fp16() { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP16); - ark::Tensor t1 = m.tensor({8192}, ark::FP16); - ark::Tensor out = m.mul(t0, t1); - - auto result = - ark::op_test("mul_fp16", m, {t0, t1}, {out}, baseline_mul); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_mul_overwrite() { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP16); - ark::Tensor t1 = m.tensor({8192}, ark::FP16); - ark::Tensor out = m.mul(t0, t1, t1); - - auto result = ark::op_test("mul_overwrite", m, {t0, t1}, {out}, - baseline_mul); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_mul_broadcast() { - { - ark::Model m; - ark::Tensor t0 = m.tensor({4, 1024}, ark::FP16); - ark::Tensor t1 = m.tensor({1, 1024}, ark::FP16); - ark::Tensor out = m.mul(t0, t1); - - auto result = ark::op_test("mul_broadcast", m, {t0, t1}, {out}, - baseline_mul); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({4, 1024}, ark::FP16); - ark::Tensor t1 = m.tensor({4, 1}, ark::FP16, {4, 2}); - ark::Tensor out = m.mul(t0, t1); - - auto result = ark::op_test("mul_broadcast", m, {t0, t1}, {out}, - baseline_mul); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({3, 1, 1024}, ark::FP16); - ark::Tensor t1 = m.tensor({1, 4, 1}, ark::FP16, {1, 4, 2}); - ark::Tensor out = m.mul(t0, t1); - - auto result = ark::op_test("mul_broadcast", m, {t0, t1}, {out}, - baseline_mul); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_mul_invalid() { - { - ark::Model m; - ark::Tensor t0 = m.tensor({1024}, ark::FP16); - ark::Tensor t1 = m.tensor({1024}, ark::FP32); - UNITTEST_THROW(m.mul(t0, t1), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP16); - ark::Tensor t1 = m.tensor({8192}, ark::FP16); - ark::Tensor out = m.tensor({8192}, ark::FP32); - UNITTEST_THROW(m.mul(t0, t1, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP16); - ark::Tensor t1 = m.tensor({8192}, ark::FP16); - ark::Tensor out = m.tensor({1024}, ark::FP16); - UNITTEST_THROW(m.mul(t0, t1, out), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_div_fp32() { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP32); - ark::Tensor t1 = m.tensor({8192}, ark::FP32); - ark::Tensor out = m.div(t0, t1); - - auto result = - ark::op_test("div_fp32", m, {t0, t1}, {out}, baseline_div); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_div_invalid() { - { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP32); - ark::Tensor t1 = m.tensor({8192}, ark::FP16); - UNITTEST_THROW(m.div(t0, t1), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP32); - ark::Tensor t1 = m.tensor({8192}, ark::FP32); - ark::Tensor out = m.tensor({8192}, ark::FP16); - UNITTEST_THROW(m.div(t0, t1, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({8192}, ark::FP32); - ark::Tensor t1 = m.tensor({8192}, ark::FP32); - ark::Tensor out = m.tensor({1024}, ark::FP16); - UNITTEST_THROW(m.div(t0, t1, out), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -int main() { - ark::init(); - UNITTEST(test_add_fp32); - UNITTEST(test_add_fp16); - UNITTEST(test_add_bf16); - UNITTEST(test_add_overwrite); - UNITTEST(test_add_broadcast); - UNITTEST(test_add_offset); - UNITTEST(test_add_invalid); - UNITTEST(test_sub_fp32); - UNITTEST(test_sub_invalid); - UNITTEST(test_mul_fp32); - UNITTEST(test_mul_fp16); - UNITTEST(test_mul_overwrite); - UNITTEST(test_mul_broadcast); - UNITTEST(test_mul_invalid); - UNITTEST(test_div_fp32); - UNITTEST(test_div_invalid); - return ark::unittest::SUCCESS; -} diff --git a/ark/ops/ops_cast_test.cpp b/ark/ops/ops_cast_test.cpp deleted file mode 100644 index 8404e07f..00000000 --- a/ark/ops/ops_cast_test.cpp +++ /dev/null @@ -1,322 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -#include "ops_test_common.hpp" - -template -void baseline_cast(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &, int) { - ToType *out = static_cast(outputs[0]); - FromType *input = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0]; - for (ark::DimType i = 0; i < osh.nelems(); ++i) { - out[i] = ToType(input[i]); - } -}; - -template -void baseline_cast_from_byte(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &, int) { - ToType *out = static_cast(outputs[0]); - // input is a byte array, but force read it as ToType. - ToType *input = reinterpret_cast(inputs[0]); - ark::Dims osh = output_shapes[0]; - for (ark::DimType i = 0; i < osh.nelems(); ++i) { - out[i] = input[i]; - } -}; - -template -void baseline_cast_to_byte(std::vector &outputs, - const std::vector &, - const std::vector &inputs, - const std::vector &input_shapes, int) { - // output is a byte array, but force write it as FromType. - FromType *out = reinterpret_cast(outputs[0]); - FromType *input = static_cast(inputs[0]); - ark::Dims ish = input_shapes[0]; - for (ark::DimType i = 0; i < ish.nelems(); ++i) { - out[i] = input[i]; - } -}; - -ark::unittest::State test_cast_fp16_to_fp32() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::FP16); - ark::Tensor out = m.cast(t, ark::FP32); - - auto result = ark::op_test("cast_fp16_to_fp32", m, {t}, {out}, - baseline_cast); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_fp16_to_int32() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::FP16); - ark::Tensor out = m.cast(t, ark::INT32); - - std::vector input_data(t.shape().nelems()); - for (size_t i = 0; i < input_data.size(); ++i) { - input_data[i] = ark::half_t(int((i + 1) % 1000)); - } - - auto result = - ark::op_test("cast_fp16_to_int32", m, {t}, {out}, - baseline_cast, {input_data.data()}); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_fp32_to_fp16() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::FP32); - ark::Tensor out = m.cast(t, ark::FP16); - - auto result = ark::op_test("cast_fp32_to_fp16", m, {t}, {out}, - baseline_cast); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_fp32_to_int32() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::FP32); - ark::Tensor out = m.cast(t, ark::INT32); - - std::vector input_data(t.shape().nelems()); - for (size_t i = 0; i < input_data.size(); ++i) { - input_data[i] = float((i + 1) % 1000); - } - - auto result = ark::op_test("cast_fp32_to_int32", m, {t}, {out}, - baseline_cast, {input_data.data()}); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_int32_to_fp32() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::INT32); - ark::Tensor out = m.cast(t, ark::FP32); - - std::vector input_data(t.shape().nelems()); - for (size_t i = 0; i < input_data.size(); ++i) { - input_data[i] = (i + 1) % 1000; - } - - auto result = ark::op_test("cast_int32_to_fp32", m, {t}, {out}, - baseline_cast, {input_data.data()}); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_int32_to_fp16() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::INT32); - ark::Tensor out = m.cast(t, ark::FP16); - - std::vector input_data(t.shape().nelems()); - for (size_t i = 0; i < input_data.size(); ++i) { - input_data[i] = (i + 1) % 1000; - } - - auto result = - ark::op_test("cast_int32_to_fp16", m, {t}, {out}, - baseline_cast, {input_data.data()}); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_byte_to_fp32() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::BYTE); - ark::Tensor out = m.cast(t, ark::FP32); - - // For preventing optimize-out - m.noop(t); - m.noop(out); - - auto result = ark::op_test("cast_byte_to_fp32", m, {t}, {out}, - baseline_cast_from_byte); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_byte_to_fp16() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::BYTE); - ark::Tensor out = m.cast(t, ark::FP16); - - // For preventing optimize-out - m.noop(t); - m.noop(out); - - auto result = ark::op_test("cast_byte_to_fp16", m, {t}, {out}, - baseline_cast_from_byte); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_byte_to_int32() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::BYTE); - ark::Tensor out = m.cast(t, ark::INT32); - - // For preventing optimize-out - m.noop(t); - m.noop(out); - - auto result = ark::op_test("cast_byte_to_int32", m, {t}, {out}, - baseline_cast_from_byte); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_fp32_to_byte() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::FP32); - ark::Tensor out = m.cast(t, ark::BYTE); - - // For preventing optimize-out - m.noop(t); - m.noop(out); - - auto result = ark::op_test("cast_fp32_to_byte", m, {t}, {out}, - baseline_cast_to_byte); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_fp16_to_byte() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::FP16); - ark::Tensor out = m.cast(t, ark::BYTE); - - // For preventing optimize-out - m.noop(t); - m.noop(out); - - auto result = ark::op_test("cast_fp16_to_byte", m, {t}, {out}, - baseline_cast_to_byte); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_int32_to_byte() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::INT32); - ark::Tensor out = m.cast(t, ark::BYTE); - - // For preventing optimize-out - m.noop(t); - m.noop(out); - - auto result = ark::op_test("cast_int32_to_byte", m, {t}, {out}, - baseline_cast_to_byte); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_bf16_to_float() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::BF16); - ark::Tensor out = m.cast(t, ark::FP32); - - auto result = ark::op_test("cast_bf16_to_float", m, {t}, {out}, - baseline_cast); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_float_to_bf16() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::FP32); - ark::Tensor out = m.cast(t, ark::BF16); - - auto result = ark::op_test("cast_float_to_bf16", m, {t}, {out}, - baseline_cast); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_cast_invalid() { - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(1), ark::BYTE); - UNITTEST_THROW(m.cast(t, ark::FP32), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor(ark::Dims(4, 1), ark::BYTE); - m.cast(t0, ark::FP32); // ok - ark::Tensor t1 = m.tensor(ark::Dims(4, 1, 1), ark::BYTE); - m.cast(t1, ark::FP32); // ok - ark::Tensor t2 = m.tensor(ark::Dims(4, 1, 1, 1), ark::BYTE); - m.cast(t2, ark::FP32); // ok - ark::Tensor t3 = m.tensor(ark::Dims(7, 1), ark::BYTE); - UNITTEST_THROW(m.cast(t3, ark::FP32), ark::ModelError); - ark::Tensor t4 = m.tensor(ark::Dims(7, 1, 1), ark::BYTE); - UNITTEST_THROW(m.cast(t4, ark::FP32), ark::ModelError); - ark::Tensor t5 = m.tensor(ark::Dims(7, 1, 1, 1), ark::BYTE); - UNITTEST_THROW(m.cast(t5, ark::FP32), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({8, 1}, ark::BYTE); - m.cast(t0, ark::FP32); // ok - ark::Tensor t1 = m.tensor({8, 1}, ark::BYTE, {13, 1}, {0, 0}, {9, 1}); - UNITTEST_THROW(m.cast(t1, ark::FP32), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({8, 1}, ark::FP16); - ark::Tensor out = m.tensor({8, 1}, ark::INT32); - UNITTEST_THROW(m.cast(t0, ark::FP32, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t0 = m.tensor({8, 1}, ark::FP16); - ark::Tensor out = m.tensor({4, 1}, ark::FP32); - UNITTEST_THROW(m.cast(t0, ark::FP32, out), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -int main() { - ark::init(); - UNITTEST(test_cast_fp16_to_fp32); - UNITTEST(test_cast_fp16_to_int32); - UNITTEST(test_cast_fp32_to_fp16); - UNITTEST(test_cast_fp32_to_int32); - UNITTEST(test_cast_int32_to_fp32); - UNITTEST(test_cast_int32_to_fp16); - UNITTEST(test_cast_byte_to_fp32); - UNITTEST(test_cast_byte_to_fp16); - UNITTEST(test_cast_byte_to_int32); - UNITTEST(test_cast_fp32_to_byte); - UNITTEST(test_cast_fp16_to_byte); - UNITTEST(test_cast_int32_to_byte); - UNITTEST(test_cast_bf16_to_float); - UNITTEST(test_cast_float_to_bf16); - UNITTEST(test_cast_invalid); - return ark::unittest::SUCCESS; -} diff --git a/ark/ops/ops_embedding_test.cpp b/ark/ops/ops_embedding_test.cpp deleted file mode 100644 index 22260529..00000000 --- a/ark/ops/ops_embedding_test.cpp +++ /dev/null @@ -1,122 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -#include -#include - -#include "ark/random.hpp" -#include "ops_test_common.hpp" - -template -void baseline_embedding(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, int) { - T *out = static_cast(outputs[0]); - int *in = static_cast(inputs[0]); - T *weight = static_cast(inputs[1]); - - ark::Dims osh = output_shapes[0].dims4(); - ark::Dims wsh = input_shapes[1].dims4(); - - assert(osh[3] == wsh[3]); - - int in_idx = 0; - for (ark::DimType n = 0; n < osh[0]; ++n) { - for (ark::DimType c = 0; c < osh[1]; ++c) { - for (ark::DimType h = 0; h < osh[2]; ++h) { - int weight_idx = in[in_idx++]; - if (weight_idx < 0) { - weight_idx += wsh[2]; - } - T *ptr = &weight[weight_idx * wsh[3]]; - for (ark::DimType w = 0; w < osh[3]; ++w) { - out[n * osh[1] * osh[2] * osh[3] + c * osh[2] * osh[3] + - h * osh[3] + w] = ptr[w]; - } - } - } - } -}; - -template -ark::unittest::State test_embedding() { - const int num_emb = 100; - const int emb_dim = 4096; - - ark::DataType weight_type; - if (std::is_same::value) { - weight_type = ark::FP32; - } else { - weight_type = ark::FP16; - } - - ark::Model m; - ark::Tensor ti = m.tensor(ark::Dims(8, 3, 64), ark::INT32); - ark::Tensor tw = m.tensor(ark::Dims(num_emb, emb_dim), weight_type); - ark::Tensor to = m.embedding(ti, tw); - - std::vector ti_data; - for (auto i = 0; i < ti.shape().nelems(); ++i) { - // Random indices in [0, num_emb) - int rand_idx = ark::rand() % num_emb; - if (i % 9 == 0) { - // test negative tokens (padding) - rand_idx = -rand_idx; - } - ti_data.push_back(rand_idx); - } - std::vector tw_data(tw.shape().nelems()); - for (auto i = 0; i < tw.shape().nelems(); ++i) { - tw_data[i] = ark::random(-1.0, 1.0); - } - std::string type_str = ""; - if (std::is_same::value) { - type_str = "fp32"; - } else if (std::is_same::value) { - type_str = "fp16"; - } else if (std::is_same::value) { - type_str = "bf16"; - } - auto result = - ark::op_test("embedding_" + type_str, m, {ti, tw}, {to}, - baseline_embedding, {ti_data.data(), tw_data.data()}); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_embedding_fp32() { return test_embedding(); } - -ark::unittest::State test_embedding_fp16() { - return test_embedding(); -} - -ark::unittest::State test_embedding_bf16() { - return test_embedding(); -} - -ark::unittest::State test_embedding_invalid() { - { - ark::Model m; - ark::Tensor ti = m.tensor(ark::Dims(4, 8, 3, 64), ark::INT32); - ark::Tensor tw = m.tensor(ark::Dims(100, 1024), ark::FP32); - UNITTEST_THROW(m.embedding(ti, tw), ark::ModelError); - } - { - ark::Model m; - ark::Tensor ti = m.tensor(ark::Dims(8, 3, 64), ark::INT32); - ark::Tensor tw = m.tensor(ark::Dims(2, 100, 1024), ark::FP32); - UNITTEST_THROW(m.embedding(ti, tw), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -int main() { - ark::init(); - UNITTEST(test_embedding_fp32); - UNITTEST(test_embedding_fp16); - UNITTEST(test_embedding_bf16); - UNITTEST(test_embedding_invalid); - return ark::unittest::SUCCESS; -} diff --git a/ark/ops/ops_layernorm_test.cpp b/ark/ops/ops_layernorm_test.cpp deleted file mode 100644 index 91dacd0a..00000000 --- a/ark/ops/ops_layernorm_test.cpp +++ /dev/null @@ -1,227 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -#include - -#include "ops_test_common.hpp" - -// Baseline: LayerNorm with affine (gamma, beta). -// For each row (last dim = W), compute: -// mean = sum(x) / W -// var = sum((x - mean)^2) / W -// out = gamma * (x - mean) / sqrt(var + eps) + beta -// eps = 1e-5 (standard). -template -void baseline_layernorm(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - T *gamma = static_cast(inputs[1]); - T *beta = static_cast(inputs[2]); - - ark::Dims osh = output_shapes[0]; - ark::DimType total = osh.nelems(); - ark::DimType W = osh[-1]; - ark::DimType num_rows = total / W; - constexpr float eps = 1e-5f; - - for (ark::DimType row = 0; row < num_rows; ++row) { - T *row_in = input + row * W; - T *row_out = out + row * W; - - // mean - float mean = 0; - for (ark::DimType j = 0; j < W; ++j) { - mean += static_cast(row_in[j]); - } - mean /= static_cast(W); - - // variance - float var = 0; - for (ark::DimType j = 0; j < W; ++j) { - float diff = static_cast(row_in[j]) - mean; - var += diff * diff; - } - var /= static_cast(W); - - float inv_std = 1.0f / std::sqrt(var + eps); - - // normalize + affine - for (ark::DimType j = 0; j < W; ++j) { - float normed = (static_cast(row_in[j]) - mean) * inv_std; - row_out[j] = T(static_cast(gamma[j]) * normed + - static_cast(beta[j])); - } - } -} - -ark::unittest::State test_layernorm_fp32() { - ark::Model m; - ark::Tensor input = m.tensor({4, 1024}, ark::FP32); - ark::Tensor gamma = m.tensor({1024}, ark::FP32); - ark::Tensor beta = m.tensor({1024}, ark::FP32); - ark::Tensor out = m.layernorm(input, gamma, beta); - - auto result = ark::op_test("layernorm_fp32", m, {input, gamma, beta}, {out}, - baseline_layernorm); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-4f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_layernorm_fp16() { - ark::Model m; - ark::Tensor input = m.tensor({2, 768}, ark::FP16); - ark::Tensor gamma = m.tensor({768}, ark::FP16); - ark::Tensor beta = m.tensor({768}, ark::FP16); - ark::Tensor out = m.layernorm(input, gamma, beta); - - auto result = ark::op_test("layernorm_fp16", m, {input, gamma, beta}, {out}, - baseline_layernorm); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-2f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_layernorm_bf16() { - ark::Model m; - ark::Tensor input = m.tensor({2, 768}, ark::BF16); - ark::Tensor gamma = m.tensor({768}, ark::BF16); - ark::Tensor beta = m.tensor({768}, ark::BF16); - ark::Tensor out = m.layernorm(input, gamma, beta); - - auto result = ark::op_test("layernorm_bf16", m, {input, gamma, beta}, {out}, - baseline_layernorm); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 5e-2f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_layernorm_batch() { - // Higher-dimensional input: [B, S, D] - ark::Model m; - ark::Tensor input = m.tensor({2, 8, 512}, ark::FP32); - ark::Tensor gamma = m.tensor({512}, ark::FP32); - ark::Tensor beta = m.tensor({512}, ark::FP32); - ark::Tensor out = m.layernorm(input, gamma, beta); - - auto result = ark::op_test("layernorm_batch", m, {input, gamma, beta}, - {out}, baseline_layernorm); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-4f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_layernorm_small_row() { - // Small last dim — tests edge case where W < warp size - ark::Model m; - ark::Tensor input = m.tensor({8, 16}, ark::FP32); - ark::Tensor gamma = m.tensor({16}, ark::FP32); - ark::Tensor beta = m.tensor({16}, ark::FP32); - ark::Tensor out = m.layernorm(input, gamma, beta); - - auto result = ark::op_test("layernorm_small_row", m, {input, gamma, beta}, - {out}, baseline_layernorm); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-4f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_layernorm_non_pow2() { - // Non-power-of-2 W — exercises boundary guards (idx_w + j < W) - ark::Model m; - ark::Tensor input = m.tensor({4, 127}, ark::FP32); - ark::Tensor gamma = m.tensor({127}, ark::FP32); - ark::Tensor beta = m.tensor({127}, ark::FP32); - ark::Tensor out = m.layernorm(input, gamma, beta); - - auto result = ark::op_test("layernorm_non_pow2", m, {input, gamma, beta}, - {out}, baseline_layernorm); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-4f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_layernorm_w1() { - // W=1 boundary: variance=0, epsilon dominates - ark::Model m; - ark::Tensor input = m.tensor({4, 1}, ark::FP32); - ark::Tensor gamma = m.tensor({1}, ark::FP32); - ark::Tensor beta = m.tensor({1}, ark::FP32); - ark::Tensor out = m.layernorm(input, gamma, beta); - - auto result = ark::op_test("layernorm_w1", m, {input, gamma, beta}, {out}, - baseline_layernorm); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-4f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_layernorm_large_w() { - // Large inner dim — exercises numerical stability for W=4096 - ark::Model m; - ark::Tensor input = m.tensor({1, 1, 1, 4096}, ark::FP32); - ark::Tensor gamma = m.tensor({4096}, ark::FP32); - ark::Tensor beta = m.tensor({4096}, ark::FP32); - ark::Tensor out = m.layernorm(input, gamma, beta); - - auto result = ark::op_test("layernorm_large_w", m, {input, gamma, beta}, - {out}, baseline_layernorm); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-4f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_layernorm_invalid() { - // gamma shape mismatch - { - ark::Model m; - ark::Tensor input = m.tensor({4, 1024}, ark::FP32); - ark::Tensor gamma = m.tensor({512}, ark::FP32); // wrong size - ark::Tensor beta = m.tensor({1024}, ark::FP32); - UNITTEST_THROW(m.layernorm(input, gamma, beta), ark::ModelError); - } - // beta shape mismatch - { - ark::Model m; - ark::Tensor input = m.tensor({4, 1024}, ark::FP32); - ark::Tensor gamma = m.tensor({1024}, ark::FP32); - ark::Tensor beta = m.tensor({512}, ark::FP32); // wrong size - UNITTEST_THROW(m.layernorm(input, gamma, beta), ark::ModelError); - } - // data type mismatch (gamma) - { - ark::Model m; - ark::Tensor input = m.tensor({4, 1024}, ark::FP32); - ark::Tensor gamma = m.tensor({1024}, ark::FP16); // wrong type - ark::Tensor beta = m.tensor({1024}, ark::FP32); - UNITTEST_THROW(m.layernorm(input, gamma, beta), ark::ModelError); - } - // output shape mismatch - { - ark::Model m; - ark::Tensor input = m.tensor({4, 1024}, ark::FP32); - ark::Tensor gamma = m.tensor({1024}, ark::FP32); - ark::Tensor beta = m.tensor({1024}, ark::FP32); - ark::Tensor bad_out = m.tensor({4, 512}, ark::FP32); // wrong shape - UNITTEST_THROW(m.layernorm(input, gamma, beta, bad_out), - ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -int main() { - ark::init(); - UNITTEST(test_layernorm_fp32); - UNITTEST(test_layernorm_fp16); - UNITTEST(test_layernorm_bf16); - UNITTEST(test_layernorm_batch); - UNITTEST(test_layernorm_small_row); - UNITTEST(test_layernorm_non_pow2); - UNITTEST(test_layernorm_w1); - UNITTEST(test_layernorm_large_w); - UNITTEST(test_layernorm_invalid); - return ark::unittest::SUCCESS; -} diff --git a/ark/ops/ops_math_test.cpp b/ark/ops/ops_math_test.cpp deleted file mode 100644 index 74b351f6..00000000 --- a/ark/ops/ops_math_test.cpp +++ /dev/null @@ -1,379 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -#include - -#include "ark/model.hpp" -#include "ops_test_common.hpp" -#include "unittest/unittest_utils.h" - -float gelu(float x) { - return 0.5 * x * (1 + tanh(sqrt(2 / M_PI) * (x + 0.044715 * pow(x, 3)))); -} - -template -void baseline_gelu(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &, int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0]; - for (ark::DimType i = 0; i < osh.nelems(); ++i) { - out[i] = gelu(input[i]); - } -}; - -template -void baseline_exp(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &, int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0]; - for (ark::DimType i = 0; i < osh.nelems(); ++i) { - out[i] = std::exp(input[i]); - } -}; - -float relu(float x) { return x > 0 ? x : 0; } - -template -void baseline_relu(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &, int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0]; - for (ark::DimType i = 0; i < osh.nelems(); ++i) { - out[i] = relu(input[i]); - } -}; - -template -void baseline_rsqrt(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &, int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0]; - for (ark::DimType i = 0; i < osh.nelems(); ++i) { - out[i] = 1.0f / std::sqrt(input[i]); - } -}; - -float sigmoid(float x) { return 1 / (1 + std::exp(-x)); } - -template -void baseline_sigmoid(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &, int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0]; - for (ark::DimType i = 0; i < osh.nelems(); ++i) { - out[i] = sigmoid(input[i]); - } -}; - -template -void baseline_sqrt(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &, int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0]; - for (ark::DimType i = 0; i < osh.nelems(); ++i) { - out[i] = std::sqrt(input[i]); - } -}; - -std::vector stable_positive_radical_inputs(ark::DimType nelems) { - static const float kInputs[] = {0.25f, 1.0f, 4.0f, 16.0f}; - std::vector data(nelems); - for (ark::DimType i = 0; i < nelems; ++i) { - data[i] = kInputs[i % 4]; - } - return data; -} - -ark::unittest::State test_gelu_fp32() { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::FP32); - ark::Tensor out = m.gelu(t); - - auto result = - ark::op_test("gelu_fp32", m, {t}, {out}, baseline_gelu); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-6f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_gelu_bf16() { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::BF16); - ark::Tensor out = m.gelu(t); - - auto result = ark::op_test("gelu_bf16", m, {t}, {out}, - baseline_gelu); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-6f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_gelu_invalid() { - { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::BF16); - ark::Tensor out = m.tensor({4, 2, 1024}, ark::FP32); - UNITTEST_THROW(m.gelu(t, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::BF16); - ark::Tensor out = m.tensor({4, 4, 1024}, ark::BF16); - UNITTEST_THROW(m.gelu(t, out), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_exp_fp32() { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::FP32); - ark::Tensor out = m.exp(t); - - auto result = ark::op_test("exp_fp32", m, {t}, {out}, baseline_exp); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-5f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_exp_fp16() { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::FP16); - ark::Tensor out = m.exp(t); - - auto result = - ark::op_test("exp_fp16", m, {t}, {out}, baseline_exp); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-2f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_exp_bf16() { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::BF16); - ark::Tensor out = m.exp(t); - - auto result = - ark::op_test("exp_bf16", m, {t}, {out}, baseline_exp); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-2f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_exp_invalid() { - { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::BF16); - ark::Tensor out = m.tensor({4, 2, 1024}, ark::FP32); - UNITTEST_THROW(m.exp(t, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::BF16); - ark::Tensor out = m.tensor({4, 4, 1024}, ark::BF16); - UNITTEST_THROW(m.exp(t, out), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_relu_fp32() { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::FP32); - ark::Tensor out = m.relu(t); - - auto result = - ark::op_test("relu_fp32", m, {t}, {out}, baseline_relu); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_relu_fp16() { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::FP16); - ark::Tensor out = m.relu(t); - - auto result = - ark::op_test("relu_fp16", m, {t}, {out}, baseline_relu); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_relu_bf16() { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::BF16); - ark::Tensor out = m.relu(t); - - auto result = ark::op_test("relu_bf16", m, {t}, {out}, - baseline_relu); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_relu_invalid() { - { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::BF16); - ark::Tensor out = m.tensor({4, 2, 1024}, ark::FP32); - UNITTEST_THROW(m.relu(t, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::BF16); - ark::Tensor out = m.tensor({4, 4, 1024}, ark::BF16); - UNITTEST_THROW(m.relu(t, out), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_math_rsqrt_fp32() { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::FP32); - ark::Tensor out = m.rsqrt(t); - auto data = stable_positive_radical_inputs(t.shape().nelems()); - - auto result = ark::op_test("math_rsqrt_fp32", m, {t}, {out}, - baseline_rsqrt, {data.data()}); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-4f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_math_rsqrt_fp16() { - ark::Model m; - ark::Tensor t = m.tensor({1, 64, 1}, ark::FP16); - ark::Tensor out = m.rsqrt(t); - - std::vector data(64, 4); - - auto result = ark::op_test("math_rsqrt_fp16", m, {t}, {out}, - baseline_rsqrt, {data.data()}); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-4f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_sigmoid_fp32() { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::FP32); - ark::Tensor out = m.sigmoid(t); - - auto result = - ark::op_test("sigmoid_fp32", m, {t}, {out}, baseline_sigmoid); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-5f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_sigmoid_bf16() { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::BF16); - ark::Tensor out = m.sigmoid(t); - - auto result = ark::op_test("sigmoid_bf16", m, {t}, {out}, - baseline_sigmoid); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-2f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_sigmoid_invalid() { - { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::BF16); - ark::Tensor out = m.tensor({4, 2, 1024}, ark::FP32); - UNITTEST_THROW(m.sigmoid(t, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::BF16); - ark::Tensor out = m.tensor({4, 4, 1024}, ark::BF16); - UNITTEST_THROW(m.sigmoid(t, out), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_math_sqrt_fp32() { - ark::Model m; - ark::Tensor t = m.tensor({4, 2, 1024}, ark::FP32); - ark::Tensor out = m.sqrt(t); - auto data = stable_positive_radical_inputs(t.shape().nelems()); - - auto result = ark::op_test("math_sqrt_fp32", m, {t}, {out}, - baseline_sqrt, {data.data()}); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-6f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_math_sqrt_fp16_small_last_dim() { - ark::Model m; - ark::Tensor t = m.tensor({4, 1024, 1}, ark::FP16, {4, 1024, 2}); - ark::Tensor out = m.sqrt(t); - auto fp32_data = stable_positive_radical_inputs(t.shape().nelems()); - std::vector data(fp32_data.begin(), fp32_data.end()); - - auto result = ark::op_test("math_sqrt_fp16_small_last_dim", m, {t}, {out}, - baseline_sqrt, {data.data()}); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-4f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_math_sqrt_invalid() { - { - ark::Model model; - ark::Tensor input = model.tensor({1, 3, 16, 8192}, ark::FP32); - ark::Tensor output = model.tensor({1, 3, 16, 8192}, ark::FP16); - UNITTEST_THROW(model.sqrt(input, output), ark::ModelError); - } - { - ark::Model model; - ark::Tensor input = model.tensor({1, 3, 16, 8192}, ark::FP32); - ark::Tensor output = model.tensor({1, 3, 16, 1024}, ark::FP32); - UNITTEST_THROW(model.sqrt(input, output), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -int main() { - ark::init(); - UNITTEST(test_gelu_fp32); - UNITTEST(test_gelu_bf16); - UNITTEST(test_gelu_invalid); - UNITTEST(test_exp_fp32); - UNITTEST(test_exp_fp16); - UNITTEST(test_exp_invalid); - UNITTEST(test_relu_fp32); - UNITTEST(test_relu_fp16); - UNITTEST(test_relu_bf16); - UNITTEST(test_relu_invalid); - UNITTEST(test_math_rsqrt_fp32); - UNITTEST(test_math_rsqrt_fp16); - UNITTEST(test_sigmoid_fp32); - UNITTEST(test_sigmoid_bf16); - UNITTEST(test_sigmoid_invalid); - UNITTEST(test_math_sqrt_fp32); - UNITTEST(test_math_sqrt_fp16_small_last_dim); - UNITTEST(test_math_sqrt_invalid); - return ark::unittest::SUCCESS; -} diff --git a/ark/ops/ops_matmul_test.cpp b/ark/ops/ops_matmul_test.cpp deleted file mode 100644 index 3426df13..00000000 --- a/ark/ops/ops_matmul_test.cpp +++ /dev/null @@ -1,718 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -#include - -#include "gpu/gpu.hpp" -#include "logging.hpp" -#include "model/model_node.hpp" -#include "model/model_op.hpp" -#include "ops_test_common.hpp" - -#if defined(ARK_CUDA) - -#include - -typedef cublasHandle_t blasHandle; -typedef cublasStatus_t blasStatus; -typedef cublasOperation_t blasOperation; -typedef cudaDataType blasDataType; -typedef cublasComputeType_t blasComputeType; -constexpr auto blasSuccess = CUBLAS_STATUS_SUCCESS; -constexpr auto BLAS_OP_N = CUBLAS_OP_N; -constexpr auto BLAS_OP_T = CUBLAS_OP_T; -constexpr auto BLAS_R_32F = CUDA_R_32F; -constexpr auto BLAS_R_16F = CUDA_R_16F; -constexpr auto BLAS_R_16BF = CUDA_R_16BF; -constexpr auto BLAS_COMPUTE_32F = CUBLAS_COMPUTE_32F; -constexpr auto BLAS_COMPUTE_32F_FAST_TF32 = CUBLAS_COMPUTE_32F_FAST_TF32; -constexpr auto BLAS_COMPUTE_16F = CUBLAS_COMPUTE_16F; - -inline auto blasGemmEx(blasHandle handle, blasOperation transA, - blasOperation transB, int m, int n, int k, - const void *alpha, const void *A, blasDataType Atype, - int lda, const void *B, blasDataType Btype, int ldb, - const void *beta, void *C, blasDataType Ctype, int ldc, - blasComputeType computeType) { - return cublasGemmEx(handle, transA, transB, m, n, k, alpha, A, Atype, lda, - B, Btype, ldb, beta, C, Ctype, ldc, computeType, - CUBLAS_GEMM_DEFAULT); -} - -inline auto blasGemmStridedBatchedEx( - blasHandle handle, blasOperation transA, blasOperation transB, int m, int n, - int k, const void *alpha, const void *A, blasDataType Atype, int lda, - int strideA, const void *B, blasDataType Btype, int ldb, int strideB, - const void *beta, void *C, blasDataType Ctype, int ldc, int strideC, - int batchCount, blasComputeType computeType) { - return cublasGemmStridedBatchedEx( - handle, transA, transB, m, n, k, alpha, A, Atype, lda, strideA, B, - Btype, ldb, strideB, beta, C, Ctype, ldc, strideC, batchCount, - computeType, CUBLAS_GEMM_DEFAULT); -} - -#elif defined(ARK_ROCM) - -#include - -typedef rocblas_handle blasHandle; -typedef rocblas_status blasStatus; -typedef rocblas_operation blasOperation; -typedef rocblas_datatype blasDataType; -typedef rocblas_datatype blasComputeType; -constexpr auto blasSuccess = rocblas_status_success; -constexpr auto BLAS_OP_N = rocblas_operation_none; -constexpr auto BLAS_OP_T = rocblas_operation_transpose; -constexpr auto BLAS_R_32F = rocblas_datatype_f32_r; -constexpr auto BLAS_R_16F = rocblas_datatype_f16_r; -constexpr auto BLAS_R_16BF = rocblas_datatype_bf16_r; -constexpr auto BLAS_COMPUTE_32F = rocblas_datatype_f32_r; -[[maybe_unused]] constexpr auto BLAS_COMPUTE_32F_FAST_TF32 = - rocblas_datatype_f32_r; -[[maybe_unused]] constexpr auto BLAS_COMPUTE_16F = rocblas_datatype_f16_r; - -inline auto blasGemmEx(blasHandle handle, blasOperation transA, - blasOperation transB, int m, int n, int k, - const void *alpha, const void *A, blasDataType Atype, - int lda, const void *B, blasDataType Btype, int ldb, - const void *beta, void *C, blasDataType Ctype, int ldc, - blasComputeType computeType) { - return rocblas_gemm_ex(handle, transA, transB, m, n, k, alpha, A, Atype, - lda, B, Btype, ldb, beta, C, Ctype, ldc, C, Ctype, - ldc, computeType, rocblas_gemm_algo_standard, 0, 0); -} - -inline auto blasGemmStridedBatchedEx( - blasHandle handle, blasOperation transA, blasOperation transB, int m, int n, - int k, const void *alpha, const void *A, blasDataType Atype, int lda, - int strideA, const void *B, blasDataType Btype, int ldb, int strideB, - const void *beta, void *C, blasDataType Ctype, int ldc, int strideC, - int batchCount, blasComputeType computeType) { - return rocblas_gemm_strided_batched_ex( - handle, transA, transB, m, n, k, alpha, A, Atype, lda, strideA, B, - Btype, ldb, strideB, beta, C, Ctype, ldc, strideC, C, Ctype, ldc, - strideC, batchCount, computeType, rocblas_gemm_algo_standard, 0, 0); -} - -#endif - -ARK_GPU_DEFINE_FUNC_ALIAS(blasCreate, cublasCreate, rocblas_create_handle); -ARK_GPU_DEFINE_FUNC_ALIAS(blasDestroy, cublasDestroy, rocblas_destroy_handle); - -class BlasHandle { - public: - BlasHandle() { - if (blasCreate(&handle_) != blasSuccess) { - throw std::runtime_error("Failed to create blas handle"); - } - } - - ~BlasHandle() { - if (blasDestroy(handle_) != blasSuccess) { - // do nothing. - } - } - - blasHandle get() const { return handle_; } - - private: - blasHandle handle_; -}; - -static BlasHandle globalBlasHandle; - -template -void blas_matmul(int m, int n, int k, const DataType *a, int lda, - const DataType *b, int ldb, DataType *c, int ldc, - int batch_size = 1) { - static_assert(std::is_same_v || - std::is_same_v || - std::is_same_v, - "Unsupported data type"); - - auto blasH = globalBlasHandle.get(); - blasStatus status; - blasOperation optypeA = (blasOperation)BlasOpTypeA; - blasOperation optypeB = (blasOperation)BlasOpTypeB; - -#if defined(ARK_CUDA) - using CompType = - typename std::conditional_t, - ark::half_t, float>; - blasComputeType ctype = - std::is_same_v - ? BLAS_COMPUTE_32F_FAST_TF32 - : (std::is_same_v ? BLAS_COMPUTE_16F - : BLAS_COMPUTE_32F); -#elif defined(ARK_ROCM) - // CK uses only fp32 compute type for fp16/bf16 - using CompType = float; - blasComputeType ctype = BLAS_COMPUTE_32F; -#endif - CompType alpha = 1; - CompType beta = 0; - - blasDataType dtype = - std::is_same_v - ? BLAS_R_32F - : (std::is_same_v ? BLAS_R_16F - : BLAS_R_16BF); - if (batch_size == 1) { - status = blasGemmEx(blasH, optypeB, optypeA, n, m, k, &alpha, b, dtype, - ldb, a, dtype, lda, &beta, c, dtype, ldc, ctype); - if (status != blasSuccess) { - throw std::runtime_error("Failed to call blasGemmEx"); - } - } else { - status = blasGemmStridedBatchedEx( - blasH, optypeB, optypeA, n, m, k, &alpha, b, dtype, ldb, n * k, a, - dtype, lda, k * m, &beta, c, dtype, ldc, n * m, batch_size, ctype); - if (status != blasSuccess) { - throw std::runtime_error("Failed to call blasGemmStridedBatchedEx"); - } - } -} - -template -void baseline_matmul_nn(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, int) { - auto out_shape_dims4 = output_shapes[0].dims4(); - - // baseline inputs & outputs have no padding - int m = out_shape_dims4[2]; - int n = out_shape_dims4[3]; - int k = input_shapes[0].dims4()[3]; - int lda = k; - int ldb = n; - int ldc = n; - - int batch_size = out_shape_dims4[0] * out_shape_dims4[1]; - - auto memA = ark::to_gpu(inputs[0], input_shapes[0].nelems() * sizeof(T)); - auto memB = ark::to_gpu(inputs[1], input_shapes[1].nelems() * sizeof(T)); - auto memC = ark::to_gpu(outputs[0], output_shapes[0].nelems() * sizeof(T)); - - T *devA = static_cast(memA.get()); - T *devB = static_cast(memB.get()); - T *devC = static_cast(memC.get()); - - blas_matmul(m, n, k, devA, lda, devB, ldb, devC, ldc, - batch_size); - ark::sync_gpu(); - - // copy back to host - ark::from_gpu(memC, outputs[0]); -} - -template -void baseline_matmul_nt(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, int) { - auto out_shape_dims4 = output_shapes[0].dims4(); - - // baseline inputs & outputs have no padding - int m = out_shape_dims4[2]; - int n = out_shape_dims4[3]; - int k = input_shapes[0].dims4()[3]; - int lda = k; - int ldb = k; - int ldc = n; - - int batch_size = out_shape_dims4[0] * out_shape_dims4[1]; - - auto memA = ark::to_gpu(inputs[0], input_shapes[0].nelems() * sizeof(T)); - auto memB = ark::to_gpu(inputs[1], input_shapes[1].nelems() * sizeof(T)); - auto memC = ark::to_gpu(outputs[0], output_shapes[0].nelems() * sizeof(T)); - - T *devA = static_cast(memA.get()); - T *devB = static_cast(memB.get()); - T *devC = static_cast(memC.get()); - - blas_matmul(m, n, k, devA, lda, devB, ldb, devC, ldc, - batch_size); - ark::sync_gpu(); - - // copy back to host - ark::from_gpu(memC, outputs[0]); -} - -template -void baseline_matmul_tn(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, int) { - auto out_shape_dims4 = output_shapes[0].dims4(); - - // baseline inputs & outputs have no padding - int m = out_shape_dims4[2]; - int n = out_shape_dims4[3]; - int k = input_shapes[0].dims4()[2]; - int lda = m; - int ldb = n; - int ldc = n; - - int batch_size = out_shape_dims4[0] * out_shape_dims4[1]; - - auto memA = ark::to_gpu(inputs[0], input_shapes[0].nelems() * sizeof(T)); - auto memB = ark::to_gpu(inputs[1], input_shapes[1].nelems() * sizeof(T)); - auto memC = ark::to_gpu(outputs[0], output_shapes[0].nelems() * sizeof(T)); - - T *devA = static_cast(memA.get()); - T *devB = static_cast(memB.get()); - T *devC = static_cast(memC.get()); - - blas_matmul(m, n, k, devA, lda, devB, ldb, devC, ldc, - batch_size); - ark::sync_gpu(); - - // copy back to host - ark::from_gpu(memC, outputs[0]); -} - -template -void baseline_matmul_tt(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, int) { - auto out_shape_dims4 = output_shapes[0].dims4(); - - // baseline inputs & outputs have no padding - int m = out_shape_dims4[2]; - int n = out_shape_dims4[3]; - int k = input_shapes[0].dims4()[2]; - int lda = m; - int ldb = k; - int ldc = n; - - int batch_size = out_shape_dims4[0] * out_shape_dims4[1]; - - auto memA = ark::to_gpu(inputs[0], input_shapes[0].nelems() * sizeof(T)); - auto memB = ark::to_gpu(inputs[1], input_shapes[1].nelems() * sizeof(T)); - auto memC = ark::to_gpu(outputs[0], output_shapes[0].nelems() * sizeof(T)); - - T *devA = static_cast(memA.get()); - T *devB = static_cast(memB.get()); - T *devC = static_cast(memC.get()); - - blas_matmul(m, n, k, devA, lda, devB, ldb, devC, ldc, - batch_size); - ark::sync_gpu(); - - // copy back to host - ark::from_gpu(memC, outputs[0]); -} - -ark::unittest::State test_matmul_model() { - // Hidden dimension of the dense layer. - unsigned int units = 1024; - // Input dimension of the dense layer. - unsigned int in_dim = 1024; - // Extra dimension of the input. CHANNEL=1 for 2D inputs. - unsigned int channel = 128; - // Batch size of the input. - unsigned int batch_size = 1; - - ark::Model m; - ark::Tensor input = m.tensor({batch_size, channel, in_dim}, ark::FP16); - ark::Tensor weight = m.tensor({in_dim, units}, ark::FP16); - m.matmul(input, weight); - - UNITTEST_TRUE(m.verify()); - auto compressed = m.compress(); - UNITTEST_TRUE(compressed.verify()); - - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_matmul_fp16() { - { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(128, 64), ark::FP16); - ark::Tensor b = m.tensor(ark::Dims(64, 256), ark::FP16); - ark::Tensor c = m.matmul(a, b); - - auto result = ark::op_test("matmul_fp16", m, {a, b}, {c}, - baseline_matmul_nn); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 64)); - } - { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(4096, 2048), ark::FP16); - ark::Tensor b = m.tensor(ark::Dims(2048, 16384), ark::FP16); - ark::Tensor c = m.matmul(a, b); - - std::vector p_ones_a(a.shape().nelems(), - ark::half_t(0.1f)); - std::vector p_ones_b(b.shape().nelems(), - ark::half_t(0.1f)); - - auto result = ark::op_test("matmul_fp16", m, {a, b}, {c}, - baseline_matmul_nn, - {p_ones_a.data(), p_ones_b.data()}); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 2048)); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_matmul_fp32() { - { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(128, 64), ark::FP32); - ark::Tensor b = m.tensor(ark::Dims(64, 256), ark::FP32); - ark::Tensor c = m.matmul(a, b); - - auto result = ark::op_test("matmul_fp32", m, {a, b}, {c}, - baseline_matmul_nn); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 64)); - } - { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(4096, 8192), ark::FP32); - ark::Tensor b = m.tensor(ark::Dims(8192, 16384), ark::FP32); - ark::Tensor c = m.matmul(a, b); - - std::vector p_ones_a(a.shape().nelems(), float(0.1f)); - std::vector p_ones_b(b.shape().nelems(), float(0.1f)); - - auto result = ark::op_test("matmul_fp32", m, {a, b}, {c}, - baseline_matmul_nn, - {p_ones_a.data(), p_ones_b.data()}); - UNITTEST_LOG(result); - // TODO: #199 -#if defined(ARK_CUDA) - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 8192)); -#endif // defined(ARK_CUDA) - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_matmul_bf16() { - { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(128, 64), ark::BF16); - ark::Tensor b = m.tensor(ark::Dims(64, 256), ark::BF16); - ark::Tensor c = m.matmul(a, b); - - auto result = ark::op_test("matmul_bf16", m, {a, b}, {c}, - baseline_matmul_nn); - UNITTEST_LOG(result); - UNITTEST_TRUE( - result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 64)); - } - { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(4096, 256), ark::BF16); - ark::Tensor b = m.tensor(ark::Dims(256, 16384), ark::BF16); - ark::Tensor c = m.matmul(a, b); - - auto result = ark::op_test("matmul_bf16", m, {a, b}, {c}, - baseline_matmul_nn); - UNITTEST_LOG(result); - UNITTEST_TRUE( - result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 256)); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_matmul_fp16_nt() { - { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(128, 64), ark::FP16); - ark::Tensor b = m.tensor(ark::Dims(256, 64), ark::FP16); - ark::Tensor c = m.matmul(a, b, ark::NullTensor, false, true); - - auto result = ark::op_test("matmul_fp16_nt", m, {a, b}, {c}, - baseline_matmul_nt); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 64)); - } - { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(4096, 2048), ark::FP16); - ark::Tensor b = m.tensor(ark::Dims(16384, 2048), ark::FP16); - ark::Tensor c = m.matmul(a, b, ark::NullTensor, false, true); - - auto result = ark::op_test("matmul_fp16_nt", m, {a, b}, {c}, - baseline_matmul_nt); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 2048)); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_matmul_fp16_tn() { - { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(64, 128), ark::FP16); - ark::Tensor b = m.tensor(ark::Dims(64, 256), ark::FP16); - ark::Tensor c = m.matmul(a, b, ark::NullTensor, true, false); - - auto result = ark::op_test("matmul_fp16_tn", m, {a, b}, {c}, - baseline_matmul_tn); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 64)); - } - { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(2048, 4096), ark::FP16); - ark::Tensor b = m.tensor(ark::Dims(2048, 16384), ark::FP16); - ark::Tensor c = m.matmul(a, b, ark::NullTensor, true, false); - - auto result = ark::op_test("matmul_fp16_tn", m, {a, b}, {c}, - baseline_matmul_tn); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 2048)); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_matmul_fp16_tt() { - { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(64, 128), ark::FP16); - ark::Tensor b = m.tensor(ark::Dims(256, 64), ark::FP16); - ark::Tensor c = m.matmul(a, b, ark::NullTensor, true, true); - - auto result = ark::op_test("matmul_fp16_tt", m, {a, b}, {c}, - baseline_matmul_tt); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 64)); - } - { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(2048, 4096), ark::FP16); - ark::Tensor b = m.tensor(ark::Dims(16384, 2048), ark::FP16); - ark::Tensor c = m.matmul(a, b, ark::NullTensor, true, true); - - auto result = ark::op_test("matmul_fp16_tt", m, {a, b}, {c}, - baseline_matmul_tt); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 2048)); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_matmul_fp16_batched() { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(3, 7, 128, 128), ark::FP16); - ark::Tensor b = m.tensor(ark::Dims(3, 7, 128, 256), ark::FP16); - ark::Tensor c = m.matmul(a, b); - - auto result = ark::op_test("matmul_fp16_batched", m, {a, b}, {c}, - baseline_matmul_nn); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 128)); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_matmul_fp16_batched_padded() { - ark::Model m; - ark::Tensor a = - m.tensor({3, 7, 2, 9}, ark::FP16, {3, 7, 128, 64}, {}, {3, 7, 128, 64}); - ark::Tensor b = - m.tensor({3, 7, 9, 2}, ark::FP16, {3, 7, 64, 256}, {}, {3, 7, 64, 256}); - ark::Tensor c = m.tensor({3, 7, 2, 2}, ark::FP16, {3, 7, 128, 256}, {}, - {3, 7, 128, 256}); - m.matmul(a, b, c); - - auto result = ark::op_test("matmul_fp16_batched_padded", m, {a, b}, {c}, - baseline_matmul_nn); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 9)); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_matmul_fp16_offset() { - ark::Model m; - ark::Tensor a = - m.tensor({1, 128, 64}, ark::FP16, {1, 128, 256}, {0, 0, 64}); - ark::Tensor b = m.tensor({1, 64, 128}, ark::FP16, {1, 128, 256}, {0, 64, 0}, - {1, 64, 256}); - ark::Tensor c = m.tensor({1, 128, 128}, ark::FP16, {2, 256, 384}, - {1, 64, 128}, {1, 128, 256}); - m.matmul(a, b, c); - - auto result = ark::op_test("matmul_fp16_offset", m, {a, b}, {c}, - baseline_matmul_nn); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1f, 64)); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_matmul_invalid() { - ark::Model m; - ark::Tensor a = m.tensor(ark::Dims(128, 64), ark::FP16); - ark::Tensor b = m.tensor(ark::Dims(128, 256), ark::FP16); - UNITTEST_THROW(m.matmul(a, b), ark::ModelError); - - ark::Tensor c = m.tensor(ark::Dims(3, 3, 128, 128), ark::FP16); - ark::Tensor d = m.tensor(ark::Dims(2, 3, 128, 128), ark::FP16); - UNITTEST_THROW(m.matmul(c, d), ark::ModelError); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_matmul_host_ops() { - // Host-only: exercise default_config, impl_name, and impl_args for - // matmul variants without a GPU. - { - // Matmul + select_tile_config (128x256 is tile-aligned) - ark::Model m; - ark::Tensor a = m.tensor({128, 64}, ark::FP16); - ark::Tensor b = m.tensor({64, 256}, ark::FP16); - ark::Tensor c = m.matmul(a, b); - - auto nodes = m.nodes(); - for (auto &node : nodes) { - auto &op = node->op; - if (op->is_virtual()) continue; - auto cfg = op->default_config(ark::ARCH_CUDA_80); - UNITTEST_FALSE(cfg.empty()); - auto name = op->impl_name(cfg); - UNITTEST_FALSE(name.empty()); - (void)op->impl_args(cfg); - } - } - { - // MatmulGelu - ark::Model m; - ark::Tensor a = m.tensor({128, 64}, ark::FP16); - ark::Tensor b = m.tensor({64, 128}, ark::FP16); - ark::Tensor c = m.matmul_gelu(a, b); - - auto nodes = m.nodes(); - for (auto &node : nodes) { - auto &op = node->op; - if (op->is_virtual()) continue; - auto cfg = op->default_config(ark::ARCH_CUDA_80); - UNITTEST_FALSE(cfg.empty()); - auto name = op->impl_name(cfg); - UNITTEST_TRUE(name.find("matmul_gelu") != std::string::npos); - (void)op->impl_args(cfg); - } - } - { - // MatmulScale - ark::Model m; - ark::Tensor a = m.tensor({128, 64}, ark::FP16); - ark::Tensor b = m.tensor({64, 128}, ark::FP16); - ark::Tensor c = m.matmul_scale(a, b, 0.125f); - - auto nodes = m.nodes(); - for (auto &node : nodes) { - auto &op = node->op; - if (op->is_virtual()) continue; - auto cfg = op->default_config(ark::ARCH_CUDA_80); - UNITTEST_FALSE(cfg.empty()); - auto name = op->impl_name(cfg); - UNITTEST_TRUE(name.find("matmul_scale") != std::string::npos); - (void)op->impl_args(cfg); - } - } - { - // MatmulAdd - ark::Model m; - ark::Tensor a = m.tensor({128, 64}, ark::FP16); - ark::Tensor b = m.tensor({64, 128}, ark::FP16); - ark::Tensor res = m.tensor({128, 128}, ark::FP16); - ark::Tensor c = m.matmul_add(a, b, res); - - auto nodes = m.nodes(); - for (auto &node : nodes) { - auto &op = node->op; - if (op->is_virtual()) continue; - auto cfg = op->default_config(ark::ARCH_CUDA_80); - UNITTEST_FALSE(cfg.empty()); - auto name = op->impl_name(cfg); - UNITTEST_FALSE(name.empty()); - (void)op->impl_args(cfg); - } - } - { - // Mma + Store - ark::Model m; - ark::Tensor a = m.tensor({128, 64}, ark::FP16); - ark::Tensor b = m.tensor({64, 128}, ark::FP16); - ark::Tensor c = m.mma(a, b); - ark::Tensor d = m.store(ark::NullTensor, c); - - auto nodes = m.nodes(); - for (auto &node : nodes) { - auto &op = node->op; - if (op->is_virtual()) continue; - auto cfg = op->default_config(ark::ARCH_CUDA_80); - UNITTEST_FALSE(cfg.empty()); - auto name = op->impl_name(cfg); - UNITTEST_FALSE(name.empty()); - (void)op->impl_args(cfg); - } - } - { - // select_tile_config with small shapes (different tile candidates) - ark::Model m; - ark::Tensor a = m.tensor({64, 64}, ark::FP16); - ark::Tensor b = m.tensor({64, 64}, ark::FP16); - m.matmul(a, b); - - auto nodes = m.nodes(); - for (auto &node : nodes) { - auto &op = node->op; - if (op->is_virtual()) continue; - auto cfg = op->default_config(ark::ARCH_CUDA_80); - UNITTEST_FALSE(cfg.empty()); - } - } - { - // select_tile_config with larger shape to trigger bigger tile - ark::Model m; - ark::Tensor a = m.tensor({256, 128}, ark::FP16); - ark::Tensor b = m.tensor({128, 256}, ark::FP16); - m.matmul(a, b); - - auto nodes = m.nodes(); - for (auto &node : nodes) { - auto &op = node->op; - if (op->is_virtual()) continue; - auto cfg = op->default_config(ark::ARCH_CUDA_80); - UNITTEST_FALSE(cfg.empty()); - } - } - return ark::unittest::SUCCESS; -} - -int main() { - ark::init(); - UNITTEST(test_matmul_host_ops); - UNITTEST(test_matmul_model); - UNITTEST(test_matmul_fp16); - UNITTEST(test_matmul_fp32); - UNITTEST(test_matmul_bf16); - UNITTEST(test_matmul_fp16_nt); - UNITTEST(test_matmul_fp16_tn); - UNITTEST(test_matmul_fp16_tt); - UNITTEST(test_matmul_fp16_batched); - UNITTEST(test_matmul_fp16_batched_padded); - UNITTEST(test_matmul_fp16_offset); - UNITTEST(test_matmul_invalid); - return ark::unittest::SUCCESS; -} diff --git a/ark/ops/ops_reduce_test.cpp b/ark/ops/ops_reduce_test.cpp deleted file mode 100644 index 637c8dae..00000000 --- a/ark/ops/ops_reduce_test.cpp +++ /dev/null @@ -1,472 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -#include -#include -#include - -#include "ops_test_common.hpp" - -template -void baseline_reduce_sum_axis0(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, - int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - - ark::Dims osh = output_shapes[0]; - ark::Dims ish = input_shapes[0].dims4(); - - if (KeepDim) { - assert(osh[0] == 1); - } else { - osh.insert(0, 1); - } - osh = osh.dims4(); - - for (ark::DimType c = 0; c < ish[1]; ++c) { - for (ark::DimType h = 0; h < ish[2]; ++h) { - for (ark::DimType w = 0; w < ish[3]; ++w) { - float sum = 0; - for (ark::DimType n = 0; n < ish[0]; ++n) { - sum += float(input[n * ish[1] * ish[2] * ish[3] + - c * ish[2] * ish[3] + h * ish[3] + w]); - } - out[c * osh[2] * osh[3] + h * osh[3] + w] = T(sum); - } - } - } -} - -template -void baseline_reduce_sum_axis1(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, - int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - - ark::Dims osh = output_shapes[0]; - ark::Dims ish = input_shapes[0].dims4(); - - if (KeepDim) { - assert(osh[1] == 1); - } else { - osh.insert(1, 1); - } - osh = osh.dims4(); - - for (ark::DimType n = 0; n < ish[0]; ++n) { - for (ark::DimType h = 0; h < ish[2]; ++h) { - for (ark::DimType w = 0; w < ish[3]; ++w) { - float sum = 0; - for (ark::DimType c = 0; c < ish[1]; ++c) { - sum += float(input[n * ish[1] * ish[2] * ish[3] + - c * ish[2] * ish[3] + h * ish[3] + w]); - } - out[n * osh[1] * osh[2] * osh[3] + h * osh[3] + w] = T(sum); - } - } - } -} - -template -void baseline_reduce_sum_axis2(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, - int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - - ark::Dims osh = output_shapes[0]; - ark::Dims ish = input_shapes[0].dims4(); - - if (KeepDim) { - assert(osh[2] == 1); - } else { - osh.insert(2, 1); - } - osh = osh.dims4(); - - for (ark::DimType n = 0; n < ish[0]; ++n) { - for (ark::DimType c = 0; c < ish[1]; ++c) { - for (ark::DimType w = 0; w < ish[3]; ++w) { - float sum = 0; - for (ark::DimType h = 0; h < ish[2]; ++h) { - sum += float(input[n * ish[1] * ish[2] * ish[3] + - c * ish[2] * ish[3] + h * ish[3] + w]); - } - out[n * osh[1] * osh[2] * osh[3] + c * osh[2] * osh[3] + w] = - T(sum); - } - } - } -}; - -template -void baseline_reduce_sum_axis3(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, - int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - - ark::Dims osh = output_shapes[0]; - ark::Dims ish = input_shapes[0].dims4(); - - if (KeepDim) { - assert(osh[3] == 1); - } else { - osh.insert(3, 1); - } - osh = osh.dims4(); - - for (ark::DimType n = 0; n < ish[0]; ++n) { - for (ark::DimType c = 0; c < ish[1]; ++c) { - for (ark::DimType h = 0; h < ish[2]; ++h) { - float sum = 0; - for (ark::DimType w = 0; w < ish[3]; ++w) { - sum += float(input[n * ish[1] * ish[2] * ish[3] + - c * ish[2] * ish[3] + h * ish[3] + w]); - } - out[n * osh[1] * osh[2] * osh[3] + c * osh[2] * osh[3] + - h * osh[3]] = T(sum); - } - } - } -}; - -template -void baseline_reduce_max_axis3(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, - int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - - ark::Dims osh = output_shapes[0]; - ark::Dims ish = input_shapes[0].dims4(); - - if (KeepDim) { - assert(osh[3] == 1); - } else { - osh.insert(3, 1); - } - osh = osh.dims4(); - - for (ark::DimType n = 0; n < ish[0]; ++n) { - for (ark::DimType c = 0; c < ish[1]; ++c) { - for (ark::DimType h = 0; h < ish[2]; ++h) { - float max_val = std::numeric_limits::lowest(); - for (ark::DimType w = 0; w < ish[3]; ++w) { - float val = - float(input[n * ish[1] * ish[2] * ish[3] + - c * ish[2] * ish[3] + h * ish[3] + w]); - max_val = std::max(max_val, val); - } - out[n * osh[1] * osh[2] * osh[3] + c * osh[2] * osh[3] + - h * osh[3]] = T(max_val); - } - } - } -}; - -template -void baseline_reduce_mean_axis3(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, - int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - - ark::Dims osh = output_shapes[0]; - ark::Dims ish = input_shapes[0].dims4(); - - if (KeepDim) { - assert(osh[3] == 1); - } else { - osh.insert(3, 1); - } - osh = osh.dims4(); - - for (ark::DimType n = 0; n < ish[0]; ++n) { - for (ark::DimType c = 0; c < ish[1]; ++c) { - for (ark::DimType h = 0; h < ish[2]; ++h) { - float mean = 0; - for (ark::DimType w = 0; w < ish[3]; ++w) { - mean += float(input[n * ish[1] * ish[2] * ish[3] + - c * ish[2] * ish[3] + h * ish[3] + w]); - } - mean /= ish[3]; - out[n * osh[1] * osh[2] * osh[3] + c * osh[2] * osh[3] + - h * osh[3]] = T(mean); - } - } - } -}; - -ark::unittest::State test_reduce_sum_axis0() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 1024), ark::FP32); - ark::Tensor out = m.reduce_sum(t, /*axis=*/0); - - auto result = ark::op_test("reduce_sum_axis0", m, {t}, {out}, - baseline_reduce_sum_axis0); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1, t.shape()[0])); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_reduce_sum_axis1() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(1, 2, 4, 1024), ark::FP32); - ark::Tensor out = m.reduce_sum(t, /*axis=*/1); - - auto result = ark::op_test("reduce_sum_axis1", m, {t}, {out}, - baseline_reduce_sum_axis1); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1, t.shape()[1])); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_reduce_sum_axis2() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(1, 1, 7, 8192), ark::FP32); - ark::Tensor out = m.reduce_sum(t, /*axis=*/2); - - auto result = ark::op_test("reduce_sum_axis2", m, {t}, {out}, - baseline_reduce_sum_axis2); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1, t.shape()[2])); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_reduce_sum_axis3() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(1, 1, 2, 8192), ark::FP32); - ark::Tensor out = m.reduce_sum(t, /*axis=*/3); - - auto result = ark::op_test("reduce_sum_axis3", m, {t}, {out}, - baseline_reduce_sum_axis3); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1, t.shape()[3])); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_reduce_sum_axis3_padded() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(1, 1, 2, 8192), ark::FP32); - ark::Tensor out = - m.tensor(ark::Dims(1, 1, 2, 1), ark::FP32, ark::Dims(1, 1, 2, 32)); - out = m.reduce_sum(t, /*axis=*/3, true, out); - - auto result = ark::op_test("reduce_sum_axis3_padded", m, {t}, {out}, - baseline_reduce_sum_axis3); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1, t.shape()[3])); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_reduce_sum_fp16() { - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 1024), ark::FP16); - ark::Tensor out = m.reduce_sum(t, /*axis=*/0); - - auto result = ark::op_test("reduce_sum_fp16_axis0", m, {t}, {out}, - baseline_reduce_sum_axis0); - UNITTEST_LOG(result); - UNITTEST_TRUE( - result.max_diff[0] < - ark::reduction_abs_error_bound(0.1, t.shape()[0])); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 1024), ark::FP16); - ark::Tensor out = m.reduce_sum(t, /*axis=*/3); - - auto result = ark::op_test("reduce_sum_fp16_axis3", m, {t}, {out}, - baseline_reduce_sum_axis3); - UNITTEST_LOG(result); - UNITTEST_TRUE( - result.max_diff[0] < - ark::reduction_abs_error_bound(0.1, t.shape()[3])); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_reduce_sum_bf16() { - std::vector data_vec(7 * 2 * 4 * 256); - for (size_t i = 0; i < data_vec.size(); ++i) { - data_vec[i] = ark::bf16((i % 1000) * 1e-4f); - } - - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 256), ark::BF16); - ark::Tensor out = m.reduce_sum(t, /*axis=*/0); - - auto result = ark::op_test("reduce_sum_bf16_axis0", m, {t}, {out}, - baseline_reduce_sum_axis0, - {data_vec.data()}); - UNITTEST_LOG(result); - UNITTEST_TRUE( - result.max_diff[0] < - ark::reduction_abs_error_bound(0.1, t.shape()[0])); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 256), ark::BF16); - ark::Tensor out = m.reduce_sum(t, /*axis=*/3); - - auto result = ark::op_test("reduce_sum_bf16_axis3", m, {t}, {out}, - baseline_reduce_sum_axis3, - {data_vec.data()}); - UNITTEST_LOG(result); - UNITTEST_TRUE( - result.max_diff[0] < - ark::reduction_abs_error_bound(0.1, t.shape()[3])); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_reduce_sum_fp16_no_keepdims() { - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 1024), ark::FP16); - ark::Tensor out = m.reduce_sum(t, /*axis=*/0, false); - - UNITTEST_EQ(out.shape(), ark::Dims(2, 4, 1024)); - - auto result = - ark::op_test("reduce_sum_fp16_axis0", m, {t}, {out}, - baseline_reduce_sum_axis0); - UNITTEST_LOG(result); - UNITTEST_TRUE( - result.max_diff[0] < - ark::reduction_abs_error_bound(0.1, t.shape()[0])); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 1024), ark::FP16); - ark::Tensor out = m.reduce_sum(t, /*axis=*/3, false); - - UNITTEST_EQ(out.shape(), ark::Dims(7, 2, 4)); - - auto result = - ark::op_test("reduce_sum_fp16_axis3", m, {t}, {out}, - baseline_reduce_sum_axis3); - UNITTEST_LOG(result); - UNITTEST_TRUE( - result.max_diff[0] < - ark::reduction_abs_error_bound(0.1, t.shape()[3])); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_reduce_sum_invalid() { - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 1024), ark::BF16); - ark::Tensor out = m.tensor(ark::Dims(1, 2, 4, 1024), ark::FP32); - UNITTEST_THROW(m.reduce_sum(t, /*axis=*/0, true, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 1024), ark::BF16); - ark::Tensor out = m.tensor(ark::Dims(7, 2, 4, 1), ark::FP32); - UNITTEST_THROW(m.reduce_sum(t, /*axis=*/3, true, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 1024), ark::BF16); - ark::Tensor out = m.tensor(ark::Dims(1, 2, 4, 512), ark::BF16); - UNITTEST_THROW(m.reduce_sum(t, /*axis=*/0, true, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 1024), ark::BF16); - ark::Tensor out = m.tensor(ark::Dims(7, 1, 4, 1), ark::BF16); - UNITTEST_THROW(m.reduce_sum(t, /*axis=*/3, true, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 1024), ark::BF16); - ark::Tensor out = m.tensor(ark::Dims(3, 2, 4, 1024), ark::BF16); - UNITTEST_THROW(m.reduce_sum(t, /*axis=*/0, true, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 1024), ark::BF16); - ark::Tensor out = m.tensor(ark::Dims(7, 2, 4, 3), ark::BF16); - UNITTEST_THROW(m.reduce_sum(t, /*axis=*/3, true, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 1), ark::BF16); - UNITTEST_THROW(m.reduce_sum(t, /*axis=*/3, true, t), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_reduce_max_axis3() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(1, 1, 2, 8192), ark::FP32); - ark::Tensor out = m.reduce_max(t, /*axis=*/3); - - auto result = ark::op_test("reduce_max_axis3", m, {t}, {out}, - baseline_reduce_max_axis3); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_reduce_mean_axis3() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(1, 1, 2, 8192), ark::FP32); - ark::Tensor out = m.reduce_mean(t, /*axis=*/3); - - auto result = ark::op_test("reduce_mean_axis3", m, {t}, {out}, - baseline_reduce_mean_axis3); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < - ark::reduction_abs_error_bound(0.1, t.shape()[3]) / - t.shape()[3]); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_reduce_invalid() { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(7, 2, 4, 1024), ark::FP32); - UNITTEST_THROW(m.reduce_max(t, /*axis=*/-10), ark::ModelError); - return ark::unittest::SUCCESS; -} - -int main() { - ark::init(); - UNITTEST(test_reduce_sum_axis0); - UNITTEST(test_reduce_sum_axis1); - UNITTEST(test_reduce_sum_axis2); - UNITTEST(test_reduce_sum_axis3); - UNITTEST(test_reduce_sum_axis3_padded); - UNITTEST(test_reduce_sum_fp16); - UNITTEST(test_reduce_sum_bf16); - UNITTEST(test_reduce_sum_fp16_no_keepdims); - UNITTEST(test_reduce_sum_invalid); - UNITTEST(test_reduce_max_axis3); - UNITTEST(test_reduce_mean_axis3); - UNITTEST(test_reduce_invalid); - return ark::unittest::SUCCESS; -} diff --git a/ark/ops/ops_rope_test.cpp b/ark/ops/ops_rope_test.cpp deleted file mode 100644 index b0812fae..00000000 --- a/ark/ops/ops_rope_test.cpp +++ /dev/null @@ -1,103 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -#include "ark/model.hpp" -#include "ops_test_common.hpp" -#include "unittest/unittest_utils.h" - -template -void baseline_rope(std::vector &outputs, const std::vector &, - const std::vector &inputs, - const std::vector &input_shapes, int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - T *other = static_cast(inputs[1]); - - ark::Dims ish = input_shapes[0].dims4(); - - for (ark::DimType n = 0; n < ish[0]; ++n) { - for (ark::DimType c = 0; c < ish[1]; ++c) { - for (ark::DimType h = 0; h < ish[2]; ++h) { - for (ark::DimType w = 0; w < ish[3]; w += 2) { - int idx = n * ish[1] * ish[2] * ish[3] + - c * ish[2] * ish[3] + h * ish[3] + w; - T input0 = input[idx]; - T input1 = input[idx + 1]; - T other0 = other[idx]; - T other1 = other[idx + 1]; - out[idx] = input0 * other0 - input1 * other1; - out[idx + 1] = input0 * other1 + input1 * other0; - } - } - } - } -} - -ark::unittest::State test_rope_fp32() { - ark::Model model; - ark::Tensor input = model.tensor(ark::Dims(1, 32, 32, 256), ark::FP32); - ark::Tensor other = model.tensor(ark::Dims(1, 32, 32, 256), ark::FP32); - ark::Tensor out = model.rope(input, other); - auto result = ark::op_test("rope", model, {input, other}, {out}, - baseline_rope); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-6f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_rope_fp16() { - ark::Model model; - ark::Tensor input = model.tensor(ark::Dims(1, 32, 32, 256), ark::FP16); - ark::Tensor other = model.tensor(ark::Dims(1, 32, 32, 256), ark::FP16); - ark::Tensor out = model.rope(input, other); - auto result = ark::op_test("rope", model, {input, other}, {out}, - baseline_rope); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-3f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_rope_bf16() { - ark::Model model; - ark::Tensor input = model.tensor(ark::Dims(1, 32, 32, 256), ark::BF16); - ark::Tensor other = model.tensor(ark::Dims(1, 32, 32, 256), ark::BF16); - ark::Tensor out = model.rope(input, other); - auto result = ark::op_test("rope", model, {input, other}, {out}, - baseline_rope); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-3f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_rope_invalid() { - { - ark::Model model; - ark::Tensor input = model.tensor(ark::Dims(1, 32, 32, 256), ark::BF16); - ark::Tensor other = model.tensor(ark::Dims(1, 32, 32, 256), ark::FP32); - UNITTEST_THROW(model.rope(input, other), ark::ModelError); - } - { - ark::Model model; - ark::Tensor input = model.tensor(ark::Dims(1, 32, 32, 256), ark::FP32); - ark::Tensor other = model.tensor(ark::Dims(1, 32, 32, 256), ark::FP32); - ark::Tensor output = model.tensor(ark::Dims(1, 32, 32, 256), ark::FP16); - UNITTEST_THROW(model.rope(input, other, output), ark::ModelError); - } - { - ark::Model model; - ark::Tensor input = model.tensor(ark::Dims(1, 32, 32, 256), ark::FP32); - ark::Tensor other = model.tensor(ark::Dims(1, 32, 32, 256), ark::FP32); - ark::Tensor output = model.tensor(ark::Dims(1, 32, 32, 32), ark::FP32); - UNITTEST_THROW(model.rope(input, other, output), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -int main() { - ark::init(); - UNITTEST(test_rope_fp32); - UNITTEST(test_rope_fp16); - UNITTEST(test_rope_bf16); - UNITTEST(test_rope_invalid); - return ark::unittest::SUCCESS; -} diff --git a/ark/ops/ops_scalar_test.cpp b/ark/ops/ops_scalar_test.cpp deleted file mode 100644 index 47a5b40b..00000000 --- a/ark/ops/ops_scalar_test.cpp +++ /dev/null @@ -1,345 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -#include "ark/executor.hpp" -#include "ark/model.hpp" -#include "ops_test_common.hpp" -#include "unittest/unittest_utils.h" - -#define FACTOR 0.7 - -template -void baseline_scalar_add(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &, int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0]; - for (ark::DimType i = 0; i < osh.nelems(); ++i) { - out[i] = input[i] + T(FACTOR); - } -}; - -template -void baseline_scalar_sub(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &, int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0]; - for (ark::DimType i = 0; i < osh.nelems(); ++i) { - out[i] = input[i] - T(FACTOR); - } -}; - -template -void baseline_scalar_mul(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &, int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0]; - for (ark::DimType i = 0; i < osh.nelems(); ++i) { - out[i] = input[i] * T(FACTOR); - } -}; - -template -void baseline_scalar_div(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &, int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0]; - for (ark::DimType i = 0; i < osh.nelems(); ++i) { - out[i] = input[i] / T(FACTOR); - } -}; - -ark::unittest::State test_scalar_assign_fp16() { - { - ark::Model m; - ark::Tensor t = m.constant(7, ark::Dims(4, 2, 50), ark::FP16); - - ark::DefaultExecutor exe(m); - - exe.launch(); - exe.run(1); - exe.stop(); - - std::vector data(4 * 2 * 50); - exe.tensor_read(t, data); - for (auto v : data) { - UNITTEST_EQ(v, ark::half_t(7)); - } - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 50), ark::FP16); - ark::Tensor out = m.copy(7, t); - - ark::DefaultExecutor exe(m); - - std::vector data(4 * 2 * 50, 3); - exe.tensor_write(t, data); - - exe.launch(); - exe.run(1); - exe.stop(); - - data.clear(); - data.resize(4 * 2 * 50); - exe.tensor_read(t, data); - for (auto v : data) { - UNITTEST_EQ(v, ark::half_t(7)); - } - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_scalar_assign_fp32() { - { - ark::Model m; - ark::Tensor out = m.copy(7); - - ark::DefaultExecutor exe(m); - - exe.launch(); - exe.run(1); - exe.stop(); - - std::vector data(1); - exe.tensor_read(out, data); - for (auto v : data) { - UNITTEST_EQ(v, 7); - } - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_scalar_add_fp16() { - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1), ark::FP16); - ark::Tensor out = m.add(t, FACTOR); - - auto result = ark::op_test("scalar_add_fp16_small", m, {t}, {out}, - baseline_scalar_add); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::FP16); - ark::Tensor out = m.add(t, FACTOR); - - auto result = ark::op_test("scalar_add_fp16", m, {t}, {out}, - baseline_scalar_add); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_scalar_sub_fp16() { - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1), ark::FP16); - ark::Tensor out = m.sub(t, FACTOR); - - auto result = ark::op_test("scalar_sub_fp16_small", m, {t}, {out}, - baseline_scalar_sub); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::FP16); - ark::Tensor out = m.sub(t, FACTOR); - - auto result = ark::op_test("scalar_sub_fp16", m, {t}, {out}, - baseline_scalar_sub); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_scalar_mul_fp32() { - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1), ark::FP32); - ark::Tensor out = m.mul(t, FACTOR); - - auto result = ark::op_test("scalar_mul_fp32_small", m, {t}, {out}, - baseline_scalar_mul); - UNITTEST_LOG(result); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::FP32); - ark::Tensor out = m.mul(t, FACTOR); - - auto result = ark::op_test("scalar_mul_fp32", m, {t}, {out}, - baseline_scalar_mul); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_scalar_mul_fp16() { - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1), ark::FP16); - ark::Tensor out = m.mul(t, FACTOR); - - auto result = ark::op_test("scalar_mul_fp16_small", m, {t}, {out}, - baseline_scalar_mul); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::FP16); - ark::Tensor out = m.mul(t, FACTOR); - - auto result = ark::op_test("scalar_mul_fp16", m, {t}, {out}, - baseline_scalar_mul); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_scalar_mul_bf16() { - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1), ark::BF16); - ark::Tensor out = m.mul(t, FACTOR); - - auto result = ark::op_test("scalar_mul_bf16_small", m, {t}, {out}, - baseline_scalar_mul); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::BF16); - ark::Tensor out = m.mul(t, FACTOR); - - auto result = ark::op_test("scalar_mul_bf16", m, {t}, {out}, - baseline_scalar_mul); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_scalar_mul_invalid() { - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::BF16); - ark::Tensor out = m.tensor(ark::Dims(4, 2, 1024), ark::FP32); - UNITTEST_THROW(m.mul(t, 3, out), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::BF16); - ark::Tensor out = m.tensor(ark::Dims(4, 4, 1024), ark::BF16); - UNITTEST_THROW(m.mul(t, 3, out), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_scalar_mul_fp16_offset() { - { - ark::Model m; - ark::Tensor buf = m.tensor({1024}, ark::FP16); - ark::Tensor tns = m.refer(buf, {2}, {1024}, {6}); - ark::Tensor doubled = m.mul(tns, 2, tns); - ark::Tensor out = m.identity(buf, {doubled}); - - std::vector data(1024, ark::half_t(2)); - auto result = ark::op_test( - "scalar_mul_fp16_offset", m, {buf}, {out}, - [](std::vector &outputs, const std::vector &, - const std::vector &, const std::vector &, - int) { - ark::half_t *out = static_cast(outputs[0]); - for (size_t i = 0; i < 1024; ++i) { - if (i == 6 || i == 7) { - out[i] = 4; - } else { - out[i] = 2; - } - } - }, - {data.data()}); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - } - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_scalar_mul_perf() { - ark::DimType nelem = 8 * 1024 * 1024; - - ark::Model m; - ark::Tensor t = m.tensor({nelem}, ark::FP32); - ark::Tensor out = m.mul(t, 0.7); - - auto result = ark::op_test("scalar_mul_perf", m, {t}, {out}, - baseline_scalar_mul); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - - float gbps = nelem * sizeof(float) / result.msec_per_iter * 1e-6; - UNITTEST_LOG(gbps, " GB/s"); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_scalar_div_fp16() { - float rel_err_bound = ark::division_rel_error_bound(FACTOR); - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1), ark::FP16); - ark::Tensor out = m.div(t, FACTOR); - - auto result = ark::op_test("scalar_div_fp16_small", m, {t}, {out}, - baseline_scalar_div); - UNITTEST_LOG(result); - UNITTEST_LT(result.max_err_rate[0], rel_err_bound); - } - { - ark::Model m; - ark::Tensor t = m.tensor(ark::Dims(4, 2, 1024), ark::FP16); - ark::Tensor out = m.div(t, FACTOR); - - auto result = ark::op_test("scalar_div_fp16", m, {t}, {out}, - baseline_scalar_div); - UNITTEST_LOG(result); - UNITTEST_LT(result.max_err_rate[0], rel_err_bound); - } - return ark::unittest::SUCCESS; -} - -int main() { - ark::init(); - UNITTEST(test_scalar_assign_fp16); - UNITTEST(test_scalar_assign_fp32); - UNITTEST(test_scalar_add_fp16); - UNITTEST(test_scalar_sub_fp16); - UNITTEST(test_scalar_mul_fp32); - UNITTEST(test_scalar_mul_fp16); - UNITTEST(test_scalar_mul_bf16); - UNITTEST(test_scalar_mul_invalid); - UNITTEST(test_scalar_mul_fp16_offset); - UNITTEST(test_scalar_mul_perf); - UNITTEST(test_scalar_div_fp16); - return ark::unittest::SUCCESS; -} diff --git a/ark/ops/ops_softmax_test.cpp b/ark/ops/ops_softmax_test.cpp deleted file mode 100644 index 0562e3ee..00000000 --- a/ark/ops/ops_softmax_test.cpp +++ /dev/null @@ -1,184 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -#include -#include - -#include "ops_test_common.hpp" - -// Baseline: row-wise softmax. -// For each row (last dim = W): -// max_val = max(row) -// exp_sum = sum(exp(x - max_val)) -// out[j] = exp(x[j] - max_val) / exp_sum -template -void baseline_softmax(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &, int) { - T *out = static_cast(outputs[0]); - T *input = static_cast(inputs[0]); - - ark::Dims osh = output_shapes[0]; - ark::DimType total = osh.nelems(); - ark::DimType W = osh[-1]; - ark::DimType num_rows = total / W; - - for (ark::DimType row = 0; row < num_rows; ++row) { - T *row_in = input + row * W; - T *row_out = out + row * W; - - // pass 1: max - float max_val = static_cast(row_in[0]); - for (ark::DimType j = 1; j < W; ++j) { - float v = static_cast(row_in[j]); - if (v > max_val) max_val = v; - } - - // pass 2: exp and sum - float exp_sum = 0; - for (ark::DimType j = 0; j < W; ++j) { - float e = std::exp(static_cast(row_in[j]) - max_val); - exp_sum += e; - } - - // pass 3: normalize - for (ark::DimType j = 0; j < W; ++j) { - float e = std::exp(static_cast(row_in[j]) - max_val); - row_out[j] = T(e / exp_sum); - } - } -} - -ark::unittest::State test_softmax_fp32() { - ark::Model m; - ark::Tensor input = m.tensor({4, 1024}, ark::FP32); - ark::Tensor out = m.softmax(input); - - auto result = ark::op_test("softmax_fp32", m, {input}, {out}, - baseline_softmax); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-5f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_softmax_fp16() { - ark::Model m; - ark::Tensor input = m.tensor({2, 512}, ark::FP16); - ark::Tensor out = m.softmax(input); - - auto result = ark::op_test("softmax_fp16", m, {input}, {out}, - baseline_softmax); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 5e-3f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_softmax_bf16() { - ark::Model m; - ark::Tensor input = m.tensor({2, 512}, ark::BF16); - ark::Tensor out = m.softmax(input); - - auto result = ark::op_test("softmax_bf16", m, {input}, {out}, - baseline_softmax); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 5e-2f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_softmax_batch() { - // Higher-dimensional: [B, H, S, S] — attention pattern - ark::Model m; - ark::Tensor input = m.tensor({2, 12, 64, 64}, ark::FP32); - ark::Tensor out = m.softmax(input); - - auto result = ark::op_test("softmax_batch", m, {input}, {out}, - baseline_softmax); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-5f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_softmax_small_row() { - // Small last dim — tests edge case where W < warp size - ark::Model m; - ark::Tensor input = m.tensor({8, 16}, ark::FP32); - ark::Tensor out = m.softmax(input); - - auto result = ark::op_test("softmax_small_row", m, {input}, {out}, - baseline_softmax); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-5f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_softmax_non_pow2() { - // Non-power-of-2 W — exercises boundary guards (idx_w + j < W) - ark::Model m; - ark::Tensor input = m.tensor({4, 127}, ark::FP32); - ark::Tensor out = m.softmax(input); - - auto result = ark::op_test("softmax_non_pow2", m, {input}, {out}, - baseline_softmax); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-5f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_softmax_w1() { - // W=1 boundary: softmax output must be exactly 1.0 - ark::Model m; - ark::Tensor input = m.tensor({4, 1}, ark::FP32); - ark::Tensor out = m.softmax(input); - - auto result = - ark::op_test("softmax_w1", m, {input}, {out}, baseline_softmax); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-5f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_softmax_large_w() { - // Large inner dim — exercises numerical stability for W=4096 - ark::Model m; - ark::Tensor input = m.tensor({1, 1, 1, 4096}, ark::FP32); - ark::Tensor out = m.softmax(input); - - auto result = ark::op_test("softmax_large_w", m, {input}, {out}, - baseline_softmax); - UNITTEST_LOG(result); - UNITTEST_TRUE(result.max_diff[0] < 1e-5f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_softmax_invalid() { - // Output shape mismatch - { - ark::Model m; - ark::Tensor input = m.tensor({4, 1024}, ark::FP32); - ark::Tensor bad_out = m.tensor({4, 512}, ark::FP32); // wrong W - UNITTEST_THROW(m.softmax(input, bad_out), ark::ModelError); - } - // Data type mismatch - { - ark::Model m; - ark::Tensor input = m.tensor({4, 1024}, ark::FP32); - ark::Tensor bad_out = m.tensor({4, 1024}, ark::FP16); - UNITTEST_THROW(m.softmax(input, bad_out), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -int main() { - ark::init(); - UNITTEST(test_softmax_fp32); - UNITTEST(test_softmax_fp16); - UNITTEST(test_softmax_bf16); - UNITTEST(test_softmax_batch); - UNITTEST(test_softmax_small_row); - UNITTEST(test_softmax_non_pow2); - UNITTEST(test_softmax_w1); - UNITTEST(test_softmax_large_w); - UNITTEST(test_softmax_invalid); - return ark::unittest::SUCCESS; -} diff --git a/ark/ops/ops_transpose_test.cpp b/ark/ops/ops_transpose_test.cpp deleted file mode 100644 index 139e1ee6..00000000 --- a/ark/ops/ops_transpose_test.cpp +++ /dev/null @@ -1,281 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT license. - -#include "ark/model.hpp" -#include "ark/planner.hpp" -#include "model/model_json.hpp" -#include "ops_test_common.hpp" -#include "unittest/unittest_utils.h" - -#define SYNC_TEST 0 - -template -void baseline_transpose_0132(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, int) { - T *out = static_cast(outputs[0]); - T *in = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0].dims4(); - ark::Dims ish = input_shapes[0].dims4(); - for (ark::DimType n = 0; n < ish[0]; ++n) { - for (ark::DimType c = 0; c < ish[1]; ++c) { - for (ark::DimType h = 0; h < ish[2]; ++h) { - for (ark::DimType w = 0; w < ish[3]; ++w) { - // out[n][c][w][h] = in[n][c][h][w] - out[h + w * osh[3] + c * osh[2] * osh[3] + - n * osh[1] * osh[2] * osh[3]] = - in[w + h * ish[3] + c * ish[3] * ish[2] + - n * ish[3] * ish[2] * ish[1]]; - } - } - } - } -}; - -template -void baseline_transpose_0231(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, int) { - T *out = static_cast(outputs[0]); - T *in = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0].dims4(); - ark::Dims ish = input_shapes[0].dims4(); - for (ark::DimType n = 0; n < ish[0]; ++n) { - for (ark::DimType c = 0; c < ish[1]; ++c) { - for (ark::DimType h = 0; h < ish[2]; ++h) { - for (ark::DimType w = 0; w < ish[3]; ++w) { - // out[n][h][w][c] = in[n][c][h][w] - out[c + w * osh[3] + h * osh[2] * osh[3] + - n * osh[1] * osh[2] * osh[3]] = - in[w + h * ish[3] + c * ish[3] * ish[2] + - n * ish[3] * ish[2] * ish[1]]; - } - } - } - } -}; - -template -void baseline_transpose_0213(std::vector &outputs, - const std::vector &output_shapes, - const std::vector &inputs, - const std::vector &input_shapes, int) { - T *out = static_cast(outputs[0]); - T *in = static_cast(inputs[0]); - ark::Dims osh = output_shapes[0].dims4(); - ark::Dims ish = input_shapes[0].dims4(); - for (ark::DimType n = 0; n < ish[0]; ++n) { - for (ark::DimType c = 0; c < ish[1]; ++c) { - for (ark::DimType h = 0; h < ish[2]; ++h) { - for (ark::DimType w = 0; w < ish[3]; ++w) { - // out[n][h][c][w] = in[n][c][h][w] - out[w + c * osh[3] + h * osh[2] * osh[3] + - n * osh[1] * osh[2] * osh[3]] = - in[w + h * ish[3] + c * ish[3] * ish[2] + - n * ish[3] * ish[2] * ish[1]]; - } - } - } - } -}; - -template -void baseline_transpose_sync_test(std::vector &outputs, - const std::vector &, - const std::vector &inputs, - const std::vector &input_shapes, - int) { - T *out = static_cast(outputs[0]); - T *in = static_cast(inputs[0]); - ::memcpy(out, in, sizeof(T) * input_shapes[0].nelems()); -}; - -ark::unittest::State test_transpose_0132_fp32() { - ark::Model m; - ark::Tensor t = m.tensor({5, 3, 32, 128}, ark::FP32); - ark::Tensor out = m.transpose(t, {0, 1, 3, 2}); - - auto result = ark::op_test("transpose_0132_fp32", m, {t}, {out}, - baseline_transpose_0132); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_transpose_0132_fp16() { - ark::Model m; - ark::Tensor t = m.tensor({5, 3, 32, 128}, ark::FP16); - ark::Tensor out = m.transpose(t, {0, 1, 3, 2}); - - auto result = ark::op_test("transpose_0132_fp16", m, {t}, {out}, - baseline_transpose_0132); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_transpose_0132_bf16() { - ark::Model m; - ark::Tensor t = m.tensor({5, 3, 32, 128}, ark::BF16); - ark::Tensor out = m.transpose(t, {0, 1, 3, 2}); - - auto result = ark::op_test("transpose_0132_bf16", m, {t}, {out}, - baseline_transpose_0132); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_transpose_0231_fp32() { - ark::Model m; - ark::Tensor t = m.tensor({5, 3, 32, 128}, ark::FP32); - ark::Tensor out = m.transpose(t, {0, 2, 3, 1}); - - auto result = ark::op_test("transpose_0231_fp32", m, {t}, {out}, - baseline_transpose_0231); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_transpose_0231_fp16() { - ark::Model m; - ark::Tensor t = m.tensor({5, 3, 32, 128}, ark::FP16); - ark::Tensor out = m.transpose(t, {0, 2, 3, 1}); - - auto result = ark::op_test("transpose_0231_fp16", m, {t}, {out}, - baseline_transpose_0231); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_transpose_0231_bf16() { - ark::Model m; - ark::Tensor t = m.tensor({5, 3, 32, 128}, ark::BF16); - ark::Tensor out = m.transpose(t, {0, 2, 3, 1}); - - auto result = ark::op_test("transpose_0231_bf16", m, {t}, {out}, - baseline_transpose_0231); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_transpose_0213_fp32() { - ark::Model m; - ark::Tensor t = m.tensor({5, 3, 32, 128}, ark::FP32); - ark::Tensor out = m.transpose(t, {0, 2, 1, 3}); - - auto result = ark::op_test("transpose_0213_fp32", m, {t}, {out}, - baseline_transpose_0213); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_transpose_0213_fp16() { - ark::Model m; - ark::PlannerContext ctx(m); - ctx.warp_range(0, 4); - ctx.sram_range(0, 0); - ctx.sync(false); - ctx.config(ark::Json({{"NumWarps", 4}, {"SramBytes", 0}, {"Tile", {8, 64}}}) - .dump()); - - ark::Tensor t = m.tensor({5, 256, 32, 128}, ark::FP16); - ark::Tensor out = m.transpose(t, {0, 2, 1, 3}); - - auto result = ark::op_test("transpose_0213_fp16", m, {t}, {out}, - baseline_transpose_0213); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_transpose_0213_bf16() { - ark::Model m; - ark::Tensor t = m.tensor({5, 3, 32, 128}, ark::BF16); - ark::Tensor out = m.transpose(t, {0, 2, 1, 3}); - - auto result = ark::op_test("transpose_0213_bf16", m, {t}, {out}, - baseline_transpose_0213); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_transpose_sync_test() { - ark::Model m; - ark::PlannerContext shared_ctx(m); - shared_ctx.warp_range(0, 4); - shared_ctx.sram_range(0, 0); - shared_ctx.sync(false); - - ark::Tensor in, t, out; - in = m.tensor({1, 16, 2, 64}, ark::FP16); - { - ark::PlannerContext ctx(m); - ctx.config( - ark::Json({{"NumWarps", 4}, {"SramBytes", 0}, {"Tile", {8, 64}}}) - .dump()); - t = m.transpose(in, {0, 2, 1, 3}); - } - { - ark::PlannerContext ctx(m); - ctx.config( - ark::Json({{"NumWarps", 4}, {"SramBytes", 0}, {"Tile", {8, 1, 64}}}) - .dump()); - out = m.transpose(t, {0, 2, 1, 3}); - } - - auto result = ark::op_test("transpose_sync_test", m, {in}, {out}, - baseline_transpose_sync_test); - UNITTEST_LOG(result); - UNITTEST_EQ(result.max_diff[0], 0.0f); - return ark::unittest::SUCCESS; -} - -ark::unittest::State test_transpose_invalid() { - { - ark::Model m; - ark::Tensor t = m.tensor({5}, ark::FP32); - UNITTEST_THROW(m.transpose(t, {0, 2, 3, 1}), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor({5, 128}, ark::FP32); - UNITTEST_THROW(m.transpose(t, {0, 2, 3, 1}), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor({5, 128}, ark::FP32); - UNITTEST_THROW(m.transpose(t, {0, 2}), ark::ModelError); - } - { - ark::Model m; - ark::Tensor t = m.tensor({5, 128}, ark::FP32); - UNITTEST_THROW(m.transpose(t, {1, 1}), ark::ModelError); - } - return ark::unittest::SUCCESS; -} - -int main() { - ark::init(); - UNITTEST(test_transpose_0132_fp32); - UNITTEST(test_transpose_0132_fp16); - UNITTEST(test_transpose_0132_bf16); - UNITTEST(test_transpose_0231_fp32); - UNITTEST(test_transpose_0231_fp16); - UNITTEST(test_transpose_0231_bf16); - UNITTEST(test_transpose_0213_fp32); - UNITTEST(test_transpose_0213_fp16); - UNITTEST(test_transpose_0213_bf16); -#if (SYNC_TEST) - UNITTEST(test_transpose_sync_test); -#endif - UNITTEST(test_transpose_invalid); - return ark::unittest::SUCCESS; -} diff --git a/python/unittest/ops/conftest.py b/python/unittest/ops/conftest.py new file mode 100644 index 00000000..e9d26d74 --- /dev/null +++ b/python/unittest/ops/conftest.py @@ -0,0 +1,22 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +""" +Shared fixtures and helpers for ARK op numerical tests. +""" + +import sys +import os +import pytest + +# Add parent directory to path so `common` is importable +sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from common import ark + + +# Auto-initialize ARK before each test; yield allows future teardown. +@pytest.fixture(autouse=True) +def _ark_init(): + """Initialize ARK state before each test so tests start fresh.""" + ark.init() + yield diff --git a/python/unittest/ops/test_arithmetic.py b/python/unittest/ops/test_arithmetic.py new file mode 100644 index 00000000..54d1c5bd --- /dev/null +++ b/python/unittest/ops/test_arithmetic.py @@ -0,0 +1,132 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Numerical tests for arithmetic ops (add, sub, mul, div), constants, and scalar copy.""" + +import pytest + +from common import ark + +torch = pytest.importorskip("torch") + +DEVICE = "cuda:0" + + +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16] +) +def test_add(dtype): + a = torch.randn(8192, dtype=dtype, device=DEVICE) + b = torch.randn(8192, dtype=dtype, device=DEVICE) + assert torch.allclose(ark.add(a, b).eval(), a + b, atol=0, rtol=0) + + +def test_add_broadcast(): + a = torch.randn(4, 1024, dtype=torch.float16, device=DEVICE) + b = torch.randn(1, 1024, dtype=torch.float16, device=DEVICE) + assert torch.allclose(ark.add(a, b).eval(), a + b, atol=0, rtol=0) + + +def test_add_broadcast_3d(): + a = torch.randn(3, 1, 1024, dtype=torch.float16, device=DEVICE) + b = torch.randn(1, 4, 1, dtype=torch.float16, device=DEVICE) + assert torch.allclose(ark.add(a, b).eval(), a + b, atol=0, rtol=0) + + +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16] +) +def test_sub(dtype): + a = torch.randn(8192, dtype=dtype, device=DEVICE) + b = torch.randn(8192, dtype=dtype, device=DEVICE) + assert torch.allclose(ark.sub(a, b).eval(), a - b, atol=0, rtol=0) + + +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16] +) +def test_mul(dtype): + a = torch.randn(8192, dtype=dtype, device=DEVICE) + b = torch.randn(8192, dtype=dtype, device=DEVICE) + assert torch.allclose(ark.mul(a, b).eval(), a * b, atol=0, rtol=0) + + +def test_mul_broadcast(): + a = torch.randn(4, 1024, dtype=torch.float16, device=DEVICE) + b = torch.randn(1, 1024, dtype=torch.float16, device=DEVICE) + assert torch.allclose(ark.mul(a, b).eval(), a * b, atol=0, rtol=0) + + +def test_mul_broadcast_3d(): + a = torch.randn(3, 1, 1024, dtype=torch.float16, device=DEVICE) + b = torch.randn(1, 4, 1, dtype=torch.float16, device=DEVICE) + assert torch.allclose(ark.mul(a, b).eval(), a * b, atol=0, rtol=0) + + +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16] +) +def test_div(dtype): + a = torch.randn(8192, dtype=dtype, device=DEVICE) + b = torch.randn(8192, dtype=dtype, device=DEVICE).abs() + 0.01 + assert torch.allclose(ark.div(a, b).eval(), a / b, atol=0, rtol=0) + + +# Scalar operations + +FACTOR = 0.75 + + +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16] +) +@pytest.mark.parametrize("shape", [(4, 2, 1), (4, 2, 1024)]) +def test_scalar_mul(dtype, shape): + a = torch.randn(shape, dtype=dtype, device=DEVICE) + assert torch.allclose(ark.mul(a, FACTOR).eval(), a * FACTOR, atol=0, rtol=0) + + +@pytest.mark.parametrize("shape", [(4, 2, 1), (4, 2, 1024)]) +def test_scalar_add(shape): + a = torch.randn(shape, dtype=torch.float16, device=DEVICE) + assert torch.allclose(ark.add(a, FACTOR).eval(), a + FACTOR, atol=0, rtol=0) + + +@pytest.mark.parametrize("shape", [(4, 2, 1), (4, 2, 1024)]) +def test_scalar_sub(shape): + a = torch.randn(shape, dtype=torch.float16, device=DEVICE) + assert torch.allclose(ark.sub(a, FACTOR).eval(), a - FACTOR, atol=0, rtol=0) + + +@pytest.mark.parametrize("shape", [(4, 2, 1), (4, 2, 1024)]) +def test_scalar_div(shape): + a = torch.randn(shape, dtype=torch.float16, device=DEVICE) + # fp16 scalar division may differ in rounding from PyTorch + assert torch.allclose( + ark.div(a, FACTOR).eval(), a / FACTOR, atol=1e-3, rtol=1e-3 + ) + + +# Constant & scalar copy +# Scalar copy only; tensor copy tested separately. + + +def test_constant_fp16(): + out = ark.constant(7, (4, 2, 50), ark.fp16).eval() + assert (out == 7).all() + + +def test_constant_fp32(): + out = ark.constant(7, (1,), ark.fp32).eval() + assert out.item() == 7.0 + + +def test_copy_scalar_fp16(): + t = torch.zeros(4, 2, 50, dtype=torch.float16, device=DEVICE) + out = ark.copy(7.0, ark.Tensor.from_torch(t)).eval() + assert (out == 7).all() + + +def test_copy_scalar_fp32(): + out = ark.copy(7.0).eval() + assert out.item() == 7.0 diff --git a/python/unittest/ops/test_cast.py b/python/unittest/ops/test_cast.py new file mode 100644 index 00000000..40b31259 --- /dev/null +++ b/python/unittest/ops/test_cast.py @@ -0,0 +1,52 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Numerical tests for cast op.""" + +import pytest + +from common import ark + +torch = pytest.importorskip("torch") + +DEVICE = "cuda:0" + +# Note: byte↔{fp32,fp16,int32} casts not tested — ark.cast does not expose byte conversions in the Python API. + + +@pytest.mark.parametrize( + "src_dtype, dst_dtype, ark_dst", + [ + (torch.float16, torch.float32, ark.fp32), + (torch.float32, torch.float16, ark.fp16), + (torch.bfloat16, torch.float32, ark.fp32), + (torch.float32, torch.bfloat16, ark.bf16), + ], +) +def test_cast_float(src_dtype, dst_dtype, ark_dst): + a = torch.randn(4, 2, 1024, dtype=torch.float32, device=DEVICE).to( + src_dtype + ) + result = ark.cast(a, ark_dst).eval() + expected = a.to(dst_dtype) + assert result.dtype == dst_dtype + assert torch.allclose(result, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize( + "src_dtype, dst_dtype, ark_dst", + [ + (torch.float32, torch.int32, ark.int32), + (torch.int32, torch.float32, ark.fp32), + (torch.float16, torch.int32, ark.int32), + (torch.int32, torch.float16, ark.fp16), + ], +) +def test_cast_int(src_dtype, dst_dtype, ark_dst): + a = ( + torch.arange(4 * 2 * 1024, device=DEVICE).reshape(4, 2, 1024) % 1000 + ).to(src_dtype) + result = ark.cast(a, ark_dst).eval() + expected = a.to(dst_dtype) + assert result.dtype == dst_dtype + assert torch.allclose(result, expected, atol=0, rtol=0) diff --git a/python/unittest/ops/test_composite.py b/python/unittest/ops/test_composite.py new file mode 100644 index 00000000..c958111f --- /dev/null +++ b/python/unittest/ops/test_composite.py @@ -0,0 +1,47 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Numerical tests for composite ops: softmax, layernorm.""" + +import pytest + +from common import ark + +torch = pytest.importorskip("torch") +import torch.nn.functional as F + +DEVICE = "cuda:0" + + +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16] +) +@pytest.mark.parametrize("shape", [(4, 8, 256), (8, 16), (3, 13, 127), (4, 1)]) +def test_softmax(dtype, shape): + a = torch.randn(shape, dtype=dtype, device=DEVICE) + result = ark.softmax(a).eval() + expected = F.softmax(a, dim=-1) + atol = {torch.float32: 1e-5, torch.float16: 1e-3, torch.bfloat16: 5e-2}[ + dtype + ] + assert torch.allclose( + result, expected, atol=atol, rtol=1e-3 + ), f"max_diff={(result - expected).abs().max()}" + + +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16] +) +@pytest.mark.parametrize("shape", [(4, 8, 256), (8, 16), (3, 13, 127)]) +def test_layernorm(dtype, shape): + a = torch.randn(shape, dtype=dtype, device=DEVICE) + result = ark.layernorm(a, eps=1e-6).eval() + mean = a.mean(dim=-1, keepdim=True) + var = ((a - mean) ** 2).mean(dim=-1, keepdim=True) + expected = (a - mean) / torch.sqrt(var + 1e-6) + atol = {torch.float32: 1e-4, torch.float16: 1e-2, torch.bfloat16: 5e-2}[ + dtype + ] + assert torch.allclose( + result, expected, atol=atol, rtol=1e-4 + ), f"max_diff={(result - expected).abs().max()}" diff --git a/python/unittest/ops/test_embedding_rope.py b/python/unittest/ops/test_embedding_rope.py new file mode 100644 index 00000000..fbae0cf0 --- /dev/null +++ b/python/unittest/ops/test_embedding_rope.py @@ -0,0 +1,58 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Numerical tests for embedding and rope ops.""" + +import pytest + +from common import ark + +torch = pytest.importorskip("torch") +import torch.nn.functional as F + +DEVICE = "cuda:0" + + +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16] +) +def test_embedding(dtype): + vocab_size, embed_dim = 100, 64 + indices = torch.randint(0, vocab_size, (4, 8), device=DEVICE).to( + torch.int32 + ) + weight = torch.randn(vocab_size, embed_dim, dtype=dtype, device=DEVICE) + result = ark.embedding(indices, weight).eval() + expected = F.embedding(indices, weight) + assert torch.allclose( + result, expected, atol=0, rtol=0 + ), f"max_diff={(result - expected).abs().max()}" + + +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16] +) +def test_rope(dtype): + """Test rotary positional embedding against PyTorch complex-multiply reference. + ARK's rope computes element-wise complex multiplication on consecutive pairs: + c[2k] = a[2k]*b[2k] - a[2k+1]*b[2k+1] + c[2k+1] = a[2k]*b[2k+1] + a[2k+1]*b[2k] + """ + shape = (1, 1, 8, 64) + x = torch.randn(shape, dtype=dtype, device=DEVICE) + other = torch.randn(shape, dtype=dtype, device=DEVICE) + result = ark.rope(x, other).eval() + # PyTorch reference: complex multiply on paired elements + a = x.reshape(*shape[:-1], -1, 2) + b = other.reshape(*shape[:-1], -1, 2) + expected = torch.stack( + [ + a[..., 0] * b[..., 0] - a[..., 1] * b[..., 1], + a[..., 0] * b[..., 1] + a[..., 1] * b[..., 0], + ], + dim=-1, + ).reshape(shape) + atol = 1e-5 if dtype == torch.float32 else 5e-2 + assert torch.allclose( + result, expected, atol=atol, rtol=1e-3 + ), f"max_diff={(result - expected).abs().max()}" diff --git a/python/unittest/ops/test_math.py b/python/unittest/ops/test_math.py new file mode 100644 index 00000000..0af479ee --- /dev/null +++ b/python/unittest/ops/test_math.py @@ -0,0 +1,69 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Numerical tests for unary math ops: exp, gelu, relu, sigmoid, sqrt, rsqrt.""" + +import pytest + +from common import ark + +torch = pytest.importorskip("torch") +import torch.nn.functional as F + +DEVICE = "cuda:0" + + +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16] +) +def test_exp(dtype): + a = torch.randn(4, 2, 1024, dtype=dtype, device=DEVICE) + atol = 1e-5 if dtype == torch.float32 else 1e-2 + assert torch.allclose(ark.exp(a).eval(), torch.exp(a), atol=atol, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_gelu(dtype): + a = torch.randn(4, 2, 1024, dtype=dtype, device=DEVICE) + atol = 1e-5 if dtype == torch.float32 else 1e-2 + assert torch.allclose( + ark.gelu(a).eval(), F.gelu(a, approximate="tanh"), atol=atol, rtol=0 + ) + + +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16] +) +def test_relu(dtype): + a = torch.randn(4, 2, 1024, dtype=dtype, device=DEVICE) + assert torch.allclose(ark.relu(a).eval(), F.relu(a), atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_sigmoid(dtype): + a = torch.randn(4, 2, 1024, dtype=dtype, device=DEVICE) + atol = 1e-5 if dtype == torch.float32 else 1e-2 + assert torch.allclose( + ark.sigmoid(a).eval(), torch.sigmoid(a), atol=atol, rtol=0 + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +def test_sqrt(dtype): + a = torch.rand(4, 2, 1024, dtype=dtype, device=DEVICE) + 0.01 + atol = 1e-6 if dtype == torch.float32 else 1e-3 + assert torch.allclose(ark.sqrt(a).eval(), torch.sqrt(a), atol=atol, rtol=0) + + +def test_sqrt_small_last_dim(): + a = torch.rand(4, 2, 7, dtype=torch.float16, device=DEVICE) + 0.01 + assert torch.allclose(ark.sqrt(a).eval(), torch.sqrt(a), atol=1e-3, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +def test_rsqrt(dtype): + a = torch.rand(4, 2, 1024, dtype=dtype, device=DEVICE) + 0.01 + atol = 1e-4 if dtype == torch.float32 else 1e-3 + assert torch.allclose( + ark.rsqrt(a).eval(), torch.rsqrt(a), atol=atol, rtol=0 + ) diff --git a/python/unittest/ops/test_matmul.py b/python/unittest/ops/test_matmul.py new file mode 100644 index 00000000..94363559 --- /dev/null +++ b/python/unittest/ops/test_matmul.py @@ -0,0 +1,74 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Numerical tests for matmul: NN, NT, TN, TT, batched.""" + +import pytest + +from common import ark + +torch = pytest.importorskip("torch") + +DEVICE = "cuda:0" + + +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16] +) +@pytest.mark.parametrize("M,N,K", [(256, 256, 512), (64, 64, 64)]) +def test_matmul_nn(dtype, M, N, K): + a = torch.randn(M, K, dtype=dtype, device=DEVICE) + b = torch.randn(K, N, dtype=dtype, device=DEVICE) + result = ark.matmul(a, b).eval() + expected = a @ b + # fp16/bf16 matmul accumulates rounding error proportional to K. + # For unit-variance randn inputs: expected error ≈ K * eps ≈ 512 * 5e-4 ≈ 0.26. + # 3e-1 provides small headroom. Matches C++ test's reduction_abs_error_bound logic. + atol = 1e-3 if dtype == torch.float32 else 3e-1 + assert torch.allclose( + result, expected, atol=atol, rtol=1e-2 + ), f"max_diff={(result - expected).abs().max()}" + + +def test_matmul_nt(): + M, N, K = 256, 256, 512 + a = torch.randn(M, K, dtype=torch.float16, device=DEVICE) + b = torch.randn(N, K, dtype=torch.float16, device=DEVICE) + result = ark.matmul(a, b, transpose_other=True).eval() + expected = a @ b.t() + assert torch.allclose( + result, expected, atol=3e-1, rtol=1e-2 + ), f"max_diff={(result - expected).abs().max()}" + + +def test_matmul_tn(): + M, N, K = 256, 256, 512 + a = torch.randn(K, M, dtype=torch.float16, device=DEVICE) + b = torch.randn(K, N, dtype=torch.float16, device=DEVICE) + result = ark.matmul(a, b, transpose_input=True).eval() + expected = a.t() @ b + assert torch.allclose( + result, expected, atol=3e-1, rtol=1e-2 + ), f"max_diff={(result - expected).abs().max()}" + + +def test_matmul_tt(): + M, N, K = 256, 256, 512 + a = torch.randn(K, M, dtype=torch.float16, device=DEVICE) + b = torch.randn(N, K, dtype=torch.float16, device=DEVICE) + result = ark.matmul(a, b, transpose_input=True, transpose_other=True).eval() + expected = a.t() @ b.t() + assert torch.allclose( + result, expected, atol=3e-1, rtol=1e-2 + ), f"max_diff={(result - expected).abs().max()}" + + +def test_matmul_batched(): + B, M, N, K = 4, 256, 256, 512 + a = torch.randn(B, M, K, dtype=torch.float16, device=DEVICE) + b = torch.randn(B, K, N, dtype=torch.float16, device=DEVICE) + result = ark.matmul(a, b).eval() + expected = a @ b + assert torch.allclose( + result, expected, atol=3e-1, rtol=1e-2 + ), f"max_diff={(result - expected).abs().max()}" diff --git a/python/unittest/ops/test_reduce.py b/python/unittest/ops/test_reduce.py new file mode 100644 index 00000000..10353c12 --- /dev/null +++ b/python/unittest/ops/test_reduce.py @@ -0,0 +1,66 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Numerical tests for reduce ops: reduce_sum, reduce_max, reduce_mean.""" + +import pytest + +from common import ark + +torch = pytest.importorskip("torch") + +DEVICE = "cuda:0" + + +@pytest.mark.parametrize("axis", [0, 1, 2, 3]) +def test_reduce_sum_fp32(axis): + shape = [7, 2, 4, 1024] + a = torch.randn(shape, dtype=torch.float32, device=DEVICE) * 0.1 + result = ark.reduce_sum(a, axis=axis).eval() + expected = torch.sum(a, dim=axis, keepdim=True) + atol = shape[axis] * 1e-5 + assert torch.allclose( + result, expected, atol=atol, rtol=1e-4 + ), f"axis={axis}, max_diff={(result - expected).abs().max()}" + + +@pytest.mark.parametrize("axis", [0, 3]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_reduce_sum_half(axis, dtype): + shape = [7, 2, 4, 1024] + a = torch.randn(shape, dtype=dtype, device=DEVICE) * 0.1 + result = ark.reduce_sum(a, axis=axis).eval() + expected = torch.sum(a, dim=axis, keepdim=True) + atol = shape[axis] * 2e-2 + assert torch.allclose( + result, expected, atol=atol, rtol=1e-2 + ), f"axis={axis}, max_diff={(result - expected).abs().max()}" + + +def test_reduce_sum_no_keepdims(): + shape = [7, 2, 4, 1024] + a = torch.randn(shape, dtype=torch.float16, device=DEVICE) * 0.1 + result = ark.reduce_sum(a, axis=3, keepdims=False).eval() + expected = torch.sum(a, dim=3, keepdim=False) + atol = shape[3] * 2e-2 + assert torch.allclose( + result, expected, atol=atol, rtol=1e-2 + ), f"max_diff={(result - expected).abs().max()}" + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +def test_reduce_max(dtype): + a = torch.randn(1, 1, 2, 8192, dtype=dtype, device=DEVICE) + result = ark.reduce_max(a, axis=-1).eval() + expected = torch.max(a, dim=-1, keepdim=True).values + assert torch.allclose(result, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +def test_reduce_mean(dtype): + a = torch.randn(1, 1, 2, 8192, dtype=dtype, device=DEVICE) * 0.1 + result = ark.reduce_mean(a, axis=-1).eval() + expected = torch.mean(a, dim=-1, keepdim=True) + atol = 1e-4 if dtype == torch.float32 else 1e-2 + rtol = 1e-4 if dtype == torch.float32 else 1e-2 + assert torch.allclose(result, expected, atol=atol, rtol=rtol) diff --git a/python/unittest/ops/test_transpose.py b/python/unittest/ops/test_transpose.py new file mode 100644 index 00000000..9acd40c4 --- /dev/null +++ b/python/unittest/ops/test_transpose.py @@ -0,0 +1,32 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Numerical tests for transpose op.""" + +import pytest + +from common import ark + +torch = pytest.importorskip("torch") + +DEVICE = "cuda:0" + + +@pytest.mark.parametrize( + "perm, shape", + [ + ([0, 1, 3, 2], [2, 3, 64, 128]), + ([0, 2, 3, 1], [2, 3, 64, 128]), + ([0, 2, 1, 3], [2, 3, 64, 128]), + ], +) +@pytest.mark.parametrize( + "dtype", [torch.float32, torch.float16, torch.bfloat16] +) +def test_transpose(perm, shape, dtype): + a = torch.randn(shape, dtype=dtype, device=DEVICE) + result = ark.transpose(a, perm).eval() + expected = a.permute(perm).contiguous() + assert torch.allclose( + result, expected, atol=0, rtol=0 + ), f"max_diff={(result - expected).abs().max()}" diff --git a/python/unittest/test_data_type_edges.py b/python/unittest/test_data_type_edges.py new file mode 100644 index 00000000..458bd3b8 --- /dev/null +++ b/python/unittest/test_data_type_edges.py @@ -0,0 +1,43 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for data_type.py edge branches.""" + +from common import ark, pytest_ark +import pytest + + +@pytest_ark(need_torch=True) +def test_data_type_from_torch(): + """DataType.from_torch for known types.""" + import torch + + assert ark.DataType.from_torch(torch.float32) == ark.fp32 + assert ark.DataType.from_torch(torch.float16) == ark.fp16 + assert ark.DataType.from_torch(torch.bfloat16) == ark.bf16 + assert ark.DataType.from_torch(torch.int32) == ark.int32 + assert ark.DataType.from_torch(torch.int8) == ark.int8 + assert ark.DataType.from_torch(torch.uint8) == ark.uint8 + + +@pytest_ark(need_torch=True) +def test_data_type_from_torch_unknown(): + """DataType.from_torch raises ValueError for unsupported type.""" + import torch + + with pytest.raises(ValueError): + ark.DataType.from_torch(torch.float64) + + +@pytest_ark() +def test_data_type_bf16_torch_type(): + """bf16 to_torch returns bfloat16.""" + import torch + + assert ark.bf16.to_torch() == torch.bfloat16 + + +@pytest_ark() +def test_data_type_bf16_numpy_none(): + """bf16 has no numpy equivalent.""" + assert ark.bf16.to_numpy() is None diff --git a/python/unittest/test_model_edges.py b/python/unittest/test_model_edges.py new file mode 100644 index 00000000..54fe9336 --- /dev/null +++ b/python/unittest/test_model_edges.py @@ -0,0 +1,74 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for model.py edge branches not covered by test_model.py.""" + +from common import ark, pytest_ark +import pytest + + +@pytest_ark() +def test_model_set_device_id_valid(): + """set_device_id accepts 0.""" + ark.Model.set_device_id(0) + assert ark.Model.get_device_id() == 0 + + +@pytest_ark() +def test_model_set_device_id_negative(): + """set_device_id raises InvalidUsageError for negative value.""" + with pytest.raises(ark.InvalidUsageError): + ark.Model.set_device_id(-1) + + +@pytest_ark() +def test_model_set_rank(): + """set_rank / get_rank round-trips.""" + ark.Model.set_rank(3) + assert ark.Model.get_rank() == 3 + ark.Model.set_rank(0) + + +@pytest_ark() +def test_model_set_world_size(): + """set_world_size / get_world_size round-trips.""" + ark.Model.set_world_size(4) + assert ark.Model.get_world_size() == 4 + ark.Model.set_world_size(1) + + +@pytest_ark() +def test_model_str(): + """Model.__str__ returns serialized JSON.""" + t = ark.tensor([64], ark.fp16) + m = ark.Model.get_model() + s = str(m) + assert isinstance(s, str) + assert "Tensors" in s or "{" in s + + +@pytest_ark() +def test_model_serialize_pretty_false(): + """Model.serialize(pretty=False) returns compact JSON.""" + t = ark.tensor([64], ark.fp16) + m = ark.Model.get_model() + compact = m.serialize(pretty=False) + pretty = m.serialize(pretty=True) + # Compact should be shorter (no indentation) + assert len(compact) <= len(pretty) + + +@pytest_ark() +def test_set_rank_top_level(): + """ark.set_rank top-level function works.""" + ark.set_rank(1) + assert ark.Model.get_rank() == 1 + ark.set_rank(0) + + +@pytest_ark() +def test_set_world_size_top_level(): + """ark.set_world_size top-level function works.""" + ark.set_world_size(8) + assert ark.Model.get_world_size() == 8 + ark.set_world_size(1) diff --git a/python/unittest/test_module.py b/python/unittest/test_module.py new file mode 100644 index 00000000..47f3d52a --- /dev/null +++ b/python/unittest/test_module.py @@ -0,0 +1,116 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from common import ark, pytest_ark +import numpy as np +import pytest + + +@pytest_ark() +def test_module_register_parameter(): + """Test registering parameters on a module.""" + + class Linear(ark.Module): + def __init__(self): + super().__init__() + self.weight = ark.parameter([64, 64], ark.fp16) + + m = Linear() + assert "weight" in m.parameters + assert isinstance(m.parameters["weight"], ark.Parameter) + + +@pytest_ark() +def test_module_register_submodule(): + """Test registering submodules.""" + + class Block(ark.Module): + def __init__(self): + super().__init__() + + class Net(ark.Module): + def __init__(self): + super().__init__() + self.block = Block() + + net = Net() + assert "block" in net.sub_modules + assert isinstance(net.sub_modules["block"], ark.Module) + + +@pytest_ark() +def test_module_params_dict_nested(): + """Test params_dict with nested modules.""" + + class Child(ark.Module): + def __init__(self): + super().__init__() + self.w = ark.parameter([32, 32], ark.fp32) + + class Parent(ark.Module): + def __init__(self): + super().__init__() + self.child = Child() + self.bias = ark.parameter([32], ark.fp32) + + parent = Parent() + pd = parent.params_dict() + assert "child.w" in pd + assert "bias" in pd + + +@pytest_ark() +def test_module_params_dict_prefix(): + """Test params_dict with custom prefix.""" + + class M(ark.Module): + def __init__(self): + super().__init__() + self.w = ark.parameter([8], ark.fp32) + + m = M() + pd = m.params_dict(prefix="layer0.") + assert "layer0.w" in pd + + +@pytest_ark() +def test_module_call_invokes_forward(): + """Test that __call__ invokes forward.""" + + class Adder(ark.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return ark.add(x, 1.0) + + m = Adder() + t = ark.tensor([64], ark.fp16) + out = m(t) + assert out.shape() == [64] + + +@pytest_ark() +def test_module_register_module_type_error(): + """register_module raises TypeError for non-Module.""" + + class M(ark.Module): + def __init__(self): + super().__init__() + + m = M() + with pytest.raises(TypeError): + m.register_module("bad", "not a module") + + +@pytest_ark() +def test_module_register_parameter_type_error(): + """register_parameter raises TypeError for non-Parameter.""" + + class M(ark.Module): + def __init__(self): + super().__init__() + + m = M() + with pytest.raises(TypeError): + m.register_parameter("bad", "not a param") diff --git a/python/unittest/test_ops_edges.py b/python/unittest/test_ops_edges.py new file mode 100644 index 00000000..ca725f8c --- /dev/null +++ b/python/unittest/test_ops_edges.py @@ -0,0 +1,136 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for ops.py error/edge branches not covered by test_ops.py.""" + +from common import ark, pytest_ark +import pytest + + +@pytest_ark() +def test_reshape_invalid_shape_type(): + """reshape raises InvalidUsageError when shape is not a list/tuple.""" + a = ark.tensor([64, 64], ark.fp16) + with pytest.raises(ark.InvalidUsageError): + ark.reshape(a, 64) + + +@pytest_ark() +def test_reshape_exceeds_4d(): + """reshape raises InvalidUsageError for > 4 dimensions.""" + a = ark.tensor([2, 2, 2, 2], ark.fp16) + with pytest.raises(ark.InvalidUsageError): + ark.reshape(a, [1, 2, 2, 2, 2]) + + +@pytest_ark() +def test_transpose_invalid_perm_type(): + """transpose raises InvalidUsageError when perm is not a list/tuple.""" + a = ark.tensor([64, 32], ark.fp16) + with pytest.raises(ark.InvalidUsageError): + ark.transpose(a, 0) + + +@pytest_ark() +def test_transpose_exceeds_4d(): + """transpose raises InvalidUsageError for perm > 4 dimensions.""" + a = ark.tensor([2, 2, 2, 2], ark.fp16) + with pytest.raises(ark.InvalidUsageError): + ark.transpose(a, [0, 1, 2, 3, 4]) + + +@pytest_ark() +def test_identity_invalid_dep(): + """identity raises InvalidUsageError if deps contain non-Tensor.""" + a = ark.tensor([64], ark.fp16) + with pytest.raises(ark.InvalidUsageError): + ark.identity(a, deps=["not_a_tensor"]) + + +@pytest_ark() +def test_add_two_scalars_returns_float(): + """add with two float scalars returns float sum.""" + result = ark.add(2.5, 3.5) + assert result == 6.0 + + +@pytest_ark() +def test_add_two_scalars_with_output(): + """add with two float scalars and an output tensor uses copy.""" + out = ark.tensor([1], ark.fp32) + result = ark.add(2.0, 3.0, output=out) + assert result.shape() == [1] + + +@pytest_ark() +def test_ops_noop(): + """noop does not return a value.""" + a = ark.tensor([64], ark.fp16) + result = ark.noop(a) + assert result is None + + +@pytest_ark() +def test_ops_copy_scalar(): + """copy accepts a scalar value.""" + out = ark.copy(42.0) + assert out.shape() == [1] + + +@pytest_ark() +def test_ops_softmax_shape(): + """softmax composite op produces correct shape.""" + a = ark.tensor([4, 64], ark.fp16) + out = ark.softmax(a) + assert out.shape() == [4, 64] + + +@pytest_ark() +def test_ops_layernorm_shape(): + """layernorm composite op produces correct shape.""" + a = ark.tensor([4, 64], ark.fp16) + out = ark.layernorm(a) + assert out.shape() == [4, 64] + + +@pytest_ark() +def test_ops_mean_alias(): + """mean is an alias for reduce_mean.""" + from ark.ops import mean + + a = ark.tensor([4, 64], ark.fp16) + out = mean(a, axis=1, keepdims=True) + assert out.shape() == [4, 1] + + +@pytest_ark() +def test_ops_ones_zeros(): + """ones and zeros produce correct shapes.""" + o = ark.ones([8, 8]) + assert o.shape() == [8, 8] + z = ark.zeros([8, 8], ark.fp16) + assert z.shape() == [8, 8] + + +@pytest_ark() +def test_ops_send_recv(): + """send/send_done/recv produce tensors (model-only, no actual comm).""" + ark.set_world_size(2) + a = ark.tensor([64], ark.fp16) + s = ark.send(a, remote_rank=1, tag=0) + assert s.shape() == [64] + sd = ark.send_done(s) + assert sd.shape() == [64] + + out = ark.tensor([64], ark.fp16) + r = ark.recv(out, remote_rank=0, tag=0) + assert r.shape() == [64] + + +@pytest_ark() +def test_ops_all_reduce(): + """all_reduce produces correct shape.""" + ark.set_world_size(2) + a = ark.tensor([1024], ark.fp16) + out = ark.all_reduce(a, rank=0, world_size=2) + assert out.shape() == [1024] diff --git a/python/unittest/test_planner_edges.py b/python/unittest/test_planner_edges.py new file mode 100644 index 00000000..1b95e667 --- /dev/null +++ b/python/unittest/test_planner_edges.py @@ -0,0 +1,150 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for planner.py edge branches — Plan creation, from_str, from_file errors.""" + +from common import ark, pytest_ark +import json +import os +import tempfile +import pytest + + +@pytest_ark() +def test_plan_default_init(): + """Plan(None) creates a valid default plan.""" + plan = ark.Plan(None) + assert plan.rank == 0 + assert plan.world_size == 1 + assert plan.architecture == "ANY" + assert plan.num_processors == 1 + assert plan.num_warps_per_processor == 1 + assert plan.task_infos == [] + assert plan.processor_groups == [] + + +@pytest_ark() +def test_plan_str(): + """Plan.__str__ produces valid JSON.""" + plan = ark.Plan(None) + s = str(plan) + parsed = json.loads(s) + assert parsed["Rank"] == 0 + + +@pytest_ark() +def test_plan_from_str_valid(): + """Plan.from_str with valid JSON.""" + data = { + "Rank": 0, + "WorldSize": 1, + "Architecture": "ANY", + "NumProcessors": 1, + "NumWarpsPerProcessor": 1, + "TaskInfos": [], + "ProcessorGroups": [], + } + plan = ark.Plan.from_str(json.dumps(data)) + assert plan.rank == 0 + assert plan.world_size == 1 + + +@pytest_ark() +def test_plan_from_str_invalid_json(): + """Plan.from_str raises InvalidUsageError on bad JSON.""" + with pytest.raises(ark.InvalidUsageError): + ark.Plan.from_str("not valid json {{{") + + +@pytest_ark() +def test_plan_from_file_valid(): + """Plan.from_file loads a valid JSON file.""" + data = { + "Rank": 0, + "WorldSize": 2, + "Architecture": "ANY", + "NumProcessors": 4, + "NumWarpsPerProcessor": 8, + "TaskInfos": [], + "ProcessorGroups": [], + } + with tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False + ) as f: + json.dump(data, f) + path = f.name + + try: + plan = ark.Plan.from_file(path) + assert plan.world_size == 2 + assert plan.num_processors == 4 + finally: + os.unlink(path) + + +@pytest_ark() +def test_plan_from_file_not_found(): + """Plan.from_file raises InvalidUsageError for missing file.""" + with pytest.raises(ark.InvalidUsageError): + ark.Plan.from_file("/tmp/nonexistent_ark_plan_12345.json") + + +@pytest_ark() +def test_plan_from_file_invalid_json(): + """Plan.from_file raises InvalidUsageError for invalid JSON content.""" + with tempfile.NamedTemporaryFile( + mode="w", suffix=".json", delete=False + ) as f: + f.write("not valid json") + path = f.name + + try: + with pytest.raises(ark.InvalidUsageError): + ark.Plan.from_file(path) + finally: + os.unlink(path) + + +@pytest_ark() +def test_planner_context_warp_range(): + """PlannerContext with warp_range kwarg.""" + t = ark.tensor([64, 64], ark.fp16) + with ark.PlannerContext(warp_range=[0, 4]): + ark.add(t, 1.0) + + plan = ark.Planner().plan() + assert len(plan.processor_groups) >= 1 + + +@pytest_ark() +def test_planner_context_sram_range(): + """PlannerContext with sram_range kwarg.""" + t = ark.tensor([64, 64], ark.fp16) + with ark.PlannerContext(sram_range=[0, 1024]): + ark.add(t, 1.0) + + plan = ark.Planner().plan() + assert len(plan.processor_groups) >= 1 + + +@pytest_ark() +def test_planner_context_config(): + """PlannerContext with config dict kwarg does not raise on entry.""" + t = ark.tensor([64, 64], ark.fp16) + # Just verify entering the context with a config does not crash; + # we don't run the planner since invalid configs will error there. + with ark.PlannerContext(config={"NumWarps": 4, "Tile": [64, 64]}): + ark.add(t, 1.0) + # Success: context entered and exited without error + + +@pytest_ark() +def test_planner_context_dump(): + """PlannerContext.dump returns a JSON string.""" + t = ark.tensor([64, 64], ark.fp16) + with ark.PlannerContext(processor_range=[0, 8]) as ctx: + s = ctx.dump() + assert isinstance(s, str) + # Should be valid JSON + parsed = json.loads(s) + assert isinstance(parsed, (dict, list)) diff --git a/python/unittest/test_serialize.py b/python/unittest/test_serialize.py new file mode 100644 index 00000000..5022144b --- /dev/null +++ b/python/unittest/test_serialize.py @@ -0,0 +1,41 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +from common import ark, pytest_ark +import os +import tempfile +import numpy as np +import pytest +from ark.serialize import save, load + + +@pytest_ark() +def test_serialize_save_load(): + """Test round-trip save and load of a state dict.""" + state = {"w": np.ones((4, 4), dtype=np.float32)} + + with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as f: + path = f.name + + try: + save(state, path) + loaded = load(path) + assert "w" in loaded + assert np.array_equal(loaded["w"], state["w"]) + finally: + os.unlink(path) + + +@pytest_ark() +def test_serialize_save_non_dict_warns(): + """save() with non-dict still succeeds (warns).""" + with tempfile.NamedTemporaryFile(suffix=".pkl", delete=False) as f: + path = f.name + + try: + # Should not raise, just warn + save([1, 2, 3], path) + loaded = load(path) + assert loaded == [1, 2, 3] + finally: + os.unlink(path) diff --git a/python/unittest/test_tensor_edges.py b/python/unittest/test_tensor_edges.py new file mode 100644 index 00000000..37153130 --- /dev/null +++ b/python/unittest/test_tensor_edges.py @@ -0,0 +1,149 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT license. + +"""Tests for tensor.py edge/error branches.""" + +from common import ark, pytest_ark +import pytest + + +@pytest_ark() +def test_tensor_shape_strides_nelems_dtype(): + """Basic accessors on a fresh tensor.""" + t = ark.tensor([4, 64], ark.fp16) + assert t.shape() == [4, 64] + assert t.nelems() == 4 * 64 + assert t.dtype() == ark.fp16 + assert isinstance(t.strides(), list) + + +@pytest_ark() +def test_tensor_getitem_int_index(): + """Integer indexing returns a 1-element slice along indexed dim.""" + t = ark.tensor([4, 64], ark.fp16) + s = t[2] + # Only the indexed dimension is reflected in the result shape + assert s.shape() == [1] + + +@pytest_ark() +def test_tensor_getitem_slice(): + """Slice indexing.""" + t = ark.tensor([8, 32], ark.fp16) + s = t[2:6, :16] + assert s.shape() == [4, 16] + + +@pytest_ark() +def test_tensor_getitem_negative_step(): + """Slice with step=-1.""" + t = ark.tensor([8, 32], ark.fp16) + s = t[5:1:-1] + # step -1 swaps start/stop: becomes [2:6] → shape [4] + assert s.shape() == [4] + + +@pytest_ark() +def test_tensor_getitem_invalid_step(): + """Slice with step != 1 or -1 raises UnsupportedError.""" + t = ark.tensor([8, 32], ark.fp16) + with pytest.raises(ark.UnsupportedError): + t[::2] + + +@pytest_ark() +def test_tensor_getitem_too_many_dims(): + """Indexing with more dims than tensor raises InvalidUsageError.""" + t = ark.tensor([8, 32], ark.fp16) + with pytest.raises(ark.InvalidUsageError): + t[0, 0, 0] + + +@pytest_ark() +def test_tensor_getitem_invalid_type(): + """Indexing with non-int/non-slice raises InvalidUsageError.""" + t = ark.tensor([8, 32], ark.fp16) + with pytest.raises(ark.InvalidUsageError): + t["bad"] + + +@pytest_ark() +def test_tensor_hash_eq(): + """Tensor hash and equality.""" + t1 = ark.tensor([64], ark.fp16) + t2 = ark.tensor([64], ark.fp16) + # Same tensor object hashes/equals itself + assert t1 == t1 + assert hash(t1) == hash(t1) + # Different tensors are not equal + assert t1 != t2 + + +@pytest_ark() +def test_tensor_eq_non_tensor(): + """Tensor != non-Tensor object.""" + t = ark.tensor([64], ark.fp16) + assert t != "not a tensor" + assert t != 42 + + +@pytest_ark() +def test_cpp_tensor_invalid_shape(): + """_cpp_tensor raises when shape is not list/tuple.""" + from ark.tensor import _cpp_tensor + + with pytest.raises(ark.InvalidUsageError): + _cpp_tensor(shape=64) + + +@pytest_ark() +def test_cpp_tensor_invalid_strides(): + """_cpp_tensor raises when strides is not list/tuple.""" + from ark.tensor import _cpp_tensor + + with pytest.raises(ark.InvalidUsageError): + _cpp_tensor(shape=[64], strides=64) + + +@pytest_ark() +def test_cpp_tensor_invalid_offsets(): + """_cpp_tensor raises when offsets is not list/tuple.""" + from ark.tensor import _cpp_tensor + + with pytest.raises(ark.InvalidUsageError): + _cpp_tensor(shape=[64], offsets="bad") + + +@pytest_ark() +def test_cpp_tensor_invalid_padded_shape(): + """_cpp_tensor raises when padded_shape is not list/tuple.""" + from ark.tensor import _cpp_tensor + + with pytest.raises(ark.InvalidUsageError): + _cpp_tensor(shape=[64], padded_shape=128) + + +@pytest_ark() +def test_cpp_tensor_exceeds_4d(): + """_cpp_tensor raises ValueError for > 4D shape.""" + from ark.tensor import _cpp_tensor + + with pytest.raises(ValueError): + _cpp_tensor(shape=[1, 2, 3, 4, 5]) + + +@pytest_ark() +def test_tensor_is_external(): + """Regular tensor is not external.""" + t = ark.tensor([64], ark.fp16) + assert not t.is_external() + + +@pytest_ark() +def test_parameter_basic(): + """Parameter creation and attributes.""" + p = ark.parameter([32, 32], ark.fp32) + assert isinstance(p, ark.Parameter) + assert p.shape() == [32, 32] + assert p.dtype() == ark.fp32 + assert p.requires_grad is False