update isimprovementoverrocm

ROCm · Aug 21, 2024 · fdc4232 · fdc4232
1 parent 86149f6
commit fdc4232
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/src/solver/multimarginloss/forward_reduced_multimarginloss.cpp b/src/solver/multimarginloss/forward_reduced_multimarginloss.cpp
@@ -46,13 +46,29 @@ bool MultiMarginLossForward::IsImprovementOverROCm(
     const ExecutionContext& /*context*/,
     const miopen::multimarginloss::ForwardProblemDescription& problem) const
 {
-    if(problem.GetiDesc().GetLengths()[1] <= 30)
-        return true;
-    if((problem.GetiDesc().GetType() == miopenHalf ||
-        problem.GetiDesc().GetType() == miopenBFloat16) &&
-       problem.GetiDesc().IsContiguous() && problem.GetiDesc().GetLengths()[1] <= 40)
-        return true;
-    return false;
+    int C = problem.GetiDesc().GetLengths()[1];
+    if(problem.allContiguousTensor())
+    {
+        switch(problem.GetiDesc().GetType())
+        {
+        case miopenFloat: return C <= 33;
+        case miopenHalf: return C <= 43;
+        case miopenBFloat16: return C <= 44;
+        // Have not tested with other types yet
+        default: return true;
+        }
+    }
+    else
+    {
+        switch(problem.GetiDesc().GetType())
+        {
+        case miopenFloat: return C <= 31;
+        case miopenHalf: return C <= 38;
+        case miopenBFloat16: return C <= 40;
+        // Have not tested with other types yet
+        default: return true;
+        }
+    }
 }
 
 bool MultiMarginLossForward::IsApplicable(

diff --git a/src/solver/multimarginloss/forward_unreduced_multimarginloss.cpp b/src/solver/multimarginloss/forward_unreduced_multimarginloss.cpp
@@ -50,7 +50,7 @@ bool MultiMarginLossUnreducedForward::IsImprovementOverROCm(
     {
         switch(problem.GetiDesc().GetType())
         {
-        case miopenFloat: return C <= 32;
+        case miopenFloat: return C <= 33;
         case miopenHalf: return C <= 43;
         case miopenBFloat16: return C <= 44;
         // Have not tested with other types yet