Add sparse any-rank tensor support

.github/workflows/yateto-ci.yml

@@ -33,7 +33,7 @@ jobs:
           mkdir -p ./build-${build_type} && cd ./build-${build_type}
           cmake .. -GNinja -DCMAKE_BUILD_TYPE=${build_type}
           ninja
-          ninja test
+          ctest -j --rerun-failed --output-on-failure
           cd ..
         done
 
@@ -73,7 +73,7 @@ jobs:
     - name: Codegen Tests
       run: |
         cd ./tests/code-gen
-        for example in matmul minimal indices slicing; do
+        for example in matmul minimal indices slicing sparsity; do
           for build_type in Debug Release; do
             for precision in single double; do
               echo " ====== Test Config: ======"

include/yateto/TensorView.h

@@ -4,7 +4,7 @@
 #include <cassert>
 #include <cstring>
 #include <algorithm>
-#include <iterator>
+#include <functional>
 #include <limits>
 #include <type_traits>
 
@@ -47,7 +47,7 @@ namespace yateto {
     }
 
   protected:
-    uint_t m_shape[Dim];
+    uint_t m_shape[Dim]{};
   };
 
   template<typename real_t, typename uint_t>
@@ -113,9 +113,55 @@ namespace yateto {
     uint_t size() const {
       return (m_stop[Dim-1]-m_start[Dim-1]) * m_stride[Dim-1];
     }
+
+    template<typename F>
+    void forall(F&& function) {
+      uint_t entry[Dim]{};
+      std::copy(m_start, m_start + Dim, entry);
+      while (entry[Dim-1] != m_stop[Dim-1]) {
+        auto values = &operator[](entry);
+        for (uint_t i = 0; i < m_stop[0]-m_start[0]; ++i) {
+          entry[0] = i + m_start[0];
+          std::invoke(std::forward<F>(function), entry, values[i*m_stride[0]]);
+        }
+        if (Dim == 1) {
+          break;
+        }
+
+        uint_t d = 0;
+        do {
+          entry[d] = m_start[d];
+          d++;
+          ++entry[d];
+        } while (entry[d] == m_stop[d] && d < Dim-1);
+      }
+    }
+
+    template<typename F>
+    void forall(F&& function) const {
+      uint_t entry[Dim]{};
+      std::copy(m_start, m_start + Dim, entry);
+      while (entry[Dim-1] != m_stop[Dim-1]) {
+        auto values = &operator[](entry);
+        for (uint_t i = 0; i < m_stop[0]-m_start[0]; ++i) {
+          entry[0] = i + m_start[0];
+          std::invoke(std::forward<F>(function), entry, values[i*m_stride[0]]);
+        }
+        if (Dim == 1) {
+          break;
+        }
+
+        uint_t d = 0;
+        do {
+          entry[d] = m_start[d];
+          d++;
+          ++entry[d];
+        } while (entry[d] == m_stop[d] && d < Dim-1);
+      }
+    }
 
     void setZero() {
-      uint_t entry[Dim];
+      uint_t entry[Dim]{};
       std::copy(m_start, m_start + Dim, entry);
       while (entry[Dim-1] != m_stop[Dim-1]) {
         auto values = &operator[](entry);
@@ -192,7 +238,7 @@ namespace yateto {
     void copyToView(view_t& other) const {
       assert(Dim == other.dim());
 
-      uint_t entry[Dim];
+      uint_t entry[Dim]{};
       for (uint_t d = 0; d < Dim; ++d) {
         assert(this->shape(d) == other.shape(d));
         entry[d] = m_start[d];
@@ -224,9 +270,9 @@ namespace yateto {
     auto subtensor(Entry... entry) -> DenseTensorView<count_slices<uint_t, Entry...>::value, real_t, uint_t, Const> {
       static_assert(sizeof...(entry) == Dim, "Number of arguments to subtensor() does not match tensor dimension.");
       constexpr auto nSlices = count_slices<uint_t, Entry...>::value;
-      uint_t begin[Dim];
-      uint_t size[nSlices];
-      uint_t stride[nSlices];
+      uint_t begin[Dim]{};
+      uint_t size[nSlices]{};
+      uint_t stride[nSlices]{};
       extractSubtensor(begin, size, stride, entry...);
       DenseTensorView<nSlices, real_t, uint_t, Const> subtensor(&operator[](begin), size, stride);
       return subtensor;
@@ -300,9 +346,9 @@ namespace yateto {
     }
 
     data_t m_values;
-    uint_t m_start[Dim];
-    uint_t m_stop[Dim];
-    uint_t m_stride[Dim];
+    uint_t m_start[Dim]{};
+    uint_t m_stop[Dim]{};
+    uint_t m_stride[Dim]{};
   };
 
   template<typename real_t, typename uint_t, bool Const>
@@ -422,11 +468,168 @@ namespace yateto {
       }
     }
 
+    template<typename F>
+    void forall(F&& function) {
+      uint_t entry[2];
+      uint_t ncols = this->shape(1);
+      for (uint_t col = 0; col < ncols; ++col) {
+        entry[1] = col;
+        for (uint_t i = m_colPtr[col]; i < m_colPtr[col+1]; ++i) {
+          entry[0] = m_rowInd[i];
+          std::invoke(std::forward<F>(function), entry, m_values[i]);
+        }
+      }
+    }
+
+    template<typename F>
+    void forall(F&& function) const {
+      uint_t entry[2];
+      uint_t ncols = this->shape(1);
+      for (uint_t col = 0; col < ncols; ++col) {
+        entry[1] = col;
+        for (uint_t i = m_colPtr[col]; i < m_colPtr[col+1]; ++i) {
+          entry[0] = m_rowInd[i];
+          std::invoke(std::forward<F>(function), entry, m_values[i]);
+        }
+      }
+    }
+
   protected:
     data_t m_values;
     uint_t const* m_rowInd;
     uint_t const* m_colPtr;
   };
-}
+
+  template<unsigned Dim, typename real_t, typename uint_t, bool Const = false>
+  class PatternTensorView : public TensorView<Dim, real_t, uint_t> {
+  public:
+    using data_t = std::conditional_t<Const, const real_t*, real_t*>;
+    using dataref_t = std::conditional_t<Const, const real_t&, real_t&>;
+
+    explicit PatternTensorView(data_t values, std::initializer_list<uint_t> shape, uint_t const* pattern)
+      : TensorView<Dim, real_t, uint_t>(shape), m_values(values), m_pattern(pattern, shape) {
+
+      computeSize();
+    }
+
+    explicit PatternTensorView(data_t values, uint_t const shape[], uint_t const* pattern)
+      : TensorView<Dim, real_t, uint_t>(shape), m_values(values), m_pattern(pattern, shape) {
+
+      computeSize();
+    }
+
+    explicit PatternTensorView(data_t values, std::initializer_list<uint_t> shape, DenseTensorView<Dim, uint_t, uint_t, true> pattern)
+      : TensorView<Dim, real_t, uint_t>(shape), m_values(values), m_pattern(std::move(pattern)) {
+
+      computeSize();
+    }
+
+    explicit PatternTensorView(data_t values, uint_t const shape[], DenseTensorView<Dim, uint_t, uint_t, true> pattern)
+      : TensorView<Dim, real_t, uint_t>(shape), m_values(values), m_pattern(std::move(pattern)) {
+
+      computeSize();
+    }
+
+    void computeSize() {
+      m_pattern.forall([&](const auto& /*index*/, const auto& idxval) {
+        if (idxval > 0) {
+          ++m_size;
+        }
+      });
+    }
+
+    uint_t size() const {
+      return m_size;
+    }
+
+    void setZero() {
+      m_pattern.forall([&](const auto& /*index*/, const auto& idxval) {
+        if (idxval > 0) {
+          m_values[idxval - 1] = 0;
+        }
+      });
+    }
+
+    template<typename ...Args>
+    const real_t& operator()(Args... index) const {
+      static_assert((std::is_integral_v<Args> && ...));
+      const auto idx = m_pattern(index...);
+      return m_values[idx - 1];
+    }
+
+    template<typename ...Args>
+    dataref_t operator()(Args... index) {
+      static_assert((std::is_integral_v<Args> && ...));
+      const auto idx = m_pattern(index...);
+      return m_values[idx - 1];
+    }
+
+    template<typename ...Args>
+    bool isInRange(Args... index) const {
+      static_assert((std::is_integral_v<Args> && ...));
+      const auto idx = m_pattern(index...);
+      return idx > 0;
+    }
+
+    dataref_t operator[](const uint_t entry[Dim]) {
+      const auto idx = m_pattern[entry];
+      return m_values[idx - 1];
+    }
+
+    const real_t& operator[](const uint_t entry[Dim]) const {
+      const auto idx = m_pattern[entry];
+      return m_values[idx - 1];
+    }
+
+    template<typename F>
+    void forall(F&& function) {
+      m_pattern.forall([&, function = std::forward<F>(function)](const auto& index, const auto& idxval) {
+        if (idxval > 0) {
+          std::invoke(function, index, m_values[idxval - 1]);
+        }
+      });
+    }
+
+    template<typename F>
+    void forall(F&& function) const {
+      m_pattern.forall([&, function = std::forward<F>(function)](const auto& index, const auto& idxval) {
+        if (idxval > 0) {
+          std::invoke(function, index, m_values[idxval - 1]);
+        }
+      });
+    }
+
+    template<class view_t>
+    void copyToView(view_t& other) const {
+      m_pattern.forall([&](const auto& index, const auto& idxval) {
+        if (idxval > 0) {
+          other[index] = m_values[idxval - 1];
+        }
+        else {
+          other[index] = 0;
+        }
+      });
+    }
+
+    template<typename... Entry>
+    auto subtensor(Entry... entry) {
+      static_assert(sizeof...(entry) == Dim, "Number of arguments to subtensor() does not match tensor dimension.");
+      const auto patternSubtensor = m_pattern.subtensor(entry...);
+      return PatternTensorView<count_slices<uint_t, Entry...>::value, real_t, uint_t, Const>(m_values, this->m_shape, patternSubtensor);
+    }
+
+    template<typename... Entry>
+    auto subtensor(Entry... entry) const {
+      static_assert(sizeof...(entry) == Dim, "Number of arguments to subtensor() does not match tensor dimension.");
+      const auto patternSubtensor = m_pattern.subtensor(entry...);
+      return PatternTensorView<count_slices<uint_t, Entry...>::value, real_t, uint_t, true>(m_values, this->m_shape, patternSubtensor);
+    }
+
+  protected:
+    data_t m_values{nullptr};
+    DenseTensorView<Dim, uint_t, uint_t, true> m_pattern;
+    std::size_t m_size{0};
+  };
+} // namespace yateto
 
 #endif

tests/code-gen/sparsity.py

@@ -0,0 +1,64 @@
+#!/usr/bin/env python3
+
+from yateto import *
+
+from yateto.memory import *
+
+import numpy as np
+
+def checkerboard(shape):
+  # cf. https://stackoverflow.com/a/51715491
+  return np.indices(tuple(shape)).sum(axis=0) % 2
+
+def invert(pattern):
+  return 1 - pattern
+
+def add(g):
+  N = 4
+  M = 2
+  A = Tensor('A', (N,), spp=np.array([1,0,0,1]), memoryLayoutClass=PatternMemoryLayout)
+  B = Tensor('B', (N,), spp=np.array([1,0,1,0]), memoryLayoutClass=PatternMemoryLayout)
+  D = Tensor('D', (N,), spp=np.ones((N,)), memoryLayoutClass=PatternMemoryLayout)
+  E = Tensor('E', (N,), spp=checkerboard((N,)), memoryLayoutClass=PatternMemoryLayout)
+  B2 = Tensor('B2', (N, N), spp=np.ones((N, N)), memoryLayoutClass=PatternMemoryLayout)
+  B3 = Tensor('B3', (N, N), spp=np.ones((N, N)), memoryLayoutClass=CSCMemoryLayout)
+  B4 = Tensor('B4', (N, N), spp=checkerboard((N, N)), memoryLayoutClass=PatternMemoryLayout)
+  B5 = Tensor('B5', (N, N), spp=invert(checkerboard((N, N))), memoryLayoutClass=PatternMemoryLayout)
+  Z = Tensor('Z', (N, N, N), spp=np.ones((N, N, N)), memoryLayoutClass=PatternMemoryLayout)
+  C0 = Tensor('C0', (N,))
+  C = Tensor('C', (N, N))
+  C2 = Tensor('C2', (N, N, N))
+
+  YA = Tensor('YA', (M, M, M), spp=checkerboard((M, M, M)), memoryLayoutClass=PatternMemoryLayout)
+  YB = Tensor('YB', (M, M, M), spp=checkerboard((M, M, M)), memoryLayoutClass=PatternMemoryLayout)
+  YC = Tensor('YC', (M, M, M))
+  C4 = Tensor('C4', (M, M, M, M))
+
+  ZA = Tensor('ZA', (M, M, M, M, M, M), spp=checkerboard((M, M, M, M, M, M)), memoryLayoutClass=PatternMemoryLayout)
+  ZB = Tensor('ZB', (M, M, M, M, M, M), spp=invert(checkerboard((M, M, M, M, M, M))), memoryLayoutClass=PatternMemoryLayout)
+  ZC = Tensor('ZC', (M, M, M, M, M, M))
+
+  class Counter:
+    def __init__(self):
+      self.counter = 0
+
+  counter = Counter()
+
+  def _(kernel):
+    counter.counter += 1
+    g.add(f'kernel{counter.counter}', kernel)
+
+  _(C['ab'] <= A['a'])
+  _(C['ab'] <= A['a'] + B['b'])
+  _(C['ab'] <= A['a'] * B['b'])
+  _(C2['abc'] <= A['a'] + B['b'])
+  _(C2['abc'] <= B4['ac'])
+  _(C['ij'] <= B2['ik'] * B3['kj'])
+  _(C['ij'] <= B4['ik'] * B3['kj'])
+  _(C['ij'] <= B4['ik'] * B3['kj'] * B4['ij'])
+  _(C2['zij'] <= B2['ik'] * Z['kjz'])
+  _(C0['k'] <= B4['ik'] * A['i'])
+
+  _(C4['ijXY'] <= YA['Zik'] * YB['kXY'] * YC['Zkj'])
+  _(ZC['abcxyz'] <= ZA['abcijk'] * ZB['ijkxyz'])
+  _(ZC['abcxyz'] <= ZA['aibjck'] * ZB['zkyjxi'])