[release/3.4] "[AMD] guard FoldTrueCmpI from tensors (triton-lang#7281)" (triton-lang#7342)

Cherry-pick of 7281, original description:

I made a mistake and didn't check the input/output types of the
arith.cmpi so in some cases FoldTrueCmpI might fold tensors (e.g.,
masks). So we should guard against that.

Co-authored-by: Maksim Levental <[email protected]>

Author: iupaikov-amd

test/TritonGPU/amd/amd-fold-true-cmpi.mlir

@@ -156,3 +156,22 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32} {
 // CHECK-NEXT:      ttg.local_dealloc %[[VAL_24]] : !ttg.memdesc<1x32x128xf16, #[[$ATTR_4]], #[[$ATTR_5]], mutable>
 // CHECK-NEXT:      tt.return %[[VAL_58]] : tensor<128x128xf32, #[[$ATTR_2]]>
 // CHECK-NEXT:      }
+
+// -----
+
+module attributes {"ttg.num-warps" = 4 : i32} {
+  tt.func @dontfoldtensor() -> tensor<128xi1> {
+    %t0 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>
+    %t1 = tt.make_range {end = 257 : i32, start = 129 : i32} : tensor<128xi32>
+    %cmp = arith.cmpi sgt, %t1, %t0 : tensor<128xi32>
+    tt.return %cmp: tensor<128xi1>
+  }
+}
+
+// CHECK-LABEL:   tt.func @dontfoldtensor
+// CHECK-NOT:       arith.constant dense<true>
+// CHECK:           %[[VAL_0:.*]] = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32>
+// CHECK:           %[[VAL_1:.*]] = tt.make_range {end = 257 : i32, start = 129 : i32} : tensor<128xi32>
+// CHECK:           %[[VAL_2:.*]] = arith.cmpi sgt, %[[VAL_1]], %[[VAL_0]] : tensor<128xi32>
+// CHECK:           tt.return %[[VAL_2]] : tensor<128xi1>
+// CHECK:         }

third_party/amd/lib/Analysis/RangeAnalysis.cpp

@@ -601,7 +601,8 @@ struct FoldTrueCmpIOp : OpRewritePattern<arith::CmpIOp> {
 
   LogicalResult matchAndRewrite(arith::CmpIOp cmpOp,
                                 PatternRewriter &rewriter) const override {
-    if (cmpIIsStaticallyTrue(*solver, cmpOp)) {
+    if (llvm::isa<IntegerType, IndexType>(cmpOp.getType()) &&
+        cmpIIsStaticallyTrue(*solver, cmpOp)) {
       if (failed(mlir::dataflow::maybeReplaceWithConstant(*solver, rewriter,
                                                           cmpOp.getResult()))) {
         LDBG("failed to replace with constant op: " << cmpOp);