issue/228: infiniop-test框架支持0步长

src/infiniop-test/include/ops.hpp

@@ -29,8 +29,8 @@ DECLARE_INFINIOP_TEST(add)
         REGISTER_INFINIOP_TEST(random_sample) \
         REGISTER_INFINIOP_TEST(add)           \
         REGISTER_INFINIOP_TEST(mul)           \
-        REGISTER_INFINIOP_TEST(swiglu)
-}
+        REGISTER_INFINIOP_TEST(swiglu)        \
+    }
 
 namespace infiniop_test {
 

src/infiniop-test/include/tensor.hpp

@@ -58,7 +58,8 @@ private:
 public:
     Tensor(const GGUFTensorInfo *info,
            const void *ggml_ptr,
-           const GGUFKeyValue *strides_meta = nullptr);
+           const GGUFKeyValue *strides_meta = nullptr,
+           const GGUFKeyValue *shapes_meta = nullptr);
     Tensor(std::shared_ptr<Memory> memory, size_t offset,
            const std::vector<size_t> &shape,
            const std::vector<ptrdiff_t> &strides,

src/infiniop-test/src/tensor.cpp

@@ -98,7 +98,8 @@ void *Tensor::data() const {
 
 Tensor::Tensor(const GGUFTensorInfo *info,
                const void *ggml_ptr,
-               const GGUFKeyValue *strides_meta) {
+               const GGUFKeyValue *strides_meta,
+               const GGUFKeyValue *shapes_meta) {
     _ggml_type = info->ggml_type;
     _offset = 0;
     size_t ndim = static_cast<size_t>(info->ndim);
@@ -120,7 +121,6 @@ Tensor::Tensor(const GGUFTensorInfo *info,
         }
     } else {
         for (size_t i = 0; i < ndim; i++) {
-            _shape[i] = static_cast<size_t>(info->shape[ndim - 1 - i]);
             if (strides_meta->gguf_type == GGUF_TYPE_INT64) {
                 _strides[i] = (ptrdiff_t)(reinterpret_cast<const int64_t *>(
                     strides_meta->value.data())[ndim - 1 - i]);
@@ -145,6 +145,23 @@ Tensor::Tensor(const GGUFTensorInfo *info,
         contiguous_strides.data(),
         ndim,
         ggmlTypeSize(_ggml_type));
+
+    if (shapes_meta != nullptr) {
+
+        for (size_t i = 0; i < ndim; i++) {
+
+            if (shapes_meta->gguf_type == GGUF_TYPE_INT64) {
+                _shape[i] = (ptrdiff_t)(reinterpret_cast<const int64_t *>(
+                    shapes_meta->value.data())[i]);
+
+            } else if (shapes_meta->gguf_type == GGUF_TYPE_INT32) {
+                _shape[i] = (ptrdiff_t)(reinterpret_cast<const int32_t *>(
+                    shapes_meta->value.data())[i]);
+            } else {
+                throw std::runtime_error("Error Creating Tensor: Unsupported shape type");
+            }
+        }
+    }
 }
 
 Tensor::Tensor(std::shared_ptr<Memory> memory, size_t offset,

src/infiniop-test/src/test.cpp

@@ -94,10 +94,13 @@ std::shared_ptr<Result> runTest(const GGUFFileReader &gguf_reader,
             auto info = tensor_info.find("test." + std::to_string(test_id) + "." + tensor_name);
             if (info != tensor_info.end()) {
                 auto strides = meta.find("test." + std::to_string(test_id) + "." + tensor_name + ".strides");
+                auto shape = meta.find("test." + std::to_string(test_id) + "." + tensor_name + ".shape");
+
                 tensors[tensor_name] = std::make_shared<Tensor>(
                     info->second.get(),
                     gguf_reader.getGgmlStart(),
-                    strides != meta.end() ? strides->second.get() : nullptr);
+                    strides != meta.end() ? strides->second.get() : nullptr,
+                    shape != meta.end() ? shape->second.get() : nullptr);
             }
         }
         std::shared_ptr<infiniop_test::base::Test> test;

test/infiniop-test/test_generate/testcases/add.py

@@ -2,6 +2,7 @@ from ast import List
 import numpy as np
 import gguf
 from typing import List
+from numpy.lib.stride_tricks import as_strided
 
 from .. import InfiniopTestWriter, InfiniopTestCase, np_dtype_to_ggml, gguf_strides
 
@@ -12,22 +13,48 @@ def add(
 ):
     return a + b
 
+def process_tensors(a, a_stride, b, b_stride):
+
+    def _rearrange(tensor, strides):
+        if strides and 0 in strides:
+            byte_strides = tuple(s * tensor.itemsize for s in strides)
+            return as_strided(tensor, shape=tensor.shape, strides=byte_strides)
+        else:
+            return tensor
+
+    a = _rearrange(a, a_stride)
+    b = _rearrange(b, b_stride)
+
+    return a, b
+
+def get_effective_shape(shape, strides):
+
+    effective_shape = tuple(dim if stride != 0 else 1 for dim, stride in zip(shape, strides))
+    return effective_shape
 
 class AddTestCase(InfiniopTestCase):
     def __init__(
         self,
         a: np.ndarray,
+        a_rearranged:np.ndarray,
         stride_a: List[int] | None,
+        shape_a: List[int] | None,
         b: np.ndarray,
+        b_rearranged:np.ndarray,
         stride_b: List[int] | None,
+        shape_b: List[int] | None,
         c: np.ndarray,
         stride_c: List[int] | None,
     ):
         super().__init__("add")
         self.a = a
+        self.a_rearranged = a_rearranged
         self.stride_a = stride_a
+        self.shape_a = shape_a
         self.b = b
+        self.b_rearranged = b_rearranged
         self.stride_b = stride_b
+        self.shape_b = shape_b
         self.c = c
         self.stride_c = stride_c
 
@@ -40,6 +67,10 @@ class AddTestCase(InfiniopTestCase):
             test_writer.add_array(test_writer.gguf_key("b.strides"), self.stride_b)
         if self.stride_c is not None:
             test_writer.add_array(test_writer.gguf_key("c.strides"), self.stride_c)
+        if self.shape_a is not None:
+            test_writer.add_array(test_writer.gguf_key("a.shape"), self.shape_a)
+        if self.shape_b is not None:
+            test_writer.add_array(test_writer.gguf_key("b.shape"), self.shape_b)
         test_writer.add_tensor(
             test_writer.gguf_key("a"), self.a, raw_dtype=np_dtype_to_ggml(self.a.dtype)
         )
@@ -50,8 +81,8 @@ class AddTestCase(InfiniopTestCase):
             test_writer.gguf_key("c"), self.c, raw_dtype=np_dtype_to_ggml(self.c.dtype)
         )
         ans = add(
-            self.a.astype(np.float64),
-            self.b.astype(np.float64),
+            self.a_rearranged.astype(np.float64),
+            self.b_rearranged.astype(np.float64),
         )
         test_writer.add_tensor(
             test_writer.gguf_key("ans"), ans, raw_dtype=gguf.GGMLQuantizationType.F64
@@ -69,28 +100,50 @@ if __name__ == "__main__":
         # shape, a_stride, b_stride, c_stride
         ((13, 4), None, None, None),
         ((13, 4), gguf_strides(10, 1), gguf_strides(10, 1), gguf_strides(10, 1)),
+        ((13, 4), gguf_strides(0, 1), None, None),
         ((13, 4, 4), None, None, None),
         ((13, 4, 4), gguf_strides(20, 4, 1), gguf_strides(20, 4, 1), gguf_strides(20, 4, 1)),
+        ((13, 4, 4), gguf_strides(4, 0, 1), gguf_strides(0, 4, 1), None),
         ((16, 5632), None, None, None),
         ((16, 5632), gguf_strides(13312, 1), gguf_strides(13312, 1), gguf_strides(13312, 1)),
         ((4, 4, 5632), None, None, None),
         ((4, 4, 5632), gguf_strides(45056, 5632, 1), gguf_strides(45056, 5632, 1), gguf_strides(45056, 5632, 1)),
     ]
-    _TENSOR_DTYPES_ = [np.float16, np.float32]
+    _TENSOR_DTYPES_ = [np.float32] # np.float16
     for dtype in _TENSOR_DTYPES_:
         for shape, stride_a, stride_b, stride_c in _TEST_CASES_:
             a = np.random.rand(*shape).astype(dtype)
             b = np.random.rand(*shape).astype(dtype)
             c = np.random.rand(*shape).astype(dtype)
 
+            # Reverse strides to match internal layout expectations
+            reversed_stride_a = tuple(reversed(stride_a)) if stride_a else None
+            reversed_stride_b = tuple(reversed(stride_b)) if stride_b else None
+
+            a_rearranged, b_rearranged = process_tensors(a, reversed_stride_a, b, reversed_stride_b)
+
+            effective_shape_a = get_effective_shape(a_rearranged.shape, tuple(s // a.itemsize for s in a_rearranged.strides))
+            effective_shape_b = get_effective_shape(b_rearranged.shape, tuple(s // b.itemsize for s in b_rearranged.strides))
+
+            # Extract unique data region (eliminate broadcast repetition)
+            slices_a = tuple(slice(0, 1) if dim == 1 else slice(None) for dim in effective_shape_a)
+            slices_b = tuple(slice(0, 1) if dim == 1 else slice(None) for dim in effective_shape_b)
+
+            a_unique = a_rearranged[slices_a]
+            b_unique = b_rearranged[slices_b]
+
             test_case = AddTestCase(
-                a=a,
+                a=a_unique,
+                a_rearranged=a_rearranged,
                 stride_a=stride_a,
-                b=b,
+                shape_a=shape,
+                b=b_unique,
+                b_rearranged=b_rearranged,
                 stride_b=stride_b,
+                shape_b=shape,
                 c=c,
                 stride_c=stride_c,
             )
             test_cases.append(test_case)
     test_writer.add_tests(test_cases)
-    test_writer.save()
+    test_writer.save()
\ No newline at end of file