Replace select with mul

Author: nbpatel

mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp

@@ -446,17 +446,16 @@ struct LinearizeVectorSplat final
 };
 
 /// This pattern converts the CreateMaskOp to work on a linearized vector.
-//  The pattern currently supports only 2D masks with a unit outer dimension.
+/// It currently supports only 2D masks with a unit outer dimension.
 /// Following,
 ///   vector.create_mask %arg0, %arg1 : vector<1x4xi1>
 /// is converted to:
 ///   %zero = arith.constant 0 : index
-///   %cmpi = arith.cmpi sle, %arg0, %zero : index
-///   %splat = vector.splat %cmpi : vector<4xi1>
-///   %cst = arith.constant dense<false> : vector<4xi1>
-///   %mask = vector.create_mask %arg1 : vector<4xi1>
-///   %out = arith.select %splat, %cst, %mask : vector<4xi1>
-///   %out_1d = vector.shape_cast %out : vector<4xi1> to vector<1x4xi1>
+///   %cmpi = arith.cmpi sgt, %arg0, %zero : index
+///   %index = arith.index_cast %cmpi : i1 to index
+///   %mul = arith.muli %index, %arg1 : index
+///   %mask = vector.create_mask %mul : vector<4xi1>
+///   %out_1d = vector.shape_cast %mask : vector<4xi1> to vector<1x4xi1>
 struct LinearizeVectorCreateMask final
     : OpConversionPattern<vector::CreateMaskOp> {
   using OpConversionPattern::OpConversionPattern;
@@ -483,25 +482,23 @@ struct LinearizeVectorCreateMask final
     if (!dstTy)
       return rewriter.notifyMatchFailure(createMaskOp, "cannot convert type.");
 
-    // Compare the first operand with 0. If it's less than or equal to 0,
-    // create a zero mask, else strip the first operand and create a mask
-    // using the second operand.
+    // Compare the first operand with 0. If it is greater than 0, the
+    // corresponding mask element is set to true, otherwise false.
+    // The result of the comparison is then multiplied with
+    // the second operand of create_mask to get the 1D mask.
     auto firstOperand = adaptor.getOperands().front();
     auto zero = rewriter.create<mlir::arith::ConstantIndexOp>(loc, 0);
-    auto isZeroOrNegative = rewriter.create<mlir::arith::CmpIOp>(
-        loc, mlir::arith::CmpIPredicate::sle, firstOperand, zero);
-    auto isZeroOrNegativeSplat =
-        rewriter.create<mlir::vector::SplatOp>(loc, dstTy, isZeroOrNegative);
-
-    // Use a select operation to choose between the masks.
-    auto zeroMask = rewriter.create<mlir::arith::ConstantOp>(
-        loc, dstTy, rewriter.getZeroAttr(dstTy));
-    auto newMask = rewriter.create<mlir::vector::CreateMaskOp>(
-        loc, dstTy, adaptor.getOperands().back());
-    auto result = rewriter.create<mlir::arith::SelectOp>(
-        loc, isZeroOrNegativeSplat, zeroMask, newMask);
+    auto isNonZero = rewriter.create<mlir::arith::CmpIOp>(
+        loc, mlir::arith::CmpIPredicate::sgt, firstOperand, zero);
+    auto isNonZeroIndex = rewriter.create<mlir::arith::IndexCastOp>(
+        loc, rewriter.getIndexType(), isNonZero);
+    auto secondOperand = adaptor.getOperands().back();
+    auto maskSize = rewriter.create<mlir::arith::MulIOp>(
+        loc, rewriter.getIndexType(), isNonZeroIndex, secondOperand);
 
-    rewriter.replaceOp(createMaskOp, result.getResult());
+    auto newMask = rewriter.create<mlir::vector::CreateMaskOp>(
+        loc, dstTy, maskSize.getResult());
+    rewriter.replaceOp(createMaskOp, newMask);
     return success();
   }
 };

mlir/test/Dialect/Vector/linearize.mlir

