Address comments and add negative scalar unit dim test

Author: qedawkins

mlir/include/mlir/Dialect/SCF/IR/SCF.h

@@ -111,7 +111,7 @@ LoopNest buildLoopNest(OpBuilder &builder, Location loc, ValueRange lbs,
 /// `replacement` value with a different type. A callback is used to insert
 /// cast ops inside the block to account for type differences.
 using ValueTypeCastFnTy =
-    std::function<Value(OpBuilder &, Location loc, Type, Value)>;
+    llvm::function_ref<Value(OpBuilder &, Location loc, Type, Value)>;
 SmallVector<Value> replaceAndCastForOpIterArg(RewriterBase &rewriter,
                                               scf::ForOp forOp,
                                               OpOperand &operand,

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

@@ -1837,8 +1837,7 @@ struct DropUnitDimsFromScfForOp final : OpRewritePattern<scf::ForOp> {
         continue;
 
       // Create a new ForOp with that iter operand replaced.
-      mlir::scf::ValueTypeCastFnTy castFn = [](OpBuilder &b, Location loc,
-                                               Type type, Value source) {
+      auto castFn = [](OpBuilder &b, Location loc, Type type, Value source) {
         return b.create<vector::ShapeCastOp>(loc, type, source);
       };
 

mlir/test/Dialect/Vector/drop-unit-dims-with-shape-cast.mlir

@@ -282,3 +282,25 @@ func.func @scf_for_with_multiple_operands(%idx: index, %vec0: vector<1x4xf32>, %
 //       CHECK:     scf.yield %{{.*}}, %[[ADD]], %[[ADD]]
 //       CHECK:   %[[CASTBACK:.+]] = vector.shape_cast %[[LOOP]]#1 : vector<4xf32> to vector<1x4xf32>
 //       CHECK:   return %[[CASTBACK]]
+
+// -----
+
+func.func @scf_for_with_scalable_unit_dims(%vec: vector<1x[1]xf32>) -> vector<1x[1]xf32> {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  %res = scf.for %i = %c0 to %c4 step %c1 iter_args(%iter = %vec) -> vector<1x[1]xf32> {
+    %s = math.sqrt %iter : vector<1x[1]xf32>
+    scf.yield %s : vector<1x[1]xf32>
+  }
+  return %res : vector<1x[1]xf32>
+}
+
+// CHECK-LABEL: func.func @scf_for_with_scalable_unit_dims
+//  CHECK-SAME:   %[[VEC:[A-Za-z0-9]+]]: vector<1x[1]xf32>
+//       CHECK:   %[[CAST:.+]] = vector.shape_cast %[[VEC]] : vector<1x[1]xf32> to vector<[1]xf32>
+//       CHECK:   %[[LOOP:.+]] = scf.for {{.*}} iter_args(%[[ITER:.+]] = %[[CAST]])
+//       CHECK:     %[[SQRT:.+]] = math.sqrt %[[ITER]] : vector<[1]xf32>
+//       CHECK:     scf.yield %[[SQRT]]
+//       CHECK:   %[[CASTBACK:.+]] = vector.shape_cast %[[LOOP]] : vector<[1]xf32> to vector<1x[1]xf32>
+//       CHECK:   return %[[CASTBACK]]