Add tests to increase code coverage rate for Constraint API (#706)

* Add tests to increase code coverage rate for Constraint API

* Update test script based on review comments

* Update code based on review comments

Author: mona-nv

tests/cpp/integration/image_constraints.cc

@@ -328,18 +328,31 @@ void test_invalid()
   auto runtime = legate::Runtime::get_runtime();
   auto context = runtime->find_library(library_name);
 
-  auto func  = runtime->create_store(legate::Shape{10, 10}, legate::int32());
-  auto range = runtime->create_store(legate::Shape{10, 10}, legate::int64());
+  auto create_task = [&](auto func, auto range) {
+    auto task = runtime->create_task(context, static_cast<std::int64_t>(IMAGE_TESTER) + func.dim());
+    auto part_domain = task.declare_partition();
+    auto part_range  = task.declare_partition();
 
-  auto task        = runtime->create_task(context, static_cast<std::int64_t>(IMAGE_TESTER) + 1);
-  auto part_domain = task.declare_partition();
-  auto part_range  = task.declare_partition();
+    task.add_input(func, part_domain);
+    task.add_input(range, part_range);
+    task.add_constraint(legate::image(part_domain, part_range));
 
-  task.add_input(func, part_domain);
-  task.add_input(range, part_range);
-  task.add_constraint(legate::image(part_domain, part_range));
+    return task;
+  };
 
-  EXPECT_THROW(runtime->submit(std::move(task)), std::invalid_argument);
+  {
+    auto func1  = runtime->create_store(legate::Shape{10, 10}, legate::int32());
+    auto range1 = runtime->create_store(legate::Shape{10, 10}, legate::int64());
+    auto task   = create_task(func1, range1);
+    EXPECT_THROW(runtime->submit(std::move(task)), std::invalid_argument);
+  }
+
+  {
+    auto func2  = runtime->create_store(legate::Shape{4, 4}, legate::point_type(2));
+    auto range2 = runtime->create_store(legate::Shape{10}, legate::int64());
+    auto task   = create_task(func2, range2.promote(1, 1));
+    EXPECT_THROW(runtime->submit(std::move(task)), std::runtime_error);
+  }
 }
 
 TEST_P(Valid, 1D)

tests/cpp/unit/constraint.cc

@@ -17,18 +17,11 @@
 
 #include <gtest/gtest.h>
 
