[BugFix] Fix ConvDirectNaiveConvFwd not applicable after #3213 (#3284)

ROCm · Oct 1, 2024 · c233b80 · c233b80
1 parent 53a0b11
commit c233b80
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/src/conv/problem_description.cpp b/src/conv/problem_description.cpp
@@ -119,6 +119,24 @@ std::string ProblemDescription::GetAlphaBetaCaseStr() const
     }
 }
 
+void ProblemDescription::HeuristicUpdateLayouts()
+{
+    static const std::vector<std::string> supported_layouts = {"NCHW", "NHWC", "CHWN", "NCDHW"};
+
+    for(const std::string& layout : supported_layouts)
+    {
+        if(in.IsPossibleLayout4D5D(layout) && out.IsPossibleLayout4D5D(layout) &&
+           weights.IsPossibleLayout4D5D(layout))
+        {
+            in_layout      = layout;
+            weights_layout = layout;
+            out_layout     = layout;
+            return;
+        }
+    }
+    // If we did not find consistent layout, leave them as-is
+}
+
 void ProblemDescription::MakeNetworkConfig(std::string& conf_key) const
 {
     std::ostringstream ss;

diff --git a/src/include/miopen/conv/problem_description.hpp b/src/include/miopen/conv/problem_description.hpp
@@ -166,6 +166,7 @@ struct MIOPEN_INTERNALS_EXPORT ProblemDescription : ProblemDescriptionBase
           beta(beta_),
           alpha_beta_case(ClassifyAlphaBeta(alpha, beta))
     {
+        HeuristicUpdateLayouts();
     }
 
     // Conv descriptor getters
@@ -368,6 +369,8 @@ struct MIOPEN_INTERNALS_EXPORT ProblemDescription : ProblemDescriptionBase
                out.AllLengthsFitIntoInt();
     }
 
+    void HeuristicUpdateLayouts();
+
     void MakeNetworkConfig(std::string& conf_key) const;
 
     NetworkConfig MakeNetworkConfig() const override