Address Feedback

Author: nbpatel

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

@@ -1055,7 +1055,8 @@ struct UnrollCreateMaskPattern : public OpRewritePattern<vector::CreateMaskOp> {
 
     Value result = arith::ConstantOp::create(rewriter, loc, resultType,
                                              rewriter.getZeroAttr(resultType));
-    auto targetVectorType = VectorType::get(*targetShape, rewriter.getI1Type());
+    VectorType targetVectorType =
+        VectorType::get(*targetShape, rewriter.getI1Type());
     SmallVector<int64_t> strides(targetShape->size(), 1);
 
     // In each dimension (d), each unrolled vector computes its mask size as:
@@ -1068,20 +1069,20 @@ struct UnrollCreateMaskPattern : public OpRewritePattern<vector::CreateMaskOp> {
            llvm::enumerate(createMaskOp.getOperands())) {
         Value offsetVal =
             arith::ConstantIndexOp::create(rewriter, loc, offsets[i]);
-        Value adjustedMaskSize = arith::SubIOp::create(
-            rewriter, loc, originalMaskOperand, offsetVal);
+        Value adjustedMaskSize = rewriter.createOrFold<arith::SubIOp>(
+            loc, originalMaskOperand, offsetVal);
         Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
         Value unrolledDimSize =
             arith::ConstantIndexOp::create(rewriter, loc, (*targetShape)[i]);
         Value nonNegative =
-            arith::MaxSIOp::create(rewriter, loc, adjustedMaskSize, zero);
-        Value unrolledOperand =
-            arith::MinSIOp::create(rewriter, loc, nonNegative, unrolledDimSize);
+            rewriter.createOrFold<arith::MaxSIOp>(loc, adjustedMaskSize, zero);
+        Value unrolledOperand = rewriter.createOrFold<arith::MinSIOp>(
+            loc, nonNegative, unrolledDimSize);
         unrolledOperands.push_back(unrolledOperand);
       }
 
-      auto unrolledMask = vector::CreateMaskOp::create(
-          rewriter, loc, targetVectorType, unrolledOperands);
+      auto unrolledMask = rewriter.createOrFold<vector::CreateMaskOp>(
+          loc, targetVectorType, unrolledOperands);
       result = rewriter.createOrFold<vector::InsertStridedSliceOp>(
           loc, unrolledMask, result, offsets, strides);
     }

mlir/test/Dialect/Vector/vector-unroll-options.mlir

@@ -537,6 +537,22 @@ func.func @vector_create_mask(%size1: index, %size2: index) -> vector<16x16xi1>
 //       CHECK:   %[[INS11:.*]] = vector.insert_strided_slice %[[MASK11]], %[[INS10]] {offsets = [8, 8], strides = [1, 1]} : vector<8x8xi1> into vector<16x16xi1>
 //       CHECK:   return %[[INS11]] : vector<16x16xi1>
 
+func.func @vector_create_mask_constant_dim_sizes() -> vector<16x16xi1> {
+  %cst16 = arith.constant 16 : index
+  %0 = vector.create_mask %cst16, %cst16 : vector<16x16xi1>
+  return %0 : vector<16x16xi1>
+}
+
+// CHECK-LABEL: func @vector_create_mask_constant_dim_sizes() -> vector<16x16xi1> {
+// CHECK:   %[[CST:.*]] = arith.constant dense<false> : vector<16x16xi1>
+// CHECK:   %[[CST_0:.*]] = arith.constant dense<true> : vector<8x8xi1>
+// CHECK:   %[[S0:.*]] = vector.insert_strided_slice %[[CST_0]], %[[CST]] {offsets = [0, 0], strides = [1, 1]} : vector<8x8xi1> into vector<16x16xi1>
+// CHECK:   %[[S1:.*]] = vector.insert_strided_slice %[[CST_0]], %[[S0]] {offsets = [0, 8], strides = [1, 1]} : vector<8x8xi1> into vector<16x16xi1>
+// CHECK:   %[[S2:.*]] = vector.insert_strided_slice %[[CST_0]], %[[S1]] {offsets = [8, 0], strides = [1, 1]} : vector<8x8xi1> into vector<16x16xi1>
+// CHECK:   %[[S3:.*]] = vector.insert_strided_slice %[[CST_0]], %[[S2]] {offsets = [8, 8], strides = [1, 1]} : vector<8x8xi1> into vector<16x16xi1>
+// CHECK:   return %[[S3]] : vector<16x16xi1>
+
+
 func.func @shape_cast_1D(%v: vector<16xf32>) -> vector<2x2x4xf32> {
   %0 = vector.shape_cast %v : vector<16xf32> to vector<2x2x4xf32>
   return %0 : vector<2x2x4xf32>