[MLIR][Affine][VectorOps] Vectorize uniform values in SuperVectorizer

[dcaballe]

Summary

This patch adds basic support for vectorization of uniform values to SuperVectorizer.
For now, only invariant values to the target vector loops are considered uniform. This
enables the vectorization of loops that use function arguments and external definitions
to the vector loops. We could extend uniform support in the future if we implement some
kind of divergence analysis algorithm.

[rriddle]

mlir/include/mlir/IR/Builders.h

337

nit: Spell out auto here.

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

1009

nit: Drop trivial braces here.

[aartbik]

mlir/include/mlir/IR/Builders.h

332

This is very nitty, but the comment order does not match the code order.
Would you mind using the same order, it just reads better in my mind....

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

939

can't we just simply return false right away, and get rid of the local bool altogether?

950

use "auto" less, VectorType and Value are reasonable types here

mlir/test/Dialect/Affine/SuperVectorize/uniform_divergent.mlir

16

hmm, this is a new style to me, would you mind either moving all CHECKs of a method together, or at least move these into the function body itself

(or correct me if I am wrong and point me to some place that uses this style)

[dcaballe]

Thanks! Addressing feedback.

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

939

Good catch! Thanks!

950

Per LLVM's coding standards (https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable):

Don’t “almost always” use auto, but do use auto with initializers like cast<Foo>(...) or other places where the type is already obvious from the context.

This is a bit subjective but IMO these two cases fall into the same category as cast, isa, dyn_cast, since the type is already spelled out in the same line:

VectorType vectorTy = getVectorType...
BroadcastOp bcast = builder.create<BroadcastOp>...

Is MLIR following different guidelines in this regard? Otherwise, I would prefer to stick to the guidelines (whatever they are) since this is one of those topics where every reviewer/developer has a different opinion on :).

mlir/test/Dialect/Affine/SuperVectorize/uniform_divergent.mlir

16

You can find similar tests, for example, here:
mlir/test/Transforms/loop-fusion.mlir
mlir/test/Transforms/buffer-placement.mlir

There are probably more. What is exactly the concern? That LABEL check is separated from the other CHECKs?
For loop transformations I personally think it's important that we keep both the input and the output as clean as possible so that we can clearly see the loop structure of them. Mixing CHECK rules with the input IR wouldn't help in that regard. I can move the LABEL next to the other rules but the LABEL check is usually placed right before the function op since it works as a barrier across tests.

[nicolasvasilache]

Thanks for extending the SuperVectorizer !

[dcaballe]

Thanks! Any other comments, @aartbik?

[aartbik]

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

950

Fine with me. Thanks for checking the standard.

mlir/test/Dialect/Affine/SuperVectorize/uniform_divergent.mlir

16

I am fine with this if you prefer. The concern was that breaking the CHECKs this way makes it a bit harder to connect the pieces in your head when reading, but I agree that is subjective.

[mehdi_amini]

mlir/include/mlir/IR/Builders.h

343

This API is a bit problematic in the ambiguity it creates for SingleResult operation which implicitly convert to Value as well.

Op op;
builder.setInsertionPointAfter(op);

Will fail after this change.

[dcaballe]

mlir/include/mlir/IR/Builders.h

343

Good point! Should we rename it to setInsertionPointAfterValue, for example?

mlir/include/mlir/IR/Builders.h

public:
   }
 
   /// Sets the insertion point to the node after the specified operation, which
   /// will cause subsequent insertions to go right after it.
   void setInsertionPointAfter(Operation *op) {
     setInsertionPoint(op->getBlock(), ++Block::iterator(op));
   }
 
