Allow build_CS_3D_BLOCK_RC to also have composite sub-schemes.

Don't pad transpose kernels with more than 3 lengths Co-authored-by: Steve Leung <[email protected]>
ROCm · Jun 24, 2022 · ba13730 · ba13730
1 parent 878f63e
commit ba13730
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 9 deletions.
diff --git a/clients/tests/accuracy_test_adhoc.cpp b/clients/tests/accuracy_test_adhoc.cpp
@@ -36,6 +36,9 @@ std::vector<std::vector<size_t>> adhoc_sizes = {
     // SBRC 192 with special param
     {192, 192, 192},
     {192, 84, 84},
+
+    // Failure with build_CS_3D_BLOCK_RC
+    {680, 128, 128},
 };
 
 const static std::vector<std::vector<size_t>> stride_range = {{1}};

diff --git a/library/src/assignment_policy.cpp b/library/src/assignment_policy.cpp
@@ -1103,6 +1103,14 @@ void AssignmentPolicy::PadPlan(ExecPlan& execPlan)
                 // SBCR plans combine higher dimensions in ways that confuse padding
                 if(u.node.scheme == CS_KERNEL_STOCKHAM_BLOCK_CR)
                     return;
+                // transpose kernels don't handle arbitrary strides,
+                // and with 4 or more lengths either choice of
+                // padding dim will trigger incorrect behaviour
+                if((u.node.scheme == CS_KERNEL_TRANSPOSE
+                    || u.node.scheme == CS_KERNEL_TRANSPOSE_XY_Z
+                    || u.node.scheme == CS_KERNEL_TRANSPOSE_Z_XY)
+                   && u.node.length.size() > 3)
+                    return;
             }
 
             // Ensure that if we're forced to pad along one dimension

diff --git a/library/src/tree_node_3D.cpp b/library/src/tree_node_3D.cpp
@@ -473,13 +473,6 @@ void BLOCKRC3DNode::AssignParams_internal()
             node->oDist = node->outStride[2] * node->length[0];
             break;
         }
-        case CS_KERNEL_STOCKHAM:
-        {
-            node->outStride = node->inStride;
-            node->oDist     = node->iDist;
-            node->AssignParams();
-            break;
-        }
         case CS_KERNEL_TRANSPOSE_XY_Z:
         {
             node->outStride.push_back(1);
@@ -497,8 +490,12 @@ void BLOCKRC3DNode::AssignParams_internal()
             break;
         }
         default:
-            // build_CS_3D_BLOCK_RC should not have created any other node types
-            throw std::runtime_error("Scheme Assertion Failed, unexpected node scheme.");
+        {
+            node->outStride = node->inStride;
+            node->oDist     = node->iDist;
+            node->AssignParams();
+            break;
+        }
         }
         prev_outStride = node->outStride;
         prev_oDist     = node->oDist;