Split up the and operation over the cycles

  - allows for the shifts to only compute the shift on the required
    cycle and not the entire input
  - removes the initial vector.shuffle as we can operate directly on the
    input and merge them after

Author: inbelic

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

@@ -837,23 +837,14 @@ static bool isCyclic(SmallVector<T> xs, int64_t cycleLen) {
   return true;
 }
 
-static SmallVector<int64_t> constructShuffles(int64_t inputSize,
-                                              int64_t numCycles,
+static SmallVector<int64_t> constructShuffles(int64_t numCycles,
                                               int64_t cycleLen, int64_t idx) {
-  // If idx == 1, then the first operand of the shuffle will be the mask which
-  // will have the original size. So we need to step through the mask with a
-  // stride of cycleSize.
-  // When idx > 1, then the first operand will be the size of (idx * cycleSize)
-  // and so we take the first idx elements of the input and then append the
-  // strided mask value.
-  int64_t inputStride = idx == 1 ? cycleLen : idx;
-
   SmallVector<int64_t> shuffles;
   for (int64_t cycle = 0; cycle < numCycles; cycle++) {
     for (int64_t inputIdx = 0; inputIdx < idx; inputIdx++) {
-      shuffles.push_back(cycle * inputStride + inputIdx);
+      shuffles.push_back(cycle * idx + inputIdx);
     }
-    shuffles.push_back(inputSize + cycle * cycleLen + idx);
+    shuffles.push_back(numCycles * idx + cycle);
   }
   return shuffles;
 }
@@ -917,47 +908,31 @@ Value BitCastRewriter::splatRewriteStep(
     PatternRewriter &rewriter, Location loc, Value initialValue,
     Value runningResult, const BitCastRewriter::Metadata &metadata) {
 
-  // Initial result will be the Shifted Mask which will have the shuffles size.
-  int64_t inputSize = metadata.shuffles.size();
-  int64_t numCycles = inputSize / cycleLen;
-
-  auto shuffleOp = rewriter.create<vector::ShuffleOp>(
-      loc, initialValue, initialValue, metadata.shuffles);
-
-  // Intersect with the mask.
-  VectorType shuffledVectorType = shuffleOp.getResultVectorType();
-  auto constOp = rewriter.create<arith::ConstantOp>(
-      loc, DenseElementsAttr::get(shuffledVectorType, metadata.masks));
-  Value andValue = rewriter.create<arith::AndIOp>(loc, shuffleOp, constOp);
-
+  int64_t numCycles = metadata.shuffles.size() / cycleLen;
+  ShapedType vectorType = dyn_cast<ShapedType>(initialValue.getType());
   Value result;
   for (int64_t idx = 0; idx < cycleLen; idx++) {
+    // Intersect with the mask.
+    auto constOp = rewriter.create<arith::ConstantOp>(
+        loc, DenseElementsAttr::get(vectorType, metadata.masks[idx]));
+    Value andValue = rewriter.create<arith::AndIOp>(loc, initialValue, constOp);
+
     auto shiftRightConstantOp = rewriter.create<arith::ConstantOp>(
-        loc, SplatElementsAttr::get(shuffledVectorType,
-                                    metadata.shiftRightAmounts[idx]));
+        loc,
+        SplatElementsAttr::get(vectorType, metadata.shiftRightAmounts[idx]));
     Value shiftedRight =
         rewriter.create<arith::ShRUIOp>(loc, andValue, shiftRightConstantOp);
 
     auto shiftLeftConstantOp = rewriter.create<arith::ConstantOp>(
-        loc, SplatElementsAttr::get(shuffledVectorType,
-                                    metadata.shiftLeftAmounts[idx]));
+        loc,
+        SplatElementsAttr::get(vectorType, metadata.shiftLeftAmounts[idx]));
     Value shiftedLeft =
         rewriter.create<arith::ShLIOp>(loc, shiftedRight, shiftLeftConstantOp);
 
-    if (result) {
-      SmallVector<int64_t> shuffles =
-          constructShuffles(inputSize, numCycles, cycleLen, idx);
-      result = rewriter.create<vector::ShuffleOp>(loc, result, shiftedLeft,
-                                                  shuffles);
-
-      // After the first shuffle in the chain, the size of the input result will
-      // grow as we append more shuffles together to reconstruct the
-      // shuffledVectorType size. Each iteration they will retain numCycles more
-      // elements.
-      inputSize = numCycles * (idx + 1);
-    } else {
-      result = shiftedLeft;
-    }
+    SmallVector<int64_t> shuffles = constructShuffles(numCycles, cycleLen, idx);
+    result = result ? rewriter.create<vector::ShuffleOp>(loc, result,
+                                                         shiftedLeft, shuffles)
+                    : shiftedLeft;
   }
 
   return result;

