[coco] Bag-level weight allocation (#1932)

This commit introduces a new allocate API in PlainWieghtContext<T> which
allows users to allocate a weight for each bag.

Signed-off-by: Jonghyun Park <jh1302.park@samsung.com>

diff --git a/contrib/coco/generic/include/coco/IR/PlainWeightContext.h b/contrib/coco/generic/include/coco/IR/PlainWeightContext.h
index 062cbdb..be50c39 100644 (file)
--- a/contrib/coco/generic/include/coco/IR/PlainWeightContext.h
+++ b/contrib/coco/generic/include/coco/IR/PlainWeightContext.h
@@ -37,6 +37,19 @@ template <typename T> struct PlainWeightContext
 {
   virtual ~PlainWeightContext() = default;
 
+  /**
+   * @brief Allocate a weight space for a given blob
+   *
+   * @require the following code SHOULD work for any bag "b":
+   *   PlainWeightContext<T> ctx;
+   *
+   *   ctx.allocate(b);
+   *   auto span = ctx.weight(b);
+   *   assert(span.data() != nullptr);
+   *   assert(span.size() == bag->size());
+   */
+  virtual void allocate(const Bag *) = 0;
+  // WARN Depercated
   virtual void allocate(const KernelObject *) = 0;
 
   /**

diff --git a/contrib/coco/generic/src/IR/Data.cpp b/contrib/coco/generic/src/IR/Data.cpp
index 34e57e3..ec0a297 100644 (file)
--- a/contrib/coco/generic/src/IR/Data.cpp
+++ b/contrib/coco/generic/src/IR/Data.cpp
@@ -120,6 +120,12 @@ public:
   PlainWeightContextImpl(PlainWeightContextImpl &&) = delete;
 
 public:
+  void allocate(const coco::Bag *bag) override
+  {
+    assert(bag != nullptr);
+    _blob->allocate(bag, sizeof(T));
+  }
+
   void allocate(const coco::KernelObject *obj)
   {
     assert(obj != nullptr);

diff --git a/contrib/coco/generic/src/IR/Data.test.cpp b/contrib/coco/generic/src/IR/Data.test.cpp
index 5c8b857..cd3bea6 100644 (file)
--- a/contrib/coco/generic/src/IR/Data.test.cpp
+++ b/contrib/coco/generic/src/IR/Data.test.cpp
@@ -33,6 +33,34 @@ TEST(IR_DATA, construct)
   ASSERT_EQ(mutable_ptr->f32(), immutable_ptr->f32());
 }
 
+TEST(IR_DATA, allocate_and_link_bag)
+{
+  auto m = coco::Module::create();
+  auto d = coco::Data::create();
+
+  // Create a bag
+  auto bag = m->entity()->bag()->create(9);
+
+  // weight(...) SHOULD return a null-span for an invalid bag
+  {
+    auto span = d->f32()->weight(bag);
+
+    ASSERT_EQ(span.data(), nullptr);
+    ASSERT_EQ(span.size(), 0);
+  }
+
+  // Allocate a weight space
+  d->f32()->allocate(bag);
+
+  // weight(...) SHOULD return a valid for a valid bag
+  {
+    auto span = d->f32()->weight(bag);
+
+    ASSERT_NE(span.data(), nullptr);
+    ASSERT_EQ(span.size(), bag->size());
+  }
+}
+
 TEST(IR_DATA, allocate_and_link_kernel)
 {
   using nncc::core::ADT::kernel::num_elements;