[mlir] Modify SuperVectorize to generate select op->combiner op

[ayzhuang]

Summary

Insert the select op before the combiner op when vectorizing a
reduction loop that needs a mask, so the vectorized reduction loop
can pass isLoopParallel check and be transformed correctly in later
passes.

[dcaballe]

Thanks, Amy! In general, moving the select op before the operation makes more sense to me since the op might have side effects that would make the execution incorrect otherwise.
I hope we can replace all of this once we have the masking proposal in place.

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp

1447

I'm wondering if the matchReduction utility in SliceAnalysis.h could simplify part of this code. Would you mind giving it a try?

[ayzhuang]

Address review comments to use matchReduction.

[dcaballe]

Thanks!

[ayzhuang]

@dcaballe Thank you for the review. I'll merge it tomorrow if there are no more comments.

mlir/lib/Dialect/Affine/Transforms/SuperVectorize.cpp

static Operation *vectorizeAffineYieldOp(AffineYieldOp yieldOp,
                                          VectorizationState &state) {
   Operation *newYieldOp = widenOp(yieldOp, state);
   Operation *newParentOp = state.builder.getInsertionBlock()->getParentOp();
 
   // If there is a mask for this loop then we must prevent garbage values from
   // being added to the accumulator by inserting `select` operations, for
   // example:
   //
-  //   %res = arith.addf %acc, %val : vector<128xf32>
-  //   %res_masked = select %mask, %res, %acc : vector<128xi1>, vector<128xf32>
-  //   affine.yield %res_masked : vector<128xf32>
+  //   %val_masked = select %mask, %val, %neutralCst : vector<128xi1>,
+  //   vector<128xf32>
+  //   %res = arith.addf %acc, %val_masked : vector<128xf32>
+  //   affine.yield %res : vector<128xf32>
   //
   if (Value mask = state.vecLoopToMask.lookup(newParentOp)) {
     state.builder.setInsertionPoint(newYieldOp);
     for (unsigned i = 0; i < newYieldOp->getNumOperands(); ++i) {
       Value result = newYieldOp->getOperand(i);
       Value iterArg = cast<AffineForOp>(newParentOp).getRegionIterArgs()[i];
-      Value maskedResult = state.builder.create<arith::SelectOp>(
-          result.getLoc(), mask, result, iterArg);
+      Operation *defOp = result.getDefiningOp();
+      assert(defOp && defOp->getNumOperands() == 2 && "must be a binary op");
+      assert((defOp->getOperand(0) == iterArg ||
+              defOp->getOperand(1) == iterArg) &&
+             "must use iterArg");
+      Value input = defOp->getOperand(0) == iterArg ? defOp->getOperand(1)
+                                                    : defOp->getOperand(0);
+      // IterOperands are neutral element vectors.
+      Value neutralVal = cast<AffineForOp>(newParentOp).getIterOperands()[i];
+      state.builder.setInsertionPoint(defOp);
+      Value maskedInput = state.builder.create<arith::SelectOp>(
+          input.getLoc(), mask, input, neutralVal);
       LLVM_DEBUG(
-          dbgs() << "\n[early-vect]+++++ masking a yielded vector value: "
-                 << maskedResult);
-      newYieldOp->setOperand(i, maskedResult);
+          dbgs() << "\n[early-vect]+++++ masking an input to a binary op that"
+                    "produces value for a yield Op: "
+                 << maskedInput);
+      defOp->replaceUsesOfWith(input, maskedInput);
     }
   }
 
   state.builder.setInsertionPointAfter(newParentOp);
   return newYieldOp;
 }
 
 /// Encodes Operation-specific behavior for vectorization. In general we

mlir/test/Dialect/Affine/SuperVectorize/vectorize_reduction.mlir

 // CHECK:       #[[$map0:.*]] = affine_map<([[d0:.*]]) -> (-[[d0]] + 500)>
 // CHECK-LABEL: @vecdim_reduction_masked
 // CHECK:       affine.for %{{.*}} = 0 to 256 {
 // CHECK:         %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
 // CHECK:         %[[vred:.*]] = affine.for %[[iv:.*]] = 0 to 500 step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {
 // CHECK:           %[[elems_left:.*]] = affine.apply #[[$map0]](%[[iv]])
 // CHECK:           %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>
 // CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
-// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[ld]] : vector<128xf32>
-// CHECK:           %[[new_acc:.*]] = arith.select %[[mask]], %[[add]], %[[red_iter]] : vector<128xi1>, vector<128xf32>
-// CHECK:           affine.yield %[[new_acc]] : vector<128xf32>
+// CHECK:           %[[select:.*]] = arith.select %[[mask]], %[[ld]], %[[vzero]] : vector<128xi1>, vector<128xf32>
+// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[select]] : vector<128xf32>
+// CHECK:           affine.yield %[[add]] : vector<128xf32>
 // CHECK:         }
 // CHECK:         %[[final_sum:.*]] = vector.reduction <add>, %[[vred:.*]] : vector<128xf32> into f32
 // CHECK:         affine.store %[[final_sum]], %{{.*}} : memref<256xf32>
 // CHECK:       }
 
 // -----
 
 // The number of iteration is not known, so a mask has to be applied.

 // CHECK:       #[[$map1:.*]] = affine_map<([[d0:.*]]){{\[}}[[s0:.*]]{{\]}} -> (-[[d0]] + [[s0]])>
 // CHECK-LABEL: @vecdim_reduction_masked_unknown_ub
 // CHECK:       affine.for %{{.*}} = 0 to 256 {
 // CHECK:         %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
 // CHECK:         %[[vred:.*]] = affine.for %[[iv:.*]] = 0 to %[[bnd:.*]] step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {
 // CHECK:           %[[elems_left:.*]] = affine.apply #[[$map1]](%[[iv]])[%[[bnd]]]
 // CHECK:           %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>
 // CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
-// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[ld]] : vector<128xf32>
-// CHECK:           %[[new_acc:.*]] = arith.select %[[mask]], %[[add]], %[[red_iter]] : vector<128xi1>, vector<128xf32>
-// CHECK:           affine.yield %[[new_acc]] : vector<128xf32>
+// CHECK:           %[[select:.*]] = arith.select %[[mask]], %[[ld]], %[[vzero]] : vector<128xi1>, vector<128xf32>
+// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[select]] : vector<128xf32>
+// CHECK:           affine.yield %[[add]] : vector<128xf32>
 // CHECK:         }
 // CHECK:         %[[final_sum:.*]] = vector.reduction <add>, %[[vred:.*]] : vector<128xf32> into f32
 // CHECK:         affine.store %[[final_sum]], %{{.*}} : memref<256xf32>
 // CHECK:       }
 
 // -----
 
 // The lower bound is nonzero, but the number of iterations is divisible by the
affine.for %i = 0 to 256 {
    }
    affine.store %final_red, %out[%i] : memref<256xf32>
  }
  return
 }
 
 // CHECK:       #[[$map2:.*]] = affine_map<([[d0:.*]]) -> (-[[d0]] + 512)>
 // CHECK-LABEL: @vecdim_reduction_masked_unknown_lb
+// CHECK:         %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
 // CHECK:         %{{.*}} = affine.for %[[iv:.*]] = %[[lb:.*]] to 512 step 128 iter_args(%[[red_iter:.*]] = {{.*}}) -> (vector<128xf32>) {
 // CHECK:           %[[elems_left:.*]] = affine.apply #[[$map2]](%[[iv]])
 // CHECK:           %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>
 // CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
-// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[ld]] : vector<128xf32>
-// CHECK:           %[[new_acc:.*]] = arith.select %[[mask]], %[[add]], %[[red_iter]] : vector<128xi1>, vector<128xf32>
-// CHECK:           affine.yield %[[new_acc]] : vector<128xf32>
+// CHECK:           %[[select:.*]] = arith.select %[[mask]], %[[ld]], %[[vzero]] : vector<128xi1>, vector<128xf32>
+// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[select]] : vector<128xf32>
+// CHECK:           affine.yield %[[add]] : vector<128xf32>
 
 // -----
 
 // The upper bound is a minimum expression.
 
 func @vecdim_reduction_complex_ub(%in: memref<256x512xf32>, %out: memref<256xf32>, %M: index, %N: index) {
  %cst = arith.constant 0.000000e+00 : f32
  affine.for %i = 0 to 256 {
    %final_red = affine.for %j = 0 to min affine_map<(d0, d1) -> (d0, d1*2)>(%M, %N) iter_args(%red_iter = %cst) -> (f32) {
      %ld = affine.load %in[%i, %j] : memref<256x512xf32>
      %add = arith.addf %red_iter, %ld : f32
      affine.yield %add : f32
    }
    affine.store %final_red, %out[%i] : memref<256xf32>
  }
  return
 }
 
 // CHECK:       #[[$map3:.*]] = affine_map<([[d0:.*]], [[d1:.*]]) -> ([[d0]], [[d1]] * 2)>
 // CHECK:       #[[$map3_sub:.*]] = affine_map<([[d0:.*]], [[d1:.*]]) -> ([[d0]] - [[d1]])>
 // CHECK-LABEL: @vecdim_reduction_complex_ub