mlir/test/Dialect/Vector/vector-rewrite-narrow-types.mlir

@@ -230,14 +230,15 @@ func.func @f3ext(%a: vector<5xi8>) -> vector<8xi17> {
 }
 
 // CHECK-LABEL: func.func @fext_splat1(
-//  CHECK-SAME: %[[A:[0-9a-z]*]]: vector<2xi8>) -> vector<4xi16> {
+//  CHECK-SAME: %[[ARG:[0-9a-z]*]]: vector<2xi8>) -> vector<4xi16> {
 func.func @fext_splat1(%a: vector<2xi8>) -> vector<4xi16> {
-  // CHECK-DAG: %[[MASK:.*]] = arith.constant dense<[15, -16, 15, -16]> : vector<4xi8>
-  // CHECK-DAG: %[[SHR_CST:.*]] = arith.constant dense<4> : vector<4xi8>
-  // CHECK: %[[V0:.*]] = vector.shuffle %[[A]], %[[A]] [0, 0, 1, 1] : vector<2xi8>, vector<2xi8>
-  // CHECK: %[[A0:.*]] = arith.andi %[[V0]], %[[MASK]] : vector<4xi8>
-  // CHECK: %[[SHR0:.*]] = arith.shrui %[[A0]], %[[SHR_CST]] : vector<4xi8>
-  // CHECK: %[[V1:.*]] = vector.shuffle %[[A0]], %[[SHR0]] [0, 5, 2, 7] : vector<4xi8>, vector<4xi8>
+  // CHECK-DAG: %[[MASK0:.*]] = arith.constant dense<15> : vector<2xi8>
+  // CHECK-DAG: %[[MASK1:.*]] = arith.constant dense<-16> : vector<2xi8>
+  // CHECK-DAG: %[[SHR_CST:.*]] = arith.constant dense<4> : vector<2xi8>
+  // CHECK-DAG: %[[A0:.*]] = arith.andi %[[ARG]], %[[MASK0]] : vector<2xi8>
+  // CHECK-DAG: %[[A1:.*]] = arith.andi %[[ARG]], %[[MASK1]] : vector<2xi8>
+  // CHECK: %[[SHR0:.*]] = arith.shrui %[[A1]], %[[SHR_CST]] : vector<2xi8>
+  // CHECK: %[[V1:.*]] = vector.shuffle %[[A0]], %[[SHR0]] [0, 2, 1, 3] : vector<2xi8>, vector<2xi8>
   // CHECK: %[[RES:.*]] = arith.extui %[[V1]] : vector<4xi8> to vector<4xi16>
   // return %[[RES]] : vector<4xi16>
   %0 = vector.bitcast %a : vector<2xi8> to vector<4xi4>
@@ -246,20 +247,25 @@ func.func @fext_splat1(%a: vector<2xi8>) -> vector<4xi16> {
 }
 
 // CHECK-LABEL: func.func @fext_splat2(
-//  CHECK-SAME: %[[A:[0-9a-z]*]]: vector<3xi16>) -> vector<12xi32> {
+//  CHECK-SAME: %[[ARG:[0-9a-z]*]]: vector<3xi16>) -> vector<12xi32> {
 func.func @fext_splat2(%a: vector<3xi16>) -> vector<12xi32> {
-  // CHECK-DAG: %[[MASK:.*]] = arith.constant dense<[15, 240, 3840, -4096, 15, 240, 3840, -4096, 15, 240, 3840, -4096]> : vector<12xi16>
-  // CHECK-DAG: %[[SHR_CST0:.*]] = arith.constant dense<4> : vector<12xi16>
-  // CHECK-DAG: %[[SHR_CST1:.*]] = arith.constant dense<8> : vector<12xi16>
-  // CHECK-DAG: %[[SHR_CST2:.*]] = arith.constant dense<12> : vector<12xi16>
-  // CHECK: %[[V0:.*]] = vector.shuffle %[[A]], %[[A]] [0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2] : vector<3xi16>, vector<3xi16>
-  // CHECK: %[[A0:.*]] = arith.andi %[[V0]], %[[MASK]] : vector<12xi16>
-  // CHECK: %[[SHR0:.*]] = arith.shrui %[[A0]], %[[SHR_CST0]] : vector<12xi16>
-  // CHECK: %[[V1:.*]] = vector.shuffle %[[A0]], %[[SHR0]] [0, 13, 4, 17, 8, 21] : vector<12xi16>, vector<12xi16>
-  // CHECK: %[[SHR1:.*]] = arith.shrui %[[A0]], %[[SHR_CST1]] : vector<12xi16>
-  // CHECK: %[[V2:.*]] = vector.shuffle %[[V1]], %[[SHR1]] [0, 1, 8, 2, 3, 12, 4, 5, 16] : vector<6xi16>, vector<12xi16>
-  // CHECK: %[[SHR2:.*]] = arith.shrui %[[A0]], %[[SHR_CST2]] : vector<12xi16>
-  // CHECK: %[[V3:.*]] = vector.shuffle %[[V2]], %[[SHR2]] [0, 1, 2, 12, 3, 4, 5, 16, 6, 7, 8, 20] : vector<9xi16>, vector<12xi16>
+  // CHECK-DAG: %[[MASK0:.*]] = arith.constant dense<15> : vector<3xi16>
+  // CHECK-DAG: %[[MASK1:.*]] = arith.constant dense<240> : vector<3xi16>
+  // CHECK-DAG: %[[MASK2:.*]] = arith.constant dense<3840> : vector<3xi16>
+  // CHECK-DAG: %[[MASK3:.*]] = arith.constant dense<-4096> : vector<3xi16>
+  // CHECK-DAG: %[[SHR_CST0:.*]] = arith.constant dense<4> : vector<3xi16>
+  // CHECK-DAG: %[[SHR_CST1:.*]] = arith.constant dense<8> : vector<3xi16>
+  // CHECK-DAG: %[[SHR_CST2:.*]] = arith.constant dense<12> : vector<3xi16>
+  // CHECK: %[[A0:.*]] = arith.andi %[[ARG]], %[[MASK0]] : vector<3xi16>
+  // CHECK: %[[A1:.*]] = arith.andi %[[ARG]], %[[MASK1]] : vector<3xi16>
+  // CHECK: %[[SHR0:.*]] = arith.shrui %[[A1]], %[[SHR_CST0]] : vector<3xi16>
+  // CHECK: %[[V1:.*]] = vector.shuffle %[[A0]], %[[SHR0]] [0, 3, 1, 4, 2, 5] : vector<3xi16>, vector<3xi16>
+  // CHECK: %[[A2:.*]] = arith.andi %[[ARG]], %[[MASK2]] : vector<3xi16>
+  // CHECK: %[[SHR1:.*]] = arith.shrui %[[A2]], %[[SHR_CST1]] : vector<3xi16>
+  // CHECK: %[[V2:.*]] = vector.shuffle %[[V1]], %[[SHR1]] [0, 1, 6, 2, 3, 7, 4, 5, 8] : vector<6xi16>, vector<3xi16>
+  // CHECK: %[[A3:.*]] = arith.andi %[[ARG]], %[[MASK3]] : vector<3xi16>
+  // CHECK: %[[SHR2:.*]] = arith.shrui %[[A3]], %[[SHR_CST2]] : vector<3xi16>
+  // CHECK: %[[V3:.*]] = vector.shuffle %[[V2]], %[[SHR2]] [0, 1, 2, 9, 3, 4, 5, 10, 6, 7, 8, 11] : vector<9xi16>, vector<3xi16>
   // CHECK: %[[RES:.*]] = arith.extui %[[V3]] : vector<12xi16> to vector<12xi32>
   // CHEKC: return %[[RES]] : vector<12xi32>
   %0 = vector.bitcast %a : vector<3xi16> to vector<12xi4>