[mlir][Linalg] Enable vectorization of explicit broadcasts

[dcaballe]

Summary

This patch teaches isProjectedPermutation and inverseAndBroadcastProjectedPermutation
utilities to deal with maps representing an explicit broadcast, e.g., (d0, d1) -> (d0, 0).
This extension is needed to enable vectorization of such explicit broadcast in Linalg.

[ThomasRaoux]

mlir/lib/IR/AffineMap.cpp

498–528

Does that mean (d0, d1) -> (0, 0) would be a projected permuation? That doesn't sound right.
I think you need a different function to check for something like prjected permutation with broadcast like a generalization of isPermutationOfMinorIdentityWithBroadcasting?

[pifon2a]

Could you please verify if all tests still pass? For example, we might need to modify applyPermutationMap in AffineMap.h.

template <typename T>
SmallVector<T> applyPermutationMap(AffineMap map, llvm::ArrayRef<T> source) {
  assert(map.isProjectedPermutation());
  assert(map.getNumInputs() == source.size());
  SmallVector<T> result;
  result.reserve(map.getNumResults());
  for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) {
    /* fix
    if (auto expr = map.getResult(i).dyn_cast<AffineConstantExpr>()) {
      result.push_back(0);
      continue;
    }
    */
    unsigned dim = map.getDimPosition(i);
    result.push_back(source[dim]);
  }
  return result;
}

[nicolasvasilache]

I have something similar locally but I am not clear this is actually needed for now.
Would you mind holding off a bit until I clear up my stack and avoid conflicts?

[dcaballe]

Addressed Alex's feedback for now.

[dcaballe]

In D111563#3056030, @pifon2a wrote:

Could you please verify if all tests still pass? For example, we might need to modify applyPermutationMap in AffineMap.h.

Good catch! Yes, I had run all the lit tests and nothing failed. I fixed this, thanks!

In D111563#3056288, @nicolasvasilache wrote:

I have something similar locally but I am not clear this is actually needed for now.
Would you mind holding off a bit until I clear up my stack and avoid conflicts?

Sure, I'll wait for you. After looking at your code, I'm addressing the review comments since they will also apply to your code.

mlir/lib/IR/AffineMap.cpp

498–528

Good point! Note that this utility, not sure if intentionally, is already returning true for (d0, d1) -> (). (d0, d1) -> (0, 0) would be the equivalent preserving the dims. I think it boils down to what the definition of "projected permutation" is. I asked @nicolasvasilache and it seems to be ok to have zeros in a projected permutation. We could add logic to skip (d0, d1) -> () and (d0, d1) -> (0, 0) but they look like a valid corner case to me. WDYT?

[ThomasRaoux]

mlir/lib/IR/AffineMap.cpp

498–528

In my mind projected permutation meant it has a permutation of the input dimensions and some dimensions may be missing (projection), so (d0, d1) -> () sounds like a valid projected permutation but having zeros instead of dimensions seems a bit odd. @nicolasvasilache is the referenced so if he thinks it is correct to allow extra zeros it is fine with me but this seems less intuitive to me.
It would be good to look at how this helper is used right now and check if any places assume that it means that all the results are dimensions.

[nicolasvasilache]

My stack is now flushed, feel free to land this @dcaballe and thanks for your patience!

Well technically, orthogonal projection of a vector v=(v0, .. vn) on the canonical axis e_k is the vector v=(v0, .. 0_k .. vn).
So far we have liberally jumped the gun a bit by also dropping those dimensions so that other invariants agree.

This LGTM; I would just make the behavior optional with a flag to avoid surprises (for now).

[dcaballe]

Adding a flag and enable it only for vectorization.
This should be ready to go.

mlir/include/mlir/IR/AffineMap.h

 ///    affine_map<(d0, d1, d2, d3) -> (d2)>
 /// ```
 ///
 /// returns:
 ///
 /// ```mlir
 ///    affine_map<(d0) -> (0, 0, d0, 0)>
 /// ```