+// CHECK:         %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
 // CHECK:         %{{.*}} = affine.for %[[iv:.*]] = 0 to min #[[$map3]](%[[M:.*]], %[[N:.*]]) step 128 iter_args(%[[red_iter:.*]] = {{.*}}) -> (vector<128xf32>) {
 // CHECK:           %[[ub:.*]] = affine.min #[[$map3]](%[[M]], %[[N]])
 // CHECK:           %[[elems_left:.*]] = affine.apply #[[$map3_sub]](%[[ub]], %[[iv]])
 // CHECK:           %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>
 // CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
-// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[ld]] : vector<128xf32>
-// CHECK:           %[[new_acc:.*]] = arith.select %[[mask]], %[[add]], %[[red_iter]] : vector<128xi1>, vector<128xf32>
-// CHECK:           affine.yield %[[new_acc]] : vector<128xf32>
+// CHECK:           %[[select:.*]] = arith.select %[[mask]], %[[ld]], %[[vzero]] : vector<128xi1>, vector<128xf32>
+// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[select]] : vector<128xf32>
+// CHECK:           affine.yield %[[add]] : vector<128xf32>
 
 // -----
 
 // The same mask is applied to both reductions.
 
 func @vecdim_two_reductions_masked(%in: memref<256x512xf32>, %out: memref<512xf32>) {
  %cst = arith.constant 0.000000e+00 : f32
  affine.for %i = 0 to 256 {
    %final_sum, %final_expsum = affine.for %j = 0 to 500 iter_args(%sum_iter = %cst, %expsum_iter = %cst) -> (f32, f32) {
      %ld = affine.load %in[%i, %j] : memref<256x512xf32>
      %exp = math.exp %ld : f32
      %add = arith.addf %sum_iter, %ld : f32
      %eadd = arith.addf %expsum_iter, %exp : f32
      affine.yield %add, %eadd : f32, f32
    }
    affine.store %final_sum, %out[2*%i] : memref<512xf32>
    affine.store %final_expsum, %out[2*%i + 1] : memref<512xf32>
  }
  return
 }
 
 // CHECK:       #[[$map4:.*]] = affine_map<([[d0:.*]]) -> (-[[d0]] + 500)>
 // CHECK-LABEL: @vecdim_two_reductions_masked
 // CHECK:       affine.for %{{.*}} = 0 to 256 {
+// CHECK:         %[[vzero0:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
+// CHECK:         %[[vzero1:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
 // CHECK:         %{{.*}} = affine.for %[[iv:.*]] = 0 to 500 step 128 iter_args(%[[sum_iter:.*]] = {{.*}}, %[[esum_iter:.*]] = {{.*}}) -> (vector<128xf32>, vector<128xf32>) {
 // CHECK:           %[[elems_left:.*]] = affine.apply #[[$map4]](%[[iv]])
 // CHECK:           %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>
 // CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
 // CHECK:           %[[exp:.*]] = math.exp %[[ld]] : vector<128xf32>
-// CHECK:           %[[add:.*]] = arith.addf %[[sum_iter]], %[[ld]] : vector<128xf32>
-// CHECK:           %[[eadd:.*]] = arith.addf %[[esum_iter]], %[[exp]] : vector<128xf32>
-// CHECK:           %[[new_acc:.*]] = arith.select %[[mask]], %[[add]], %[[sum_iter]] : vector<128xi1>, vector<128xf32>
-// CHECK:           %[[new_eacc:.*]] = arith.select %[[mask]], %[[eadd]], %[[esum_iter]] : vector<128xi1>, vector<128xf32>
-// CHECK:           affine.yield %[[new_acc]], %[[new_eacc]] : vector<128xf32>
+// CHECK:           %[[select0:.*]] = arith.select %[[mask]], %[[ld]], %[[vzero0]] : vector<128xi1>, vector<128xf32>
+// CHECK:           %[[add:.*]] = arith.addf %[[sum_iter]], %[[select0]] : vector<128xf32>
+// CHECK:           %[[select1:.*]] = arith.select %[[mask]], %[[exp]], %[[vzero1]] : vector<128xi1>, vector<128xf32>
+// CHECK:           %[[eadd:.*]] = arith.addf %[[esum_iter]], %[[select1]] : vector<128xf32>
+// CHECK:           affine.yield %[[add]], %[[eadd]] : vector<128xf32>
 // CHECK:         }