@@ -347,39 +347,18 @@ func.func @linearize_scalable_vector_splat(%arg0: i32) -> vector<4x[2]xi32> {
 }
 
 // -----
-// ALL-LABEL: linearize_create_mask
-func.func @linearize_create_mask() -> vector<1x16xi1> {
+
+// CHECK-LABEL: linearize_create_mask
+// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index) -> vector<1x16xi1>
+func.func @linearize_create_mask(%arg0 : index, %arg1 : index) -> vector<1x16xi1> {
+  
   // CHECK: %[[C0:.*]] = arith.constant 0 : index
-  // CHECK: %[[C10:.*]] = arith.constant 10 : index
-  // CHECK: %[[C0_0:.*]] = arith.constant 0 : index
-  // CHECK: %[[CMP:.*]] = arith.cmpi sle, %[[C0]], %[[C0_0]] : index
-  // CHECK: %[[SPLAT:.*]] = vector.splat %[[CMP]] : vector<16xi1>
-  // CHECK: %[[CST:.*]] = arith.constant dense<false> : vector<16xi1>
-  // CHECK: %[[MASK_1D:.*]] = vector.create_mask %[[C10]] : vector<16xi1>
-  // CHECK: %[[SELECT:.*]] = arith.select %[[SPLAT]], %[[CST]], %[[MASK_1D]] : vector<16xi1>, vector<16xi1>
-  // CHECK: %[[CAST:.*]] = vector.shape_cast %[[SELECT]] : vector<16xi1> to vector<1x16xi1>
+  // CHECK: %[[CMP:.*]] = arith.cmpi sgt, %[[ARG0]], %[[C0]] : index
+  // CHECK: %[[INDEXCAST:.*]] = arith.index_cast %[[CMP]] : i1 to index
+  // CHECK: %[[MULI:.*]] = arith.muli %[[INDEXCAST]], %[[ARG1]] : index
+  // CHECK: %[[MASK_1D:.*]] = vector.create_mask %[[MULI]] : vector<16xi1>
+  // CHECK: %[[CAST:.*]] = vector.shape_cast %[[MASK_1D]] : vector<16xi1> to vector<1x16xi1>
   // CHECK: return %[[CAST]] : vector<1x16xi1>
-  %c0 = arith.constant 0 : index
-  %c10 = arith.constant 10 : index
-  %0 = vector.create_mask %c0, %c10 : vector<1x16xi1>
+  %0 = vector.create_mask %arg0, %arg1 : vector<1x16xi1>
   return %0 : vector<1x16xi1>
 }
-
-// -----
-// ALL-LABEL: linearize_scalable_create_mask
-func.func @linearize_scalable_create_mask() -> vector<1x[16]xi1> {
-  // CHECK: %[[C0:.*]] = arith.constant 0 : index
-  // CHECK: %[[C10:.*]] = arith.constant 10 : index
-  // CHECK: %[[C0_0:.*]] = arith.constant 0 : index
-  // CHECK: %[[CMP:.*]] = arith.cmpi sle, %[[C0]], %[[C0_0]] : index
-  // CHECK: %[[SPLAT:.*]] = vector.splat %[[CMP]] : vector<[16]xi1>
-  // CHECK: %[[CST:.*]] = arith.constant dense<false> : vector<[16]xi1>
-  // CHECK: %[[MASK_1D:.*]] = vector.create_mask %[[C10]] : vector<[16]xi1>
-  // CHECK: %[[SELECT:.*]] = arith.select %[[SPLAT]], %[[CST]], %[[MASK_1D]] : vector<[16]xi1>, vector<[16]xi1>
-  // CHECK: %[[CAST:.*]] = vector.shape_cast %[[SELECT]] : vector<[16]xi1> to vector<1x[16]xi1>
-  // CHECK: return %[[CAST]] : vector<1x[16]xi1>
-  %c0 = arith.constant 0 : index
-  %c10 = arith.constant 10 : index
-  %0 = vector.create_mask %c0, %c10 : vector<1x[16]xi1>
-  return %0 : vector<1x[16]xi1>
-}