+/// Example 4:
+///
+/// ```mlir
+///    affine_map<(d0, d1, d2) -> (d0, 0)>
+/// ```
+///
+/// returns:
+///
+/// ```mlir
+///    affine_map<(d0, d1) -> (d0, 0, 0)>
+/// ```
 AffineMap inverseAndBroadcastProjectedPermuation(AffineMap map);
 
 /// Concatenates a list of `maps` into a single AffineMap, stepping over
 /// potentially empty maps. Assumes each of the underlying map has 0 symbols.
 /// The resulting map has a number of dims equal to the max of `maps`' dims and
 /// the concatenated results as its results.
 /// Returns an empty map if all input `maps` are empty.
 ///

 
 /// Apply a permutation from `map` to `source` and return the result.
 template <typename T>
 SmallVector<T> applyPermutationMap(AffineMap map, llvm::ArrayRef<T> source) {
   assert(map.isProjectedPermutation());
   assert(map.getNumInputs() == source.size());
   SmallVector<T> result;
   result.reserve(map.getNumResults());
-  for (unsigned i = 0, e = map.getNumResults(); i < e; ++i) {
-    unsigned dim = map.getDimPosition(i);
-    result.push_back(source[dim]);
+  for (AffineExpr expr : map.getResults()) {
+    if (auto dimExpr = expr.dyn_cast<AffineDimExpr>()) {
+      result.push_back(source[dimExpr.getPosition()]);
+    } else if (auto constExpr = expr.dyn_cast<AffineConstantExpr>()) {
+      assert(constExpr.getValue() == 0 &&
+             "Unexpected constant in projected permutation map");
+      result.push_back(0);
+    } else {
+      llvm_unreachable("Unexpected result in projected permutation map");
+    }
   }
   return result;
 }
 
 inline raw_ostream &operator<<(raw_ostream &os, AffineMap map) {
   map.print(os);
   return os;
 }

mlir/lib/IR/AffineMap.cpp

SmallVector<int64_t, 4> AffineMap::compose(ArrayRef<int64_t> values) const {
   auto resMap = compose(AffineMap::get(0, 0, exprs, ctx));
   SmallVector<int64_t, 4> res;
   res.reserve(resMap.getNumResults());
   for (auto e : resMap.getResults())
     res.push_back(e.cast<AffineConstantExpr>().getValue());
   return res;
 }
 
 bool AffineMap::isProjectedPermutation() const {
   if (getNumSymbols() > 0)
     return false;
+
+  // Having more results than inputs means that results have duplicated dims or
+  // zeros that can't be mapped to input dims.
+  if (getNumResults() > getNumInputs())
+    return false;
+
   SmallVector<bool, 8> seen(getNumInputs(), false);
+  // A projected permutation can have, at most, only one instance of each input
+  // dimension in the result expressions. Zeros are allowed as long as the
+  // number of result expressions is lower or equal than the number of input
+  // expressions.
   for (auto expr : getResults()) {
     if (auto dim = expr.dyn_cast<AffineDimExpr>()) {
       if (seen[dim.getPosition()])
         return false;
       seen[dim.getPosition()] = true;
-      continue;
-    }
+    } else if (auto constExpr = expr.dyn_cast<AffineConstantExpr>()) {
+      if (constExpr.getValue() != 0)
         return false;
+    } else {
+      // Neither dim nor zero.
+      return false;
+    }
   }
+
+  // Results are either dims or zeros and zeros can be mapped to input dims.
   return true;
 }
 
 bool AffineMap::isPermutation() const {
   if (getNumDims() != getNumResults())
     return false;
   return isProjectedPermutation();
 }
 
 AffineMap AffineMap::getSubMap(ArrayRef<unsigned> resultPos) const {

 
 AffineMap mlir::inverseAndBroadcastProjectedPermuation(AffineMap map) {
   assert(map.isProjectedPermutation());
   MLIRContext *context = map.getContext();
   AffineExpr zero = mlir::getAffineConstantExpr(0, context);
   // Start with all the results as 0.
   SmallVector<AffineExpr, 4> exprs(map.getNumInputs(), zero);
   for (unsigned i : llvm::seq(unsigned(0), map.getNumResults())) {
-    // Reverse each dimension existing in the oringal map result.
+    // Skip zeros from input map. 'exprs' is already initialized to zero.
+    if (auto constExpr = map.getResult(i).dyn_cast<AffineConstantExpr>()) {
+      assert(constExpr.getValue() == 0 &&
+             "Unexpected constant in projected permutation");
+      (void)constExpr;
+      continue;
+    }
+
+    // Reverse each dimension existing in the original map result.
     exprs[map.getDimPosition(i)] = getAffineDimExpr(i, context);
   }
   return AffineMap::get(map.getNumResults(), /*symbolCount=*/0, exprs, context);
 }
 
 AffineMap mlir::concatAffineMaps(ArrayRef<AffineMap> maps) {
   unsigned numResults = 0, numDims = 0, numSymbols = 0;
   for (auto m : maps)

mlir/test/Dialect/Linalg/vectorization.mlir

%red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
                          ins(%arg0 : tensor<4x4xf32>) outs(%fill : tensor<4xf32>) {
   ^bb0(%in0: f32, %out0: f32):  // no predecessors
     %min = minf %in0, %out0 : f32
     linalg.yield %min : f32
   } -> tensor<4xf32>
   return %red : tensor<4xf32>
 }
 
+// -----
+
+// CHECK-DAG: #[[$M5:.*]] = affine_map<(d0, d1) -> (d0, 0)>
+
+// CHECK-LABEL:   func @explicit_broadcast(
+func @explicit_broadcast(%arg0: tensor<4x4xf32>, %arg1: tensor<4x1xf32>) -> tensor<4x4xf32> {
+  // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true]} : tensor<4x4xf32>, vector<4x4xf32>
+  // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true], permutation_map = #[[$M5]]} : tensor<4x1xf32>, vector<4x4xf32>
+  // CHECK: subf {{.*}} : vector<4x4xf32>
+  // CHECK: vector.transfer_write {{.*}} {in_bounds = [true, true]} : vector<4x4xf32>, tensor<4x4xf32>
+  %c0 = constant 0.0 : f32
+  %init = linalg.init_tensor [4, 4] : tensor<4x4xf32>
+  %fill = linalg.fill(%c0, %init) : f32, tensor<4x4xf32> -> tensor<4x4xf32>
+  %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                                          affine_map<(d0, d1) -> (d0, 0)>,
+                                          affine_map<(d0, d1) -> (d0, d1)>],
+   iterator_types = ["parallel", "parallel"]}
+   ins(%arg0, %arg1 : tensor<4x4xf32>, tensor<4x1xf32>)
+   outs(%fill : tensor<4x4xf32>) {
+    ^bb0(%arg7: f32, %arg8: f32, %arg9: f32):
+      %40 = subf %arg7, %arg8 : f32
+      linalg.yield %40 : f32
+    } -> tensor<4x4xf32>
+  return %red : tensor<4x4xf32>
+}
+
+// -----
+
+// CHECK-DAG: #[[$M6:.*]] = affine_map<(d0, d1) -> (d0, 0)>
+// CHECK-DAG: #[[$M7:.*]] = affine_map<(d0) -> (d0, 0)>
+
+// CHECK-LABEL:   func @fused_broadcast_red_2d
+func @fused_broadcast_red_2d(%arg0: tensor<4x4xf32>, %arg1: tensor<4x1xf32>) -> tensor<4xf32> {
+  // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true]} : tensor<4x4xf32>, vector<4x4xf32>
+  // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true], permutation_map = #[[$M6]]} : tensor<4x1xf32>, vector<4x4xf32>
+  // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true], permutation_map = #[[$M7]]} : tensor<4xf32>, vector<4x4xf32>
+  // CHECK: subf {{.*}} : vector<4x4xf32>
+  // CHECK: math.exp {{.*}} : vector<4x4xf32>
+  // CHECK: addf {{.*}} : vector<4x4xf32>
+  // CHECK: vector.multi_reduction #vector.kind<add>, {{.*}} : vector<4x4xf32> to vector<4xf32>
+  // CHECK: vector.transfer_write {{.*}} {in_bounds = [true]} : vector<4xf32>, tensor<4xf32>
+  %c0 = constant 0.0 : f32
+  %init = linalg.init_tensor [4] : tensor<4xf32>
+  %fill = linalg.fill(%c0, %init) : f32, tensor<4xf32> -> tensor<4xf32>
+  %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,
+                                          affine_map<(d0, d1) -> (d0, 0)>,
+                                          affine_map<(d0, d1) -> (d0)>],
+   iterator_types = ["parallel", "reduction"]}
+   ins(%arg0, %arg1 : tensor<4x4xf32>, tensor<4x1xf32>)
+   outs(%fill : tensor<4xf32>) {
+    ^bb0(%arg7: f32, %arg8: f32, %arg9: f32):
+      %40 = subf %arg7, %arg8 : f32
+      %41 = math.exp %40 : f32
+      %42 = addf %41, %arg9 : f32
+      linalg.yield %42 : f32
+    } -> tensor<4xf32>
+  return %red : tensor<4xf32>
+}