+  /// Sets the insertion point to the node after the specified value. If value
+  /// is a BlockArgument, sets the insertion point to the start of its block.
+  /// Otherwise, sets the insertion point to the node after the defining
+  /// operation of the value. This will cause subsequent insertions to go right
+  /// after it.
+  void setInsertionPointAfter(Value val) {
+    if (auto op = val.getDefiningOp()) {
+      setInsertionPointAfter(op);
+    } else {
+      auto blockArg = val.cast<BlockArgument>();
+      setInsertionPointToStart(blockArg.getOwner());
+    }
+  }
+
   /// Sets the insertion point to the start of the specified block.
   void setInsertionPointToStart(Block *block) {
     setInsertionPoint(block, block->begin());
   }
 
   /// Sets the insertion point to the end of the specified block.
   void setInsertionPointToEnd(Block *block) {
     setInsertionPoint(block, block->end());

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

 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
 using namespace mlir;
+using namespace vector;
 
 ///
 /// Implements a high-level vectorization strategy on a Function.
 /// The abstraction used is that of super-vectors, which provide a single,
 /// compact, representation in the vector types, information that is expected
 /// to reduce the impact of the phase ordering problem
 ///
 /// Vector granularity:
static Value vectorizeConstant(Operation *op, ConstantOp constant, Type type) {
   auto *constantOpInst = constant.getOperation();
 
   OperationState state(loc, constantOpInst->getName().getStringRef(), {},
                        {vectorType}, {b.getNamedAttr("value", attr)});
 
   return b.createOperation(state)->getResult(0);
 }
 
+/// Returns the vector type resulting from applying the provided vectorization
+/// strategy on the scalar type.
+static VectorType getVectorType(Type scalarTy,
+                                const VectorizationStrategy *strategy) {
+  assert(!scalarTy.isa<VectorType>() && "Expected scalar type");
+  return VectorType::get(strategy->vectorSizes, scalarTy);
+}
+
+/// Returns true if the provided value is vector uniform given the vectorization
+/// strategy.
+// TODO: For now, only values that are invariants to all the loops in the
+// vectorization strategy are considered vector uniforms.
+static bool isUniformDefinition(Value value,
+                                const VectorizationStrategy *strategy) {
+  bool isUniform = true;
+  for (auto loopToDim : strategy->loopToVectorDim) {
+    auto loop = cast<AffineForOp>(loopToDim.first);
+    isUniform &= loop.isDefinedOutsideOfLoop(value);
+  }
+  return isUniform;
+}
+
+/// Generates a broadcast op for the provided uniform value using the
+/// vectorization strategy in 'state'.
+static Value vectorizeUniform(Value value, VectorizationState *state) {
+  OpBuilder builder(value.getContext());
+  builder.setInsertionPointAfter(value);
+
+  auto vectorTy = getVectorType(value.getType(), state->strategy);
+  auto bcast = builder.create<BroadcastOp>(value.getLoc(), vectorTy, value);
+
+  // Add broadcast to the replacement map to reuse it for other uses.
+  state->replacementMap[value] = bcast;
+  return bcast;
+}
+
 /// Tries to vectorize a given operand `op` of Operation `op` during
 /// def-chain propagation or during terminal vectorization, by applying the
 /// following logic:
 /// 1. if the defining operation is part of the vectorizedSet (i.e. vectorized
 ///    useby -def propagation), `op` is already in the proper vector form;
 /// 2. otherwise, the `op` may be in some other vector form that fails to
 ///    vectorize atm (i.e. broadcasting required), returns nullptr to indicate
 ///    failure;
 /// 3. if the `op` is a constant, returns the vectorized form of the constant;
-/// 4. non-constant scalars are currently non-vectorizable, in particular to
+/// 4. if the `op` is uniform, returns a vector broadcast of the `op`;
+/// 5. non-constant scalars are currently non-vectorizable, in particular to
 ///    guard against vectorizing an index which may be loop-variant and needs
 ///    special handling.
 ///
 /// In particular this logic captures some of the use cases where definitions
 /// that are not scoped under the current pattern are needed to vectorize.
 /// One such example is top level function constants that need to be splatted.
 ///
 /// Returns an operand that has been vectorized to match `state`'s strategy if
static Value vectorizeOperand(Value operand, Operation *op,
   }
   // 2. TODO: broadcast needed.
   if (operand.getType().isa<VectorType>()) {
     LLVM_DEBUG(dbgs() << "-> non-vectorizable");
     return nullptr;
   }
   // 3. vectorize constant.
   if (auto constant = operand.getDefiningOp<ConstantOp>()) {
-    return vectorizeConstant(
-        op, constant,
-        VectorType::get(state->strategy->vectorSizes, operand.getType()));
+    return vectorizeConstant(op, constant,
+                             getVectorType(operand.getType(), state->strategy));
+  }
+  // 4. Uniform values.
+  if (isUniformDefinition(operand, state->strategy)) {
+    return vectorizeUniform(operand, state);
   }
-  // 4. currently non-vectorizable.
+  // 5. currently non-vectorizable.
   LLVM_DEBUG(dbgs() << "-> non-vectorizable: " << operand);
   return nullptr;
 }
 
 /// Encodes Operation-specific behavior for vectorization. In general we assume
 /// that all operands of an op must be vectorized but this is not always true.
 /// In the future, it would be nice to have a trait that describes how a
 /// particular operation vectorizes. For now we implement the case distinction

mlir/test/Dialect/Affine/SuperVectorize/uniform_divergent.mlir

+// RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=128" -split-input-file | FileCheck %s
+
+// Specific tests to check vectorization of uniform/divergent values.
+
+// CHECK-LABEL: @uniform_arg
+// CHECK-SAME:  %[[in:.*]]: memref<512xf32>,
+// CHECK-SAME:  %[[uniform:.*]]: f32
+func @uniform_arg(%in : memref<512xf32>, %uniform : f32) {
+  affine.for %i = 0 to 512 {
+    %ld = affine.load %in[%i] : memref<512xf32>
+    %add = addf %ld, %uniform : f32
+  }
+  return
+}
+
+// CHECK-NEXT: %[[bcast:.*]] = vector.broadcast %[[uniform]] : f32 to vector<128xf32>
+// CHECK-NEXT: affine.for
+// CHECK:        addf %{{.*}}, %[[bcast]] : vector<128xf32>
+
+// -----
+
+// CHECK-LABEL: @multi_use_uniform_arg
+// CHECK-SAME:  %[[in:.*]]: memref<512xf32>
+// CHECK-SAME:  %[[uniform:.*]]: f32
+func @multi_use_uniform_arg(%in : memref<512xf32>, %uniform : f32) {
+  affine.for %i = 0 to 512 {
+    %ld = affine.load %in[%i] : memref<512xf32>
+    %user0 = addf %ld, %uniform : f32
+    %user1 = addf %ld, %uniform : f32
+  }
+  return
+}
+
+// CHECK-NEXT: %[[bcast:.*]] = vector.broadcast %[[uniform]] : f32 to vector<128xf32>
+// CHECK-NOT:  vector.broadcast
+// CHECK-NEXT: affine.for
+// CHECK:        addf %{{.*}}, %[[bcast]] : vector<128xf32>
+// CHECK:        addf %{{.*}}, %[[bcast]] : vector<128xf32>
+
+// -----
+
+// CHECK-LABEL: @uniform_load
+func @uniform_load(%A : memref<?x?xf32>, %C : memref<?x?xf32>) {
+  %c0 = constant 0 : index
+  %N = dim %A, %c0 : memref<?x?xf32>
+  affine.for %i = 0 to %N {
+    %uniform_ld = affine.load %A[%i, %i] : memref<?x?xf32>
+    affine.for %j = 0 to %N {
+      %b = affine.load %A[%i, %j] : memref<?x?xf32>
+      %c = addf %uniform_ld, %b : f32
+    }
+  }
+  return
+}
+
+// CHECK:      affine.for
+// CHECK-NEXT:   %[[uniform_ld:.*]] = affine.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>
+// CHECK-NEXT:   %[[bcast:.*]] = vector.broadcast %[[uniform_ld]] : f32 to vector<128xf32>
+// CHECK-NEXT:   affine.for
+// CHECK:          addf %[[bcast]], %{{.*}} : vector<128xf32>

mlir/test/Dialect/Affine/SuperVectorize/vectorize_1d.mlir

// CHECK:  affine.for %{{.*}}{{[0-9]*}} = 0 to %{{[0-9]*}} {
    affine.for %i15 = 0 to %M { // not vectorized due to condition below
      affine.if #set0(%i15) {
        %a15 = affine.load %A[%c0, %c0] : memref<?x?xf32>
      }
    }
    return
 }
 
-// This should not vectorize and should not crash.
-// CHECK-LABEL: @vec_rejected_11
-func @vec_rejected_11(%A : memref<?x?xf32>, %C : memref<?x?xf32>) {
-  %c0 = constant 0 : index
-  %N = dim %A, %c0 : memref<?x?xf32>
-  affine.for %i = 0 to %N {
-// CHECK-NOT: vector
-    %a = affine.load %A[%i, %i] : memref<?x?xf32> // not vectorized
-    affine.for %j = 0 to %N {
-      %b = affine.load %A[%i, %j] : memref<?x?xf32> // may be vectorized
-// CHECK-NOT: vector
-      %c = addf %a, %b : f32 // not vectorized because %a wasn't
-// CHECK-NOT: vector
-      affine.store %c, %C[%i, %j] : memref<?x?xf32> // not vectorized because %c wasn't
-    }
-  }
-  return
-}
-
 // This should not vectorize due to the sequential dependence in the scf.
 // CHECK-LABEL: @vec_rejected_sequential
 func @vec_rejected_sequential(%A : memref<?xf32>) {
   %c0 = constant 0 : index
   %N = dim %A, %c0 : memref<?xf32>
   affine.for %i = 0 to %N {
     // CHECK-NOT: vector
     %a = affine.load %A[%i] : memref<?xf32>
     // CHECK-NOT: vector
     affine.store %a, %A[%i + 1] : memref<?xf32>
   }
   return
 }