[AIROCMLIR-708] Fix validity masking for padded transform indicies (#2353)

* fix validity masking for padded transform indicies

Signed-off-by: bogdan-petkovic <bogdan.petkovic@htecgroup.com>

* refactor validity check and update test IR expectations

Signed-off-by: bogdan-petkovic <bogdan.petkovic@htecgroup.com>

* move zeroConst inside addLowerDimBoundsCheck lambda

Signed-off-by: bogdan-petkovic <bogdan.petkovic@htecgroup.com>

Author: bogdan-petkovic

mlir/lib/Dialect/Rock/utility/transformMapUtils.cpp

@@ -1077,15 +1077,19 @@ Value mlir::rock::updateValidityAfter(OpBuilder &b, Location loc,
   Value isValid =
       b.createOrFold<arith::ConstantIntOp>(loc, b.getI1Type(), true);
   ArrayRef<int64_t> lowerBounds = map.getLowerBounds();
-
-  // unsigned < catches both negatives (as all negatives are > the bound)
-  // and being too large on the right.
-  auto addLowerDimUltClamp = [&](uint32_t lowerDim) {
+  // Explicitly check both bounds. Left padding can produce negative indices,
+  // while right padding can produce indices >= bound.
+  auto addLowerDimBoundsCheck = [&](uint32_t lowerDim) {
     int64_t bound = lowerBounds[lowerDim];
+    Value zeroConst = b.createOrFold<arith::ConstantIndexOp>(loc, 0);
     Value boundConst = b.createOrFold<arith::ConstantIndexOp>(loc, bound);
     Value output = outputs[lowerDim];
-    Value inBounds = arith::CmpIOp::create(b, loc, arith::CmpIPredicate::ult,
-                                           output, boundConst);
+    Value geLowerBound = arith::CmpIOp::create(
+        b, loc, arith::CmpIPredicate::sge, output, zeroConst);
+    Value ltUpperBound = arith::CmpIOp::create(
+        b, loc, arith::CmpIPredicate::slt, output, boundConst);
+    Value inBounds = b.createOrFold<arith::AndIOp>(loc, b.getI1Type(),
+                                                   geLowerBound, ltUpperBound);
     isValid =
         b.createOrFold<arith::AndIOp>(loc, b.getI1Type(), inBounds, isValid);
   };
@@ -1102,13 +1106,13 @@ Value mlir::rock::updateValidityAfter(OpBuilder &b, Location loc,
 
         if (params[leftParam] == 0 && params[rightParam] == 0)
           continue;
-        addLowerDimUltClamp(lowerDim);
+        addLowerDimBoundsCheck(lowerDim);
       }
     }
     if (type == TransformType::Embed) {
       if (!embedCanBeInvalid(map, op))
         continue;
-      addLowerDimUltClamp(op.getLowerDims()[0]);
+      addLowerDimBoundsCheck(op.getLowerDims()[0]);
     }
   }
   return isValid;

mlir/test/Dialect/Rock/lowering_transforming_for.mlir

@@ -13,6 +13,10 @@
     by [<Pad{0, 8} ["pad"] at [0] -> ["raw"] at [0]>]
     bounds = [16] -> [8]>
 
+#transform_map_pad_left = #rock.transform_map<affine_map<(d0) -> (d0 - 1)>
+    by [<Pad{1, 7} ["pad"] at [0] -> ["raw"] at [0]>]
+    bounds = [16] -> [8]>
+
 module {
 // CHECK-LABEL: func.func @no_transform_to_affine
 func.func @no_transform_to_affine() {
@@ -294,10 +298,13 @@ func.func @no_loop_loop_result(%arg0: index, %arg1: index) -> index {
 
 
 // CHECK-LABEL: func.func @bounds_check_pad
+// CHECK-DAG: %[[c0:.*]] = arith.constant 0
 // CHECK-DAG: %[[c8:.*]] = arith.constant 8
 // CHECK: affine.for %[[num:.*]] = {{.*}}to 16
-// CHECK: %[[valid:.*]] = arith.cmpi ult, %[[num]], %[[c8]]
-// CHECK: gpu.printf "%d", %1
+// CHECK-DAG: %[[ge:.*]] = arith.cmpi sge, %[[num]], %[[c0]]
+// CHECK-DAG: %[[lt:.*]] = arith.cmpi slt, %[[num]], %[[c8]]
+// CHECK: %[[valid:.*]] = arith.andi %[[ge]], %[[lt]]
+// CHECK: gpu.printf "%d", %{{.*}}
 func.func @bounds_check_pad() {
     %c0 = arith.constant 0 : index
     rock.transforming_for (%arg0) = [#transform_map_pad](%c0) (%arg1) = validity bounds [16] strides [1] {
@@ -306,4 +313,23 @@ func.func @bounds_check_pad() {
     }
     return
 }
+
+// CHECK-LABEL: func.func @bounds_check_pad_left
+// CHECK-DAG: %[[c0:.*]] = arith.constant 0
+// CHECK-DAG: %[[c8:.*]] = arith.constant 8
+// CHECK-DAG: %[[cm1:.*]] = arith.constant -1
+// CHECK: affine.for %[[num:.*]] = {{.*}}to 16
+// CHECK: %[[shifted:.*]] = arith.addi %[[num]], %[[cm1]]
+// CHECK-DAG: %[[ge:.*]] = arith.cmpi sge, %[[shifted]], %[[c0]]
+// CHECK-DAG: %[[lt:.*]] = arith.cmpi slt, %[[shifted]], %[[c8]]
+// CHECK: %[[valid:.*]] = arith.andi %[[ge]], %[[lt]]
+// CHECK: gpu.printf "%d", %{{.*}}
+func.func @bounds_check_pad_left() {
+    %c0 = arith.constant 0 : index
+    rock.transforming_for (%arg0) = [#transform_map_pad_left](%c0) (%arg1) = validity bounds [16] strides [1] {
+        %arg1_i32 = arith.extui %arg1 : i1 to i32
+        gpu.printf "%d", %arg1_i32 : i32
+    }
+    return
+}
 }