Add parsing of min (>=) constraints for mapspace.

NVlabs · Oct 5, 2023 · f1daa2e · f1daa2e
1 parent ba2afa8
commit f1daa2e
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 0 deletions.
diff --git a/include/mapping/constraints.hpp b/include/mapping/constraints.hpp
@@ -52,6 +52,7 @@ class Constraints
   // The constraints.
   std::map<unsigned, std::map<problem::Shape::FlattenedDimensionID, int>> factors_;
   std::map<unsigned, std::map<problem::Shape::FlattenedDimensionID, int>> max_factors_;
+  std::map<unsigned, std::map<problem::Shape::FlattenedDimensionID, int>> min_factors_;
   std::map<unsigned, std::pair<std::vector<problem::Shape::FlattenedDimensionID>,
                                std::vector<problem::Shape::FlattenedDimensionID>>> permutations_;
   std::map<unsigned, std::uint32_t> spatial_splits_;
@@ -73,6 +74,7 @@ class Constraints
 
   const std::map<unsigned, std::map<problem::Shape::FlattenedDimensionID, int>>& Factors() const;
   const std::map<unsigned, std::map<problem::Shape::FlattenedDimensionID, int>>& MaxFactors() const;
+  const std::map<unsigned, std::map<problem::Shape::FlattenedDimensionID, int>>& MinFactors() const;
   const std::map<unsigned, std::pair<std::vector<problem::Shape::FlattenedDimensionID>,
                                      std::vector<problem::Shape::FlattenedDimensionID>>>& Permutations() const;
   const std::map<unsigned, std::uint32_t>& SpatialSplits() const;  
@@ -113,6 +115,7 @@ class Constraints
   // Parsers.
   std::map<problem::Shape::FlattenedDimensionID, int> ParseFactors(config::CompoundConfigNode constraint);
   std::map<problem::Shape::FlattenedDimensionID, int> ParseMaxFactors(config::CompoundConfigNode constraint);
+  std::map<problem::Shape::FlattenedDimensionID, int> ParseMinFactors(config::CompoundConfigNode constraint);
   std::pair<std::vector<problem::Shape::FlattenedDimensionID>,
             std::vector<problem::Shape::FlattenedDimensionID>> ParsePermutations(config::CompoundConfigNode constraint);
   void ParseDatatypeBypassSettings(config::CompoundConfigNode constraint, unsigned level);

diff --git a/src/mapping/constraints.cpp b/src/mapping/constraints.cpp
@@ -83,6 +83,12 @@ const std::map<unsigned, std::map<problem::Shape::FlattenedDimensionID, int>>&
   return max_factors_;
 }
 
+const std::map<unsigned, std::map<problem::Shape::FlattenedDimensionID, int>>&
+  Constraints::MinFactors() const
+{
+  return min_factors_;
+}
+
 const std::map<unsigned, std::uint32_t>& 
   Constraints::MaxRemainders() const
 {
@@ -623,6 +629,21 @@ void Constraints::ParseSingleConstraint(
       max_factors_[level_id][max_factor.first] = max_factor.second;
     }
 
+    auto level_min_factors = ParseMinFactors(attributes);
+    for (auto& min_factor: level_min_factors)
+    {
+      if (min_factors_[level_id].find(min_factor.first) != min_factors_[level_id].end())
+      {
+        std::cerr << "ERROR: re-specification of min factor for dimension "
+                  << problem::GetShape()->FlattenedDimensionIDToName.at(min_factor.first)
+                  << " at level " << arch_props_.TilingLevelName(level_id)
+                  << ". This may imply a conflict between architecture and "
+                  << "mapspace constraints." << std::endl;
+        exit(1);
+      }
+      min_factors_[level_id][min_factor.first] = min_factor.second;
+    }
+
     auto level_permutations = ParsePermutations(attributes);
     if (level_permutations.first.size() > 0 || level_permutations.second.size() > 0)
     {
@@ -1021,6 +1042,65 @@ Constraints::ParseMaxFactors(config::CompoundConfigNode constraint)
   return retval;
 }
 
+//
+// Parse user min factors.
+//
+std::map<problem::Shape::FlattenedDimensionID, int>
+Constraints::ParseMinFactors(config::CompoundConfigNode constraint)
+{
+  std::map<problem::Shape::FlattenedDimensionID, int> retval;
+
+  std::string buffer;
+  if (constraint.lookupValue("factors", buffer))
+  {
+    buffer = buffer.substr(0, buffer.find("#"));
+
+    std::regex re("([A-Za-z]+)[[:space:]]*>=[[:space:]]*([0-9]+)", std::regex::extended);
+    std::smatch sm;
+    std::string str = std::string(buffer);
+
+    while (std::regex_search(str, sm, re))
+    {
+      std::string dimension_name = sm[1];
+      problem::Shape::FlattenedDimensionID dimension;
+      try
+      {
+        dimension = problem::GetShape()->FlattenedDimensionNameToID.at(dimension_name);
+      }
+      catch (const std::out_of_range& oor)
+      {
+        std::cerr << "ERROR: parsing factors: " << buffer << ": dimension " << dimension_name
+                  << " not found in problem shape." << std::endl;
+        exit(1);
+      }
+
+      int min = std::stoi(sm[2]);
+      if (min <= 0)
+      {
+        std::cerr << "ERROR: min factor must be positive in constraint: " << buffer << std::endl;
+        exit(1);
+      }
+
+      // Found all the information we need to setup a factor!
+      retval[dimension] = min;
+
+      str = sm.suffix().str();
+    }
+  }
+  if (constraint.lookupValue("default_min_factor", buffer))
+  {
+      int min = std::stoi(buffer);
+      for(auto& it : problem::GetShape()->FlattenedDimensionNameToID)
+      {
+        if(retval.find(it.second) == retval.end())
+        {
+          retval[it.second] = min;
+        }
+      }
+  }
+  return retval;
+}
+
 //
 // Parse user permutations.
 //