[ROCM] Update bounds for large f16 data-tiling ukernel (#22481)

Signed-off-by: Jorn Tuyls <[email protected]>

Author: jtuyls

compiler/plugins/target/ROCM/Dialect/ROCM/IR/ROCMAttrs.cpp

@@ -81,12 +81,21 @@ static bool checkIterationSizeConstraints(ArrayRef<int64_t> iterationSizes,
     if (indexVal < 0 || indexVal >= iterationSizes.size()) {
       return false;
     }
+    // For now, assume a dynamic dimension is very large and any division
+    // constraint is satisfied to keep the performance state on current models
+    // (llama) as is.
+    // TODO(#22370): This is not ideal and can be improved once we support value
+    // bounds on dynamic dimensions for encodings.
     if (IntegerAttr sizeMin = constraint.getSizeMin()) {
-      if (iterationSizes[indexVal] < sizeMin.getInt()) {
+      if (ShapedType::isStatic(iterationSizes[indexVal]) &&
+          iterationSizes[indexVal] < sizeMin.getInt()) {
         return false;
       }
     }
     if (IntegerAttr sizeMax = constraint.getSizeMax()) {
+      if (ShapedType::isDynamic(iterationSizes[indexVal])) {
+        return false;
+      }
       if (iterationSizes[indexVal] > sizeMax.getInt()) {
         return false;
       }
@@ -95,6 +104,9 @@ static bool checkIterationSizeConstraints(ArrayRef<int64_t> iterationSizes,
       if (sizeDiv.getInt() <= 0) {
         return false;
       }
+      if (ShapedType::isDynamic(iterationSizes[indexVal])) {
+        return true;
+      }
       if (iterationSizes[indexVal] % sizeDiv.getInt()) {
         return false;
       }

compiler/plugins/target/ROCM/builtins/mlir_ukernel/iree_uk_amdgpu_dt_matmul_f16.mlir

@@ -17,7 +17,24 @@
 util.func @pingpong_dt_large_f16(%lhs_base: !lhs_base_ty, %rhs_base: !rhs_base_ty, %unused_acc: !acc_base_ty) -> !acc_base_ty attributes {
   ukernel_info = #rocm.ukernel_info<
     match = {
-      types = [f16, f16, f32]
+      types = [f16, f16, f32],
+      iteration_sizes_constraints = [
+        #rocm.ukernel_interation_size_constraint<
+          index = 0,
+          size_min = 512,
+          size_div = 64
+        >,
+        #rocm.ukernel_interation_size_constraint<
+          index = 1,
+          size_min = 32832,
+          size_div = 64
+        >,
+        #rocm.ukernel_interation_size_constraint<
+          index = 2,
+          size_min = 512,
+          size_div = 64
+        >
+      ]
     },
     mma = #iree_gpu.data_tiled_mma_layout<
       intrinsic = MFMA_F32_16x16x16_F16,