#include "nest/VarContext.h"
#include "nest/DomainContext.h"
+#include "nest/Ret.h"
#include "nest/Block.h"
namespace nest
public:
void push(const Expr &expr);
+
+private:
+ std::shared_ptr<Ret> _ret;
+
+public:
+ // NOTE Do NOT invoke ret() before ret(expr) call
+ const Ret &ret(void) const;
+
+public:
+ // NOTE Only one ret(expr) call is allowed for each module
+ void ret(const Closure &closure);
};
} // namespace nest
#include "nest/Module.h"
+#include <cassert>
+
namespace nest
{
_block.append(stmt);
}
+void Module::ret(const Closure &clo)
+{
+ // Only one RET is allowed for each module
+ assert(_ret == nullptr);
+ _ret = std::make_shared<Ret>(clo.id(), clo.sub());
+}
+
+const Ret &Module::ret(void) const
+{
+ // Caller should NOT invoke this method before setting ret
+ assert(_ret != nullptr);
+ return *_ret;
+}
+
} // namespace nest
ASSERT_EQ(m.block().size(), 1);
ASSERT_NE(m.block().at(0)->asPush(), nullptr);
}
+
+TEST(MODULE, ret)
+{
+ nest::Module m;
+
+ auto ifm = m.domain().make({1});
+ auto ofm = m.domain().make({1});
+
+ auto ind = m.var().make();
+
+ m.push(ifm(ind));
+ m.ret(ofm(ind));
+
+ ASSERT_EQ(m.ret().id(), ofm.id());
+ ASSERT_EQ(m.ret().sub().rank(), 1);
+}