-namespace unit {
-
-using Variable        = DefaultFixture;
-using Alignment       = DefaultFixture;
-using Broadcast       = DefaultFixture;
-using ImageConstraint = DefaultFixture;
-
 namespace {
 
-constexpr const char library_name[] = "test_constraints";
+// NOLINTBEGIN(readability-magic-numbers)
 
-}  // namespace
+constexpr const char library_name[] = "test_constraints";
 
 // Dummy task to make the runtime think the store is initialized
 struct Initializer : public legate::LegateTask<Initializer> {
@@ -37,39 +30,36 @@ struct Initializer : public legate::LegateTask<Initializer> {
   static void cpu_variant(legate::TaskContext /*context*/) {}
 };
 
-void register_tasks()
-{
-  static bool prepared = false;
-  if (prepared) {
-    return;
+class Constraint : public DefaultFixture {
+ public:
+  void SetUp() override
+  {
+    DefaultFixture::SetUp();
+    auto runtime = legate::Runtime::get_runtime();
+    auto context = runtime->create_library(library_name);
+    Initializer::register_variants(context);
   }
-  prepared     = true;
-  auto runtime = legate::Runtime::get_runtime();
-  auto context = runtime->create_library(library_name);
-  Initializer::register_variants(context);
-}
+};
 
-TEST_F(Variable, BasicMethods)
+TEST_F(Constraint, Variable)
 {
-  register_tasks();
-
   auto runtime = legate::Runtime::get_runtime();
   auto context = runtime->find_library(library_name);
   auto task    = runtime->create_task(context, Initializer::TASK_ID);
 
   // Test basic properties
   auto part     = task.declare_partition();
   auto part_imp = part.impl();
-  EXPECT_FALSE(part_imp->closed());
-  EXPECT_EQ(part_imp->kind(), legate::detail::Expr::Kind::VARIABLE);
-  EXPECT_EQ(part_imp->as_literal(), nullptr);
-  EXPECT_EQ(part_imp->as_variable(), part_imp);
-  EXPECT_TRUE(part_imp->operation() != nullptr);
+  ASSERT_FALSE(part_imp->closed());
+  ASSERT_EQ(part_imp->kind(), legate::detail::Expr::Kind::VARIABLE);
+  ASSERT_EQ(part_imp->as_literal(), nullptr);
+  ASSERT_EQ(part_imp->as_variable(), part_imp);
+  ASSERT_TRUE(part_imp->operation() != nullptr);
 
   // Test equal
   auto part1(part);
   auto part1_imp = part1.impl();
-  EXPECT_EQ(*part_imp, *part1_imp);
+  ASSERT_EQ(*part_imp, *part1_imp);
   auto part2     = task.declare_partition();
   auto part2_imp = part2.impl();
 
@@ -78,16 +68,14 @@ TEST_F(Variable, BasicMethods)
   part_imp->find_partition_symbols(symbols);
   part1_imp->find_partition_symbols(symbols);
   part2_imp->find_partition_symbols(symbols);
-  EXPECT_EQ(symbols.size(), 3);
-  EXPECT_TRUE(std::find(symbols.begin(), symbols.end(), part_imp) != symbols.end());
-  EXPECT_TRUE(std::find(symbols.begin(), symbols.end(), part1_imp) != symbols.end());
-  EXPECT_TRUE(std::find(symbols.begin(), symbols.end(), part2_imp) != symbols.end());
+  ASSERT_EQ(symbols.size(), 3);
+  ASSERT_TRUE(std::find(symbols.begin(), symbols.end(), part_imp) != symbols.end());
+  ASSERT_TRUE(std::find(symbols.begin(), symbols.end(), part1_imp) != symbols.end());
+  ASSERT_TRUE(std::find(symbols.begin(), symbols.end(), part2_imp) != symbols.end());
 }
 
-TEST_F(Alignment, BasicMethods)
+TEST_F(Constraint, Alignment)
 {
-  register_tasks();
-
   auto runtime = legate::Runtime::get_runtime();
   auto context = runtime->find_library(library_name);
   auto task    = runtime->create_task(context, Initializer::TASK_ID);
@@ -96,51 +84,51 @@ TEST_F(Alignment, BasicMethods)
   auto part2 = task.declare_partition();
 
   auto aligment = legate::detail::align(part1.impl(), part2.impl());
-  EXPECT_EQ(aligment->kind(), legate::detail::Constraint::Kind::ALIGNMENT);
-  EXPECT_EQ(aligment->lhs(), part1.impl());
-  EXPECT_EQ(aligment->rhs(), part2.impl());
-  EXPECT_EQ(aligment->as_alignment(), aligment.get());
-  EXPECT_EQ(aligment->as_broadcast(), nullptr);
-  EXPECT_EQ(aligment->as_image_constraint(), nullptr);
-  EXPECT_FALSE(aligment->is_trivial());
+  ASSERT_EQ(aligment->kind(), legate::detail::Constraint::Kind::ALIGNMENT);
+  ASSERT_EQ(aligment->lhs(), part1.impl());
+  ASSERT_EQ(aligment->rhs(), part2.impl());
+  ASSERT_EQ(aligment->as_alignment(), aligment.get());
+  ASSERT_EQ(aligment->as_broadcast(), nullptr);
+  ASSERT_EQ(aligment->as_image_constraint(), nullptr);
+  ASSERT_EQ(aligment->as_scale_constraint(), nullptr);
+  ASSERT_EQ(aligment->as_bloat_constraint(), nullptr);
+  ASSERT_FALSE(aligment->is_trivial());
 
   // Test find_partition_symbols
   std::vector<const legate::detail::Variable*> symbols = {};
   aligment->find_partition_symbols(symbols);
-  EXPECT_EQ(symbols.size(), 2);
-  EXPECT_TRUE(std::find(symbols.begin(), symbols.end(), part1.impl()) != symbols.end());
-  EXPECT_TRUE(std::find(symbols.begin(), symbols.end(), part2.impl()) != symbols.end());
+  ASSERT_EQ(symbols.size(), 2);
+  ASSERT_TRUE(std::find(symbols.begin(), symbols.end(), part1.impl()) != symbols.end());
+  ASSERT_TRUE(std::find(symbols.begin(), symbols.end(), part2.impl()) != symbols.end());
 }
 
-TEST_F(Broadcast, BasicMethods)
+TEST_F(Constraint, Broadcast)
 {
-  register_tasks();
-
   auto runtime = legate::Runtime::get_runtime();
   auto context = runtime->find_library(library_name);
   auto task    = runtime->create_task(context, Initializer::TASK_ID);
   auto part1   = task.declare_partition();
 
   auto dims      = legate::from_range<std::uint32_t>(3);
   auto broadcast = legate::detail::broadcast(part1.impl(), dims);
-  EXPECT_EQ(broadcast->kind(), legate::detail::Constraint::Kind::BROADCAST);
-  EXPECT_EQ(broadcast->variable(), part1.impl());
-  EXPECT_EQ(broadcast->axes(), dims);
-  EXPECT_EQ(broadcast->as_alignment(), nullptr);
-  EXPECT_EQ(broadcast->as_broadcast(), broadcast.get());
-  EXPECT_EQ(broadcast->as_image_constraint(), nullptr);
+  ASSERT_EQ(broadcast->kind(), legate::detail::Constraint::Kind::BROADCAST);
+  ASSERT_EQ(broadcast->variable(), part1.impl());
+  ASSERT_EQ(broadcast->axes(), dims);
+  ASSERT_EQ(broadcast->as_alignment(), nullptr);
+  ASSERT_EQ(broadcast->as_broadcast(), broadcast.get());
+  ASSERT_EQ(broadcast->as_image_constraint(), nullptr);
+  ASSERT_EQ(broadcast->as_scale_constraint(), nullptr);
+  ASSERT_EQ(broadcast->as_bloat_constraint(), nullptr);
 
   // Test find_partition_symbols
   std::vector<const legate::detail::Variable*> symbols = {};
   broadcast->find_partition_symbols(symbols);
-  EXPECT_EQ(symbols.size(), 1);
-  EXPECT_TRUE(std::find(symbols.begin(), symbols.end(), part1.impl()) != symbols.end());
+  ASSERT_EQ(symbols.size(), 1);
+  ASSERT_TRUE(std::find(symbols.begin(), symbols.end(), part1.impl()) != symbols.end());
 }
 
-TEST_F(ImageConstraint, BasicMethods)
+TEST_F(Constraint, ImageConstraint)
 {
-  register_tasks();
-
   auto runtime    = legate::Runtime::get_runtime();
   auto context    = runtime->find_library(library_name);
   auto task       = runtime->create_task(context, Initializer::TASK_ID);
@@ -149,18 +137,80 @@ TEST_F(ImageConstraint, BasicMethods)
 
   auto image_constraint = legate::detail::image(
     part_func.impl(), part_range.impl(), legate::ImageComputationHint::NO_HINT);
-  EXPECT_EQ(image_constraint->kind(), legate::detail::Constraint::Kind::IMAGE);
-  EXPECT_EQ(image_constraint->var_function(), part_func.impl());
-  EXPECT_EQ(image_constraint->var_range(), part_range.impl());
-  EXPECT_EQ(image_constraint->as_alignment(), nullptr);
-  EXPECT_EQ(image_constraint->as_broadcast(), nullptr);
-  EXPECT_EQ(image_constraint->as_image_constraint(), image_constraint.get());
+  ASSERT_EQ(image_constraint->kind(), legate::detail::Constraint::Kind::IMAGE);
+  ASSERT_EQ(image_constraint->var_function(), part_func.impl());
+  ASSERT_EQ(image_constraint->var_range(), part_range.impl());
+  ASSERT_EQ(image_constraint->as_alignment(), nullptr);
+  ASSERT_EQ(image_constraint->as_broadcast(), nullptr);
+  ASSERT_EQ(image_constraint->as_image_constraint(), image_constraint.get());
+  ASSERT_EQ(image_constraint->as_scale_constraint(), nullptr);
+  ASSERT_EQ(image_constraint->as_bloat_constraint(), nullptr);
 
   // Test find_partition_symbols
   std::vector<const legate::detail::Variable*> symbols = {};
   image_constraint->find_partition_symbols(symbols);
-  EXPECT_EQ(symbols.size(), 2);
-  EXPECT_TRUE(std::find(symbols.begin(), symbols.end(), part_func.impl()) != symbols.end());
-  EXPECT_TRUE(std::find(symbols.begin(), symbols.end(), part_range.impl()) != symbols.end());
+  ASSERT_EQ(symbols.size(), 2);
+  ASSERT_TRUE(std::find(symbols.begin(), symbols.end(), part_func.impl()) != symbols.end());
+  ASSERT_TRUE(std::find(symbols.begin(), symbols.end(), part_range.impl()) != symbols.end());
+}
+
+TEST_F(Constraint, ScaleConstraint)
+{
+  auto runtime      = legate::Runtime::get_runtime();
+  auto context      = runtime->find_library(library_name);
+  auto task         = runtime->create_task(context, Initializer::TASK_ID);
+  auto smaller      = runtime->create_store({3}, legate::int64());
+  auto bigger       = runtime->create_store({5}, legate::int64());
+  auto part_smaller = task.add_output(smaller);
+  auto part_bigger  = task.add_output(bigger);
+
+  auto scale_constraint = legate::detail::scale({1}, part_smaller.impl(), part_bigger.impl());
+  ASSERT_EQ(scale_constraint->kind(), legate::detail::Constraint::Kind::SCALE);
+  ASSERT_EQ(scale_constraint->var_smaller(), part_smaller.impl());
+  ASSERT_EQ(scale_constraint->var_bigger(), part_bigger.impl());
+  ASSERT_EQ(scale_constraint->as_alignment(), nullptr);
+  ASSERT_EQ(scale_constraint->as_broadcast(), nullptr);
+  ASSERT_EQ(scale_constraint->as_image_constraint(), nullptr);
+  ASSERT_EQ(scale_constraint->as_scale_constraint(), scale_constraint.get());
+  ASSERT_EQ(scale_constraint->as_bloat_constraint(), nullptr);
+
+  // Test find_partition_symbols
+  std::vector<const legate::detail::Variable*> symbols = {};
+  scale_constraint->find_partition_symbols(symbols);
+  ASSERT_EQ(symbols.size(), 2);
+  ASSERT_TRUE(std::find(symbols.begin(), symbols.end(), part_smaller.impl()) != symbols.end());
+  ASSERT_TRUE(std::find(symbols.begin(), symbols.end(), part_bigger.impl()) != symbols.end());
 }
-}  // namespace unit
+
+TEST_F(Constraint, BloatConstraint)
+{
+  auto runtime = legate::Runtime::get_runtime();
+  auto context = runtime->find_library(library_name);
+  auto task    = runtime->create_task(context, Initializer::TASK_ID);
+  auto source  = runtime->create_store({5}, legate::int64());
+  auto bloated = runtime->create_store({5}, legate::int64());
+  runtime->issue_fill(source, legate::Scalar(std::int64_t{0}));
+  runtime->issue_fill(bloated, legate::Scalar(std::int64_t{0}));
+  auto part_source  = task.add_input(source);
+  auto part_bloated = task.add_input(bloated);
+
+  auto bloat_constraint = legate::detail::bloat(part_source.impl(), part_bloated.impl(), {1}, {3});
+  ASSERT_EQ(bloat_constraint->kind(), legate::detail::Constraint::Kind::BLOAT);
+  ASSERT_EQ(bloat_constraint->var_source(), part_source.impl());
+  ASSERT_EQ(bloat_constraint->var_bloat(), part_bloated.impl());
+  ASSERT_EQ(bloat_constraint->as_alignment(), nullptr);
+  ASSERT_EQ(bloat_constraint->as_broadcast(), nullptr);
+  ASSERT_EQ(bloat_constraint->as_image_constraint(), nullptr);
+  ASSERT_EQ(bloat_constraint->as_scale_constraint(), nullptr);
+  ASSERT_EQ(bloat_constraint->as_bloat_constraint(), bloat_constraint.get());
+
+  // Test find_partition_symbols
+  std::vector<const legate::detail::Variable*> symbols = {};
+  bloat_constraint->find_partition_symbols(symbols);
+  ASSERT_EQ(symbols.size(), 2);
+  ASSERT_TRUE(std::find(symbols.begin(), symbols.end(), part_source.impl()) != symbols.end());
+  ASSERT_TRUE(std::find(symbols.begin(), symbols.end(), part_bloated.impl()) != symbols.end());
+}
+
+// NOLINTEND(readability-magic-numbers)
+}  // namespace