[T1-2-1] S-hhhhh

include/infiniop.h

@@ -17,5 +17,7 @@
 #include "infiniop/ops/sub.h"
 #include "infiniop/ops/swiglu.h"
 #include "infiniop/tensor_descriptor.h"
+#include "infiniop/ops/reduce_mean.h"
+#include "infiniop/ops/reduce_max.h"
 
 #endif // __INFINIOP_API_H__

include/infiniop/ops/reduce_max.h

@@ -0,0 +1,27 @@
+#ifndef __INFINIOP_REDUCE_MAX_API_H__
+#define __INFINIOP_REDUCE_MAX_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopReduceMaxDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateReduceMaxDescriptor(
+    infiniopHandle_t handle,
+    infiniopReduceMaxDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t output_desc,
+    infiniopTensorDescriptor_t input_desc,
+    size_t dim);
+
+__C __export infiniStatus_t infiniopGetReduceMaxWorkspaceSize(infiniopReduceMaxDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopReduceMax(
+    infiniopReduceMaxDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *output,
+    const void *input,
+    void *stream);
+
+__C __export infiniStatus_t infiniopDestroyReduceMaxDescriptor(infiniopReduceMaxDescriptor_t desc);
+
+#endif

include/infiniop/ops/reduce_mean.h

@@ -0,0 +1,27 @@
+#ifndef __INFINIOP_REDUCE_MEAN_API_H__
+#define __INFINIOP_REDUCE_MEAN_API_H__
+
+#include "../operator_descriptor.h"
+
+typedef struct InfiniopDescriptor *infiniopReduceMeanDescriptor_t;
+
+__C __export infiniStatus_t infiniopCreateReduceMeanDescriptor(
+    infiniopHandle_t handle,
+    infiniopReduceMeanDescriptor_t *desc_ptr,
+    infiniopTensorDescriptor_t y_desc,
+    infiniopTensorDescriptor_t x_desc,
+    size_t dim);
+
+__C __export infiniStatus_t infiniopGetReduceMeanWorkspaceSize(infiniopReduceMeanDescriptor_t desc, size_t *size);
+
+__C __export infiniStatus_t infiniopReduceMean(
+    infiniopReduceMeanDescriptor_t desc,
+    void *workspace,
+    size_t workspace_size,
+    void *y,
+    const void *x,
+    void *stream);
+
+__C __export infiniStatus_t infiniopDestroyReduceMeanDescriptor(infiniopReduceMeanDescriptor_t desc);
+
+#endif

scripts/python_test.py

@@ -24,6 +24,8 @@ def run_tests(args):
         "rope.py",
         "sub.py",
         "swiglu.py",
+        "reduce_mean.py",
+        "reduce_max.py",
     ]:
         result = subprocess.run(
             f"python {test} {args} --debug", text=True, encoding="utf-8", shell=True

src/infiniop-test/include/ops.hpp

@@ -16,6 +16,8 @@ DECLARE_INFINIOP_TEST(add)
 DECLARE_INFINIOP_TEST(causal_softmax)
 DECLARE_INFINIOP_TEST(rearrange)
 DECLARE_INFINIOP_TEST(sub)
+DECLARE_INFINIOP_TEST(reduce_mean)
+DECLARE_INFINIOP_TEST(reduce_max)
 
 #define REGISTER_INFINIOP_TEST(name)                      \
     {                                                     \
@@ -43,6 +45,8 @@ DECLARE_INFINIOP_TEST(sub)
         REGISTER_INFINIOP_TEST(causal_softmax) \
         REGISTER_INFINIOP_TEST(rearrange)      \
         REGISTER_INFINIOP_TEST(sub)            \
+        REGISTER_INFINIOP_TEST(reduce_mean)    \
+        REGISTER_INFINIOP_TEST(reduce_max)     \
     }
 
 namespace infiniop_test {
@@ -64,4 +68,4 @@ bool check_names(
 
 } // namespace infiniop_test
 
-#endif
+#endif

src/infiniop-test/include/test.hpp

@@ -47,18 +47,19 @@ std::vector<std::shared_ptr<Result>> runAllTests(
     const GGUFFileReader &,
     infiniDevice_t device, int device_id,
     size_t warm_ups, size_t iterations,
-    double rtol, double atol);
+    double rtol, double atol, bool equal_nan = false);
 
 // Run a single test read from a GGUF file
 std::shared_ptr<Result> runTest(
     const GGUFFileReader &,
     infiniDevice_t device, int device_id,
     size_t warm_ups, size_t iterations,
     double rtol, double atol,
-    size_t test_id);
+    size_t test_id,
+    bool equal_nan = false);
 
 // Check if two tensors are close within given tolerance
-void allClose(std::shared_ptr<Tensor> actual, std::shared_ptr<Tensor> expected, double rtol = 1e-3, double atol = 1e-3);
+void allClose(std::shared_ptr<Tensor> actual, std::shared_ptr<Tensor> expected, double rtol = 1e-3, double atol = 1e-3, bool equal_nan = false);
 
 // Check if two tensors are equal
 void allEqual(std::shared_ptr<Tensor> actual, std::shared_ptr<Tensor> expected);
@@ -85,13 +86,14 @@ class Test {
     namespace infiniop_test::name {                                           \
     class Test : public infiniop_test::base::Test {                           \
         double _rtol, _atol;                                                  \
+        bool _equal_nan;                                                      \
                                                                               \
     public:                                                                   \
         static std::string op_name() { return #name; }                        \
         static std::shared_ptr<Test> build(                                   \
             std::unordered_map<std::string, std::vector<uint8_t>> attributes, \
             std::unordered_map<std::string, std::shared_ptr<Tensor>> tensors, \
-            double, double);                                                  \
+            double, double, bool);                                            \
                                                                               \
         static std::vector<std::string> attribute_names();                    \
         static std::vector<std::string> tensor_names();                       \
@@ -109,15 +111,16 @@ class Test {
         struct Attributes;                                                    \
         Attributes *_attributes;                                              \
         Test() = delete;                                                      \
-        Test(double rtol, double atol) : _rtol(rtol), _atol(atol) {}          \
+        Test(double rtol, double atol, bool equal_nan = false)                \
+            : _rtol(rtol), _atol(atol), _equal_nan(equal_nan) {}              \
     };                                                                        \
     }
 
 namespace infiniop_test {
 using BuilderFunc = std::function<std::shared_ptr<infiniop_test::base::Test>(
     std::unordered_map<std::string, std::vector<uint8_t>>,
     std::unordered_map<std::string, std::shared_ptr<Tensor>>,
-    double, double)>;
+    double, double, bool)>;
 
 // Testcase Registry
 // Each testcase should provid a formatted builder, attribute names, and tensor names

src/infiniop-test/src/gguf.cpp

@@ -53,7 +53,9 @@ GGUFFileReader::GGUFFileReader(const std::string &filepath) {
     try {
         _file = std::make_shared<FileMapping>(filepath);
     } catch (const std::exception &e) {
-        throw e;
+        // throw e;
+        std::cerr << "Error: " << e.what() << std::endl;
+        // throw e;
     }
     _data = _file->ptr();
     _cursor = reinterpret_cast<uint8_t *>(_data);

src/infiniop-test/src/main.cpp

@@ -1,8 +1,8 @@
 #include "gguf.hpp"
 #include "test.hpp"
+#include <cstring>
 #include <infinirt.h>
 #include <iostream>
-
 struct ParsedArgs {
     std::string file_path;                          // Mandatory argument: test.gguf file path
     infiniDevice_t device_type = INFINI_DEVICE_CPU; // Default to CPU
@@ -11,12 +11,13 @@ struct ParsedArgs {
     int iterations = 0;                             // Default to 0 if not given
     double atol = 0.001;                            // Default absolute tolerance
     double rtol = 0.001;                            // Default relative tolerance
+    bool equal_nan = false;                         // Default relative tolerance
 };
 
 void printUsage() {
     std::cout << "Usage:" << std::endl
               << std::endl;
-    std::cout << "infiniop-test <test.gguf> [--<device>[:id]] [--warmup <warmups>] [--run <iterations>] [--atol <atol>] [--rtol <rtol>]" << std::endl
+    std::cout << "infiniop-test <test.gguf> [--<device>[:id]] [--warmup <warmups>] [--run <iterations>] [--atol <atol>] [--rtol <rtol>] [--equal-nan <equal nan>]" << std::endl
               << std::endl;
     std::cout << "  <test.gguf>>" << std::endl;
     std::cout << "    Path to the test gguf file" << std::endl
@@ -36,6 +37,9 @@ void printUsage() {
     std::cout << "  --rtol <relative_tolerance>" << std::endl;
     std::cout << "    (Optional) Relative tolerance for correctness check. Default to 0.001" << std::endl
               << std::endl;
+    std::cout << "  --equal-nan <compare NaNs as equal>" << std::endl;
+    std::cout << "    (Optional) If True, then two NaNs will be considered equal. Default to False" << std::endl
+              << std::endl;
     exit(-1);
 }
 
@@ -91,6 +95,11 @@ ParsedArgs parseArgs(int argc, char *argv[]) {
             else if (arg == "--rtol" && i + 1 < argc) {
                 args.rtol = std::stod(argv[++i]);
             }
+            else if (arg == "--equal-nan" && i + 1 < argc) {
+                args.equal_nan = (strcmp(argv[++i], "True") == 0 || strcmp(argv[i], "true") == 0)
+                                   ? true
+                                   : false;
+            }
             else {
                 printUsage();
             }
@@ -119,7 +128,7 @@ int main(int argc, char *argv[]) {
             reader,
             (infiniDevice_t)args.device_type, args.device_id,
             args.warmups, args.iterations,
-            args.rtol, args.atol);
+            args.rtol, args.atol, args.equal_nan);
 
         std::cout << "=====================================" << std::endl;
         for (auto result : results) {

src/infiniop-test/src/ops/add.cpp

@@ -15,8 +15,8 @@ struct Test::Attributes {
 std::shared_ptr<Test> Test::build(
     std::unordered_map<std::string, std::vector<uint8_t>> attributes,
     std::unordered_map<std::string, std::shared_ptr<Tensor>> tensors,
-    double rtol, double atol) {
-    auto test = std::shared_ptr<Test>(new Test(rtol, atol));
+    double rtol, double atol, bool equal_nan) {
+    auto test = std::shared_ptr<Test>(new Test(rtol, atol, equal_nan));
     test->_attributes = new Attributes();
     if (tensors.find("a") == tensors.end()
         || tensors.find("b") == tensors.end()
@@ -58,7 +58,7 @@ std::shared_ptr<infiniop_test::Result> Test::run(
              return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution."));
 
     try {
-        allClose(c, _attributes->ans, _rtol, _atol);
+        allClose(c, _attributes->ans, _rtol, _atol, _equal_nan);
     } catch (const std::exception &e) {
         return TEST_FAILED(RESULT_INCORRECT, e.what());
     }
@@ -98,7 +98,7 @@ std::string Test::toString() const {
     oss << "- b: " << _attributes->b->info() << std::endl;
     oss << "- c: " << _attributes->c->info() << std::endl;
     oss << std::scientific << std::setprecision(2);
-    oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl;
+    oss << "- rtol=" << _rtol << ", atol=" << _atol << ", equal_nan=" << _equal_nan << std::endl;
     return oss.str();
 }
 

src/infiniop-test/src/ops/causal_softmax.cpp

@@ -14,8 +14,8 @@ struct Test::Attributes {
 std::shared_ptr<Test> Test::build(
     std::unordered_map<std::string, std::vector<uint8_t>> attributes,
     std::unordered_map<std::string, std::shared_ptr<Tensor>> tensors,
-    double rtol, double atol) {
-    auto test = std::shared_ptr<Test>(new Test(rtol, atol));
+    double rtol, double atol, bool equal_nan) {
+    auto test = std::shared_ptr<Test>(new Test(rtol, atol, equal_nan));
     test->_attributes = new Attributes();
     if (tensors.find("x") == tensors.end()
         || tensors.find("y") == tensors.end()
@@ -53,7 +53,7 @@ std::shared_ptr<infiniop_test::Result> Test::run(
              return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution."));
 
     try {
-        allClose(y, _attributes->ans, _rtol, _atol);
+        allClose(y, _attributes->ans, _rtol, _atol, _equal_nan);
     } catch (const std::exception &e) {
         return TEST_FAILED(RESULT_INCORRECT, e.what());
     }
@@ -92,7 +92,7 @@ std::string Test::toString() const {
     oss << "- y: " << _attributes->y->info() << std::endl;
     oss << "- ans: " << _attributes->ans->info() << std::endl;
     oss << std::scientific << std::setprecision(2);
-    oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl;
+    oss << "- rtol=" << _rtol << ", atol=" << _atol << ", equal_nan=" << _equal_nan << std::endl;
     return oss.str();
 }
 

src/infiniop-test/src/ops/clip.cpp

@@ -16,8 +16,8 @@ struct Test::Attributes {
 std::shared_ptr<Test> Test::build(
     std::unordered_map<std::string, std::vector<uint8_t>> attributes,
     std::unordered_map<std::string, std::shared_ptr<Tensor>> tensors,
-    double rtol, double atol) {
-    auto test = std::shared_ptr<Test>(new Test(rtol, atol));
+    double rtol, double atol, bool equal_nan) {
+    auto test = std::shared_ptr<Test>(new Test(rtol, atol, equal_nan));
     test->_attributes = new Attributes();
     if (tensors.find("x") == tensors.end()
         || tensors.find("min_val") == tensors.end()
@@ -64,7 +64,7 @@ std::shared_ptr<infiniop_test::Result> Test::run(
              return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution."));
 
     try {
-        allClose(y, _attributes->ans, _rtol, _atol);
+        allClose(y, _attributes->ans, _rtol, _atol, _equal_nan);
     } catch (const std::exception &e) {
         return TEST_FAILED(RESULT_INCORRECT, e.what());
     }
@@ -109,7 +109,7 @@ std::string Test::toString() const {
     oss << "- max_val: " << _attributes->max_val->info() << std::endl;
     oss << "- y: " << _attributes->y->info() << std::endl;
     oss << std::scientific << std::setprecision(2);
-    oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl;
+    oss << "- rtol=" << _rtol << ", atol=" << _atol << ", equal_nan=" << _equal_nan << std::endl;
     return oss.str();
 }
 

src/infiniop-test/src/ops/gemm.cpp

@@ -18,8 +18,8 @@ struct Test::Attributes {
 std::shared_ptr<Test> Test::build(
     std::unordered_map<std::string, std::vector<uint8_t>> attributes,
     std::unordered_map<std::string, std::shared_ptr<Tensor>> tensors,
-    double rtol, double atol) {
-    auto test = std::shared_ptr<Test>(new Test(rtol, atol));
+    double rtol, double atol, bool equal_nan) {
+    auto test = std::shared_ptr<Test>(new Test(rtol, atol, equal_nan));
     test->_attributes = new Attributes();
     if (!check_names(attributes, Test::attribute_names()) || !check_names(tensors, Test::tensor_names())) {
         throw std::runtime_error("Invalid Test");
@@ -65,7 +65,7 @@ std::shared_ptr<infiniop_test::Result> Test::run(
              return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution."));
 
     try {
-        allClose(c, _attributes->ans, _rtol, _atol);
+        allClose(c, _attributes->ans, _rtol, _atol, _equal_nan);
     } catch (const std::exception &e) {
         return TEST_FAILED(RESULT_INCORRECT, e.what());
     }
@@ -120,7 +120,7 @@ std::string Test::toString() const {
     oss << "- b: " << _attributes->b->info() << std::endl;
     oss << "- c: " << _attributes->c->info() << std::endl;
     oss << std::scientific << std::setprecision(2);
-    oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl;
+    oss << "- rtol=" << _rtol << ", atol=" << _atol << ", equal_nan=" << _equal_nan << std::endl;
     return oss.str();
 }
 

src/infiniop-test/src/ops/mul.cpp

@@ -15,8 +15,8 @@ struct Test::Attributes {
 std::shared_ptr<Test> Test::build(
     std::unordered_map<std::string, std::vector<uint8_t>> attributes,
     std::unordered_map<std::string, std::shared_ptr<Tensor>> tensors,
-    double rtol, double atol) {
-    auto test = std::shared_ptr<Test>(new Test(rtol, atol));
+    double rtol, double atol, bool equal_nan) {
+    auto test = std::shared_ptr<Test>(new Test(rtol, atol, equal_nan));
     test->_attributes = new Attributes();
     if (tensors.find("a") == tensors.end()
         || tensors.find("b") == tensors.end()
@@ -58,7 +58,7 @@ std::shared_ptr<infiniop_test::Result> Test::run(
              return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution."));
 
     try {
-        allClose(c, _attributes->ans, _rtol, _atol);
+        allClose(c, _attributes->ans, _rtol, _atol, _equal_nan);
     } catch (const std::exception &e) {
         return TEST_FAILED(RESULT_INCORRECT, e.what());
     }
@@ -98,7 +98,7 @@ std::string Test::toString() const {
     oss << "- b: " << _attributes->b->info() << std::endl;
     oss << "- c: " << _attributes->c->info() << std::endl;
     oss << std::scientific << std::setprecision(2);
-    oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl;
+    oss << "- rtol=" << _rtol << ", atol=" << _atol << ", equal_nan=" << _equal_nan << std::endl;
     return oss.str();
 }
 

src/infiniop-test/src/ops/random_sample.cpp

@@ -20,8 +20,8 @@ struct Test::Attributes {
 std::shared_ptr<Test> Test::build(
     std::unordered_map<std::string, std::vector<uint8_t>> attributes,
     std::unordered_map<std::string, std::shared_ptr<Tensor>> tensors,
-    double rtol, double atol) {
-    auto test = std::shared_ptr<Test>(new Test(rtol, atol));
+    double rtol, double atol, bool equal_nan) {
+    auto test = std::shared_ptr<Test>(new Test(rtol, atol, equal_nan));
     test->_attributes = new Attributes();
     if (!check_names(attributes, Test::attribute_names()) || !check_names(tensors, Test::tensor_names())) {
         throw std::runtime_error("Invalid Test");
@@ -70,7 +70,7 @@ std::shared_ptr<infiniop_test::Result> Test::run(
              return TEST_FAILED(OP_EXECUTION_FAILED, "Failed during execution."));
 
     try {
-        allClose(result, _attributes->ans, _rtol, _atol);
+        allClose(result, _attributes->ans, _rtol, _atol, _equal_nan);
     } catch (const std::exception &e) {
         return TEST_FAILED(RESULT_INCORRECT, e.what());
     }
@@ -117,7 +117,7 @@ std::string Test::toString() const {
     oss << "- data: " << _attributes->data->info() << std::endl;
     oss << "- result: " << _attributes->result->info() << std::endl;
     oss << std::scientific << std::setprecision(2);
-    oss << "- rtol=" << _rtol << ", atol=" << _atol << std::endl;
+    oss << "- rtol=" << _rtol << ", atol=" << _atol << ", equal_nan=" << _equal_nan << std::endl;
     return oss.str();
 }