[mlir] [VectorOps] consolidate all vector utilities to one header/cc file

[merge_guards_bot]

Unit tests: pass. 62280 tests passed, 0 failed and 827 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[aartbik]

rebased with master

[merge_guards_bot]

Unit tests: fail. 62306 tests passed, 1 failed and 838 were skipped.

failed: libc++.std/thread/thread_mutex/thread_mutex_requirements/thread_timedmutex_requirements/thread_timedmutex_recursive/try_lock_until.pass.cpp

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[andydavis1]

Great cleanup. The call sites look great. Thanks for doing this!

mlir/include/mlir/Dialect/VectorOps/VectorUtils.h

31

I think this could be named "computeStrides" as it could be applied to non-slices as well. It just takes a 'shape' and 'sizes'.

40

Since this is a library it might be good to make the name more verbose:

computeElementOffsetsFromVectorSliceOffsets

40

Given the target slice sizes of a vector...

47

nit: I prefer "computeSliceSizes", but good either way.

[aartbik]

addressed comments

[aartbik]

Thanks! I was not sure on the naming, and like your suggestions better indeed. All addressed.

[dcaballe]

mlir/include/mlir/Dialect/VectorOps/VectorUtils.h

1

Utils.h -> VectorUtils.h?

37

All these utilities returning SmallVector should be turned into void utility_name(..., SmallVectorImpl<element_type>& output_name) so that they can accept SmallVectors of arbitrary static sizes.

mlir/lib/Dialect/VectorOps/VectorTransforms.cpp

115

just curious... why is this static_cast needed?

[aartbik]

mlir/include/mlir/Dialect/VectorOps/VectorUtils.h

37

Note that this is a refactoring of existing code, so I would prefer to do that as follow up. But just in case you insist, can you point me to at least one occurrence of this elsewhere in the tree?

mlir/lib/Dialect/VectorOps/VectorTransforms.cpp

115

size_t vs int64_t (may cause some compilers to complain about sign diff in comparison)

[aartbik]

Utils -> VectorUtils typo

[merge_guards_bot]

Unit tests: pass. 62309 tests passed, 0 failed and 838 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[dcaballe]

mlir/include/mlir/Dialect/VectorOps/VectorUtils.h

37

Agree, we can address this separately.
There are a few examples here: https://github.com/llvm/llvm-project/blob/master/llvm/include/llvm/Analysis/LoopInfo.h#L256

mlir/lib/Dialect/VectorOps/VectorTransforms.cpp

115

Thanks! Sorry, I thought the cast applied to the result of the comparison. My bad.

[aartbik]

mlir/include/mlir/Dialect/VectorOps/VectorUtils.h

37

Ah, I see what you mean. Thanks for pointing this out.

[merge_guards_bot]

Unit tests: pass. 62310 tests passed, 0 failed and 838 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

mlir/include/mlir/Dialect/VectorOps/VectorUtils.h

-//===- Utils.h - VectorOps Utils ----------------------------*- C++ -*-=======//
+//===- Utils.h - VectorOps Utilities ------------------------*- C++ -*-=======//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_VECTOROPS_VECTORUTILS_H_
 #define MLIR_DIALECT_VECTOROPS_VECTORUTILS_H_
 
 #include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/DenseMap.h"
 
 namespace mlir {
 
+// Forward declarations.
 class AffineApplyOp;
 class AffineForOp;
 class AffineMap;
 class Location;
 class MemRefType;
 class OpBuilder;
 class Operation;
 class Value;
 class VectorType;
 
+/// Given the shape and sizes of a vector, returns the slice strides
+/// for each dimension.
+SmallVector<int64_t, 4> makeSliceStrides(ArrayRef<int64_t> shape,
+                                         ArrayRef<int64_t> sizes);
+
+/// Given the slice strides together with a linear index in the dimension
+/// space, returns the vector-space offsets in each dimension for a
+/// de-linearized index.
+SmallVector<int64_t, 4> delinearize(ArrayRef<int64_t> sliceStrides,
+                                    int64_t linearIndex);
+
+/// Given the sizes of a vector, together with vector-space offsets, returns
+/// the element-space offsets for each dimension.
+SmallVector<int64_t, 4> toElementSpace(ArrayRef<int64_t> sizes,
+                                       ArrayRef<int64_t> vectorOffsets);
+
+/// Given the shape, sizes, and element-space offsets of a vector, returns
+/// the slize sizes for each dimension.
+SmallVector<int64_t, 4> makeSliceSizes(ArrayRef<int64_t> shape,
+                                       ArrayRef<int64_t> sizes,
+                                       ArrayRef<int64_t> elementOffsets);
+
 /// Computes and returns the multi-dimensional ratio of `superShape` to
 /// `subShape`. This is calculated by performing a traversal from minor to major
 /// dimensions (i.e. in reverse shape order). If integral division is not
 /// possible, returns None.
 /// The ArrayRefs are assumed (and enforced) to only contain > 1 values.
 /// This constraint comes from the fact that they are meant to be used with
 /// VectorTypes, for which the property holds by construction.
 ///

mlir/lib/Analysis/CMakeLists.txt

 add_llvm_library(MLIRAnalysis STATIC
   AffineAnalysis.cpp
   AffineStructures.cpp
   CallGraph.cpp
   Dominance.cpp
   InferTypeOpInterface.cpp
   Liveness.cpp
   LoopAnalysis.cpp
   MemRefBoundCheck.cpp
   NestedMatcher.cpp
   OpStats.cpp
   SliceAnalysis.cpp
   TestMemRefDependenceCheck.cpp
   TestParallelismDetection.cpp
   Utils.cpp
-  VectorAnalysis.cpp
   Verifier.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Analysis
   )
 add_dependencies(MLIRAnalysis
   MLIRAffineOps
   MLIRCallOpInterfacesIncGen
   MLIRTypeInferOpInterfaceIncGen
   MLIRLoopOps
   MLIRVectorOps
   )
 target_link_libraries(MLIRAnalysis MLIRAffineOps MLIRLoopOps MLIRVectorOps)

mlir/lib/Dialect/VectorOps/CMakeLists.txt

 add_llvm_library(MLIRVectorOps
   DialectRegistration.cpp
   VectorOps.cpp
   VectorTransforms.cpp
+  VectorUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/VectorOps
   )
 
 add_dependencies(MLIRVectorOps MLIRVectorOpsIncGen)
 add_dependencies(MLIRVectorOps MLIRVectorTransformPatternsIncGen)
 
 target_link_libraries(MLIRVectorOps MLIRIR)

mlir/lib/Dialect/VectorOps/VectorOps.cpp

 //===- VectorOps.cpp - MLIR Super Vectorizer Operations -------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements convenience types for working with super-vectorization
 // operations, in particular super-vector loads and stores.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/VectorOps/VectorOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
+#include "mlir/Dialect/VectorOps/VectorUtils.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Support/Functional.h"
 #include "mlir/Support/LLVM.h"
auto is_vector_type_of_rank = [&](Type t) {
     return t.isa<VectorType>() && t.cast<VectorType>().getRank() == rank;
   };
   if (!llvm::all_of(tupleType.getTypes(), is_vector_type_of_rank))
     return op->emitError("requires vector tuple elements of rank ") << rank;
   // Check that 'sizes' and 'strides' are of size == 'rank'.
   if (sizes.size() != rank || strides.size() != rank)
     return op->emitError("requires sizes and strides of rank ") << rank;
 
-  // Compute the number of slices in each dimension.
-  // TODO(andydavis) Move this into a slice generation helper function.
-  auto shape = vectorType.getShape();
-  SmallVector<int64_t, 4> dimSliceCounts(rank);
-  for (unsigned i = 0; i < rank; ++i)
-    dimSliceCounts[i] = ceilDiv(shape[i], sizes[i]);
-  // Compute the strides between slices in each dimension.
-  SmallVector<int64_t, 4> sliceStrides(rank);
-  sliceStrides[rank - 1] = 1;
-  for (int i = rank - 2; i >= 0; --i)
-    sliceStrides[i] = sliceStrides[i + 1] * dimSliceCounts[i + 1];
-
   // Generate each slice shape based on 'sizes', 'strides' and 'vectorType',
   // and verify that the same matches the corresponding tuple element 'i'.
+  auto shape = vectorType.getShape();
+  auto sliceStrides = makeSliceStrides(shape, sizes);
   for (int64_t i = 0, e = tupleType.size(); i < e; ++i) {
-    // De-linearize w.r.t. 'sliceStrides'.
-    SmallVector<int64_t, 4> vectorOffsets(rank);
-    int64_t linearIndex = i;
-    for (unsigned j = 0; j < rank; ++j) {
-      vectorOffsets[j] = linearIndex / sliceStrides[j];
-      linearIndex %= sliceStrides[j];
-    }
-    // Convert from unrolled vector-space offsets to element-space offsets.
-    auto offsets = mlir::functional::zipMap(
-        [](int64_t v1, int64_t v2) { return v1 * v2; }, vectorOffsets, sizes);
-    // Initialize 'sliceSizes' to target 'sizes'
-    SmallVector<int64_t, 4> sliceSizes(sizes.begin(), sizes.end());
-    for (unsigned j = 0; j < rank; ++j) {
-      // Based on 'offsets' and 'shape' clip some dim sizes for partial tiles.
-      sliceSizes[j] = std::min(sliceSizes[j], shape[j] - offsets[j]);
-    }
+    auto vectorOffsets = delinearize(sliceStrides, i);
+    auto elementOffsets = toElementSpace(sizes, vectorOffsets);
+    auto sliceSizes = makeSliceSizes(shape, sizes, elementOffsets);
     // Create slice VectorType type.
     auto sliceVectorType =
         VectorType::get(sliceSizes, vectorType.getElementType());
     // Verify that 'sliceVectorType' matches tupleType.getTypes(i)
     if (sliceVectorType != tupleType.getType(i))
       return op->emitError("invalid tuple element type ") << sliceVectorType;
   }
   return success();

mlir/lib/Dialect/VectorOps/VectorTransforms.cpp

 #include "mlir/IR/Function.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Support/Functional.h"
-#include "mlir/Support/MathExtras.h"
 #include "mlir/Support/STLExtras.h"
 
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "vector-to-vector"
 

 static int64_t linearize(ArrayRef<int64_t> offsets, ArrayRef<int64_t> basis) {
   assert(offsets.size() == basis.size());
   int64_t linearIndex = 0;
   for (unsigned idx = 0, e = basis.size(); idx < e; ++idx)
     linearIndex += offsets[idx] * basis[idx];
   return linearIndex;
 }
 
-/// Given a shape with sizes greater than 0 along all dimensions, returns the
-/// delinearized components of linearIndex along shape.
-static SmallVector<int64_t, 8> delinearize(int64_t linearIndex,
-                                           ArrayRef<int64_t> basis) {
-  SmallVector<int64_t, 8> res;
-  res.reserve(basis.size());
-  for (unsigned idx = 0, e = basis.size(); idx < e; ++idx) {
-    assert(basis[idx] > 0);
-    res.push_back(linearIndex / basis[idx]);
-    linearIndex %= basis[idx];
-  }
-  // Sanity check.
-  assert(linearIndex == 0 && "linear index remainder must be 0");
-  return res;
-}
-
 // Clones `op` into a new operations that takes `operands` and returns
 // `resultTypes`.
 static Operation *cloneOpWithOperandsAndTypes(PatternRewriter &builder,
                                               Location loc, Operation *op,
                                               ArrayRef<Value> operands,
                                               ArrayRef<Type> resultTypes) {
   OperationState res(loc, op->getName().getStringRef(), operands, resultTypes,
                      op->getAttrs());

 // of 'vectorType' unrolled by 'sizes' and 'strides'.
 // TODO(andydavis) Move this to a utility function and share it with
 // Extract/InsertSlicesOp verification.
 static TupleType generateExtractSlicesOpResultType(VectorType vectorType,
                                                    ArrayRef<int64_t> sizes,
                                                    ArrayRef<int64_t> strides,
                                                    PatternRewriter &builder) {
   assert(llvm::all_of(strides, [](int64_t s) { return s == 1; }));
-  unsigned rank = vectorType.getRank();
-  assert(sizes.size() == rank);
-  assert(strides.size() == rank);
+  assert(static_cast<int64_t>(sizes.size()) == vectorType.getRank());
+  assert(static_cast<int64_t>(strides.size()) == vectorType.getRank());
 
   // Compute shape ratio of 'shape' and 'sizes'.
   auto shape = vectorType.getShape();
   auto maybeDimSliceCounts = shapeRatio(shape, sizes);
   assert(maybeDimSliceCounts.hasValue());
   auto sliceDimCounts = *maybeDimSliceCounts;
 
   // Compute strides w.r.t number of slices in each dimension.
-  auto basis = computeStrides(sliceDimCounts);
+  auto sliceStrides = computeStrides(sliceDimCounts);
   int64_t sliceCount = computeMaxLinearIndex(sliceDimCounts);
   SmallVector<Type, 4> vectorTypes(sliceCount);
   for (unsigned i = 0; i < sliceCount; ++i) {
-    // De-linearize w.r.t. 'basis'.
-    auto vectorOffsets = delinearize(i, basis);
-    // Convert from unrolled vector-space offsets to element-space offsets.
-    auto offsets = zipMap([](int64_t v1, int64_t v2) { return v1 * v2; },
-                          vectorOffsets, sizes);
-    // Initialize 'sliceSizes' to target 'sizes'
-    SmallVector<int64_t, 4> sliceSizes(sizes.begin(), sizes.end());
-    for (unsigned j = 0; j < rank; ++j) {
-      // Based on 'offsets' and 'shape' clip some dim sizes for partial tiles.
-      sliceSizes[j] = std::min(sliceSizes[j], shape[j] - offsets[j]);
-    }
+    auto vectorOffsets = delinearize(sliceStrides, i);
+    auto elementOffsets = toElementSpace(sizes, vectorOffsets);
+    auto sliceSizes = makeSliceSizes(shape, sizes, elementOffsets);
     // Create Vector type and add to 'vectorTypes[i]'.
     vectorTypes[i] = VectorType::get(sliceSizes, vectorType.getElementType());
   }
   return TupleType::get(vectorTypes, builder.getContext());
 }
 
 // UnrolledVectorState aggregates per-operand/result vector state required for
 // unrolling.
static Value unrollSingleResultStructuredOp(Operation *op,
   for (unsigned i = 0; i < numVectors; ++i) {
     int64_t operandIndex = vectors[i].operandIndex;
     auto operand = operandIndex >= 0 ? op->getOperand(operandIndex) : nullptr;
     initUnrolledVectorState(vectors[i].type, operand, vectors[i].indexMap,
                             targetShape, unrolledVectorState[i], builder);
   }
   // Compute number of total unrolled instances.
   auto numUnrolledInstances = computeMaxLinearIndex(unrollFactors);
-  auto basis = computeStrides(unrollFactors);
+  auto sliceStrides = computeStrides(unrollFactors);
 
   auto &resultValueState = unrolledVectorState[resultIndex];
   auto unrolledResultType = VectorType::get(resultValueState.unrolledShape,
                                             shapedType.getElementType());
 
   // Initialize caches for intermediate vector results.
   std::vector<SmallVector<Value, 4>> caches(numVectors);
   for (unsigned i = 0; i < numVectors; ++i)
     caches[i].resize(unrolledVectorState[i].numInstances);
 
   // Unroll 'numUnrolledInstances' of 'op', storing results in 'caches'.
   for (unsigned i = 0; i < numUnrolledInstances; ++i) {
-    // De-linearize w.r.t. 'basis'.
-    auto vectorOffsets = delinearize(i, basis);
-    // Convert from unrolled vector-space offsets to element-space offsets.
-    auto offsets = zipMap([](int64_t v1, int64_t v2) { return v1 * v2; },
-                          vectorOffsets, targetShape);
+    auto vectorOffsets = delinearize(sliceStrides, i);
+    auto elementOffsets = toElementSpace(targetShape, vectorOffsets);
     // Get cached slice (or create slice) for each operand at 'offsets'.
     SmallVector<Value, 3> operands;
     operands.resize(op->getNumOperands());
     for (unsigned i = 0; i < numVectors; ++i) {
       int64_t operandIndex = vectors[i].operandIndex;
       if (operandIndex < 0)
         continue; // Output
       auto operand = op->getOperand(operandIndex);
       operands[operandIndex] = getOrCreateUnrolledVectorSlice(
-          op->getLoc(), unrolledVectorState[i], vectorOffsets, offsets,
+          op->getLoc(), unrolledVectorState[i], vectorOffsets, elementOffsets,
           vectors[i].indexMap, operand, caches[i], builder);
     }
     // Create op on sliced vector arguments.
     auto resultVector =
         cloneOpWithOperandsAndTypes(builder, op->getLoc(), op, operands,
                                     unrolledResultType)
             ->getResult(0);
 

 generateTransferOpSlices(VectorType vectorType, TupleType tupleType,
                          ArrayRef<int64_t> sizes, ArrayRef<int64_t> strides,
                          ArrayRef<Value> indices, PatternRewriter &rewriter,
                          function_ref<void(unsigned, ArrayRef<Value>)> fn) {
   // Compute strides w.r.t. to slice counts in each dimension.
   auto maybeDimSliceCounts = shapeRatio(vectorType.getShape(), sizes);
   assert(maybeDimSliceCounts.hasValue());
   auto sliceDimCounts = *maybeDimSliceCounts;
-  auto basis = computeStrides(sliceDimCounts);
+  auto sliceStrides = computeStrides(sliceDimCounts);
 
   int64_t numSlices = tupleType.size();
   unsigned numSliceIndices = indices.size();
   auto *ctx = rewriter.getContext();
   for (unsigned i = 0; i < numSlices; ++i) {
-    // De-linearize w.r.t. 'basis'.
-    auto vectorOffsets = delinearize(i, basis);
-    // Convert from unrolled vector-space offsets to element-space offsets.
-    auto offsets = zipMap([](int64_t v1, int64_t v2) { return v1 * v2; },
-                          vectorOffsets, sizes);
+    auto vectorOffsets = delinearize(sliceStrides, i);
+    auto elementOffsets = toElementSpace(sizes, vectorOffsets);
     // Compute 'sliceIndices' by adding 'sliceOffsets[i]' to 'indices[i]'.
     SmallVector<Value, 4> sliceIndices(numSliceIndices);
     for (auto it : llvm::enumerate(indices)) {
       auto expr = getAffineDimExpr(0, ctx) +
-                  getAffineConstantExpr(offsets[it.index()], ctx);
+                  getAffineConstantExpr(elementOffsets[it.index()], ctx);
       auto map = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr);
       sliceIndices[it.index()] = rewriter.create<AffineApplyOp>(
           it.value().getLoc(), map, ArrayRef<Value>(it.value()));
     }
     // Call 'fn' to generate slice 'i' at 'sliceIndices'.
     fn(i, sliceIndices);
   }
 }

 ///   %b = vector.strided_slice %0
 ///   ..
 ///   %x = vector.tuple %a, %b, ..
 class ExtractSlicesOpLowering
     : public OpRewritePattern<vector::ExtractSlicesOp> {
 public:
   using OpRewritePattern<vector::ExtractSlicesOp>::OpRewritePattern;
 
-  // TODO(ajcbik): refactor slice utilities out into VectorUtils.h
   PatternMatchResult matchAndRewrite(vector::ExtractSlicesOp op,
                                      PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 
     VectorType vectorType = op.getSourceVectorType();
-    int64_t rank = vectorType.getRank();
     auto shape = vectorType.getShape();
 
     SmallVector<int64_t, 4> sizes;
     op.getSizes(sizes);
     SmallVector<int64_t, 4> strides;
     op.getStrides(strides); // all-ones at the moment
 
-    // Compute the number of slices in each dimension.
-    SmallVector<int64_t, 4> sliceDimCounts(rank);
-    for (int64_t r = 0; r < rank; ++r)
-      sliceDimCounts[r] = ceilDiv(shape[r], sizes[r]);
-
     // For each element in the tuple, generate the proper strided slice.
-    auto basis = computeStrides(sliceDimCounts);
     TupleType tupleType = op.getResultTupleType();
     int64_t tupleSize = tupleType.size();
     SmallVector<Value, 4> tupleValues(tupleSize);
+    auto sliceStrides = makeSliceStrides(shape, sizes);
     for (int64_t i = 0; i < tupleSize; ++i) {
-      // De-linearize w.r.t. 'basis'.
-      auto vectorOffsets = delinearize(i, basis);
-      // Convert from unrolled vector-space offsets to element-space offsets.
-      auto elementOffsets = mlir::functional::zipMap(
-          [](int64_t v1, int64_t v2) { return v1 * v2; }, vectorOffsets, sizes);
-      // Compute the size of each slice.
-      SmallVector<int64_t, 4> sliceSizes(rank);
-      for (int64_t r = 0; r < rank; ++r)
-        sliceSizes[r] = std::min(sizes[r], shape[r] - elementOffsets[r]);
+      auto vectorOffsets = delinearize(sliceStrides, i);
+      auto elementOffsets = toElementSpace(sizes, vectorOffsets);
+      auto sliceSizes = makeSliceSizes(shape, sizes, elementOffsets);
       // Insert in tuple.
       tupleValues[i] = rewriter.create<vector::StridedSliceOp>(
           loc, op.vector(), elementOffsets, sliceSizes, strides);
     }
 
     rewriter.replaceOpWithNewOp<vector::TupleOp>(op, tupleType, tupleValues);
     return matchSuccess();
   }

 //    %t2 = vector.tuple_get %0, 1
 ///   %r2 = vector.insert_strided_slice %r1, %t2
 ///   ..
 ///   %x  = ..
 class InsertSlicesOpLowering : public OpRewritePattern<vector::InsertSlicesOp> {
 public:
   using OpRewritePattern<vector::InsertSlicesOp>::OpRewritePattern;
 
-  // TODO(ajcbik): refactor slice utilities out into VectorUtils.h
   PatternMatchResult matchAndRewrite(vector::InsertSlicesOp op,
                                      PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 
     VectorType vectorType = op.getResultVectorType();
-    int64_t rank = vectorType.getRank();
     auto shape = vectorType.getShape();
 
     SmallVector<int64_t, 4> sizes;
     op.getSizes(sizes);
     SmallVector<int64_t, 4> strides;
     op.getStrides(strides); // all-ones at the moment
 
-    // Compute the number of slices in each dimension.
-    SmallVector<int64_t, 4> sliceDimCounts(rank);
-    for (int64_t r = 0; r < rank; ++r)
-      sliceDimCounts[r] = ceilDiv(shape[r], sizes[r]);
-
     // Prepare result.
     auto elemType = vectorType.getElementType();
     Value zero = rewriter.create<ConstantOp>(loc, elemType,
                                              rewriter.getZeroAttr(elemType));
     Value result = rewriter.create<SplatOp>(loc, vectorType, zero);
 
     // For each element in the tuple, extract the proper strided slice.
-    auto basis = computeStrides(sliceDimCounts);
     TupleType tupleType = op.getSourceTupleType();
     int64_t tupleSize = tupleType.size();
-    SmallVector<Value, 4> tupleValues(tupleSize);
+    auto sliceStrides = makeSliceStrides(shape, sizes);
     for (int64_t i = 0; i < tupleSize; ++i) {
-      // De-linearize w.r.t. 'basis'.
-      auto vectorOffsets = delinearize(i, basis);
-      // Convert from unrolled vector-space offsets to element-space offsets.
-      auto elementOffsets = mlir::functional::zipMap(
-          [](int64_t v1, int64_t v2) { return v1 * v2; }, vectorOffsets, sizes);
-      // Compute the size of each slice.
-      SmallVector<int64_t, 4> sliceSizes(rank);
-      for (int64_t r = 0; r < rank; ++r)
-        sliceSizes[r] = std::min(sizes[r], shape[r] - elementOffsets[r]);
+      auto vectorOffsets = delinearize(sliceStrides, i);
+      auto elementOffsets = toElementSpace(sizes, vectorOffsets);
       // Extract from tuple into the result.
       auto index = rewriter.getI64IntegerAttr(i);
       auto tupleGet = rewriter.create<vector::TupleGetOp>(
           loc, tupleType.getType(i), op.getOperand(), index);
       result = rewriter.create<vector::InsertStridedSliceOp>(
           loc, tupleGet, result, elementOffsets, strides);
     }
 

mlir/lib/Dialect/VectorOps/VectorUtils.cpp

-//===- VectorAnalysis.cpp - Analysis for Vectorization --------------------===//
+//===- VectorUtils.cpp - MLIR Utilities for VectorOps   ------------------===//
 //
-// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// Part of the MLIR Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
+//
+// This file implements utility methods for working with the VectorOps dialect.
+//
+//===----------------------------------------------------------------------===//
 
-#include "mlir/Analysis/AffineAnalysis.h"
+#include "mlir/Dialect/VectorOps/VectorUtils.h"
 #include "mlir/Analysis/LoopAnalysis.h"
 #include "mlir/Dialect/AffineOps/AffineOps.h"
 #include "mlir/Dialect/StandardOps/Ops.h"
 #include "mlir/Dialect/VectorOps/VectorOps.h"
-#include "mlir/Dialect/VectorOps/VectorUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Support/Functional.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Support/MathExtras.h"
 #include "mlir/Support/STLExtras.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 
-///
-/// Implements Analysis functions specific to vectors which support
-/// the vectorization and vectorization materialization passes.
-///
+using llvm::SetVector;
 
-using namespace mlir;
+namespace mlir {
 
-using llvm::SetVector;
+SmallVector<int64_t, 4> makeSliceStrides(ArrayRef<int64_t> shape,
+                                         ArrayRef<int64_t> sizes) {
+  int64_t rank = shape.size();
+  // Compute the count for each dimension.
+  SmallVector<int64_t, 4> sliceDimCounts(rank);
+  for (int64_t r = 0; r < rank; ++r)
+    sliceDimCounts[r] = ceilDiv(shape[r], sizes[r]);
+  // Use that to compute the slice stride for each dimension.
+  SmallVector<int64_t, 4> sliceStrides(rank);
+  sliceStrides[rank - 1] = 1;
+  for (int64_t r = rank - 2; r >= 0; --r)
+    sliceStrides[r] = sliceStrides[r + 1] * sliceDimCounts[r + 1];
+  return sliceStrides;
+}
+
+SmallVector<int64_t, 4> delinearize(ArrayRef<int64_t> sliceStrides,
+                                    int64_t index) {
+  int64_t rank = sliceStrides.size();
+  SmallVector<int64_t, 4> vectorOffsets(rank);
+  for (int64_t r = 0; r < rank; ++r) {
+    assert(sliceStrides[r] > 0);
+    vectorOffsets[r] = index / sliceStrides[r];
+    index %= sliceStrides[r];
+  }
+  return vectorOffsets;
+}
+
+SmallVector<int64_t, 4> toElementSpace(ArrayRef<int64_t> sizes,
+                                       ArrayRef<int64_t> vectorOffsets) {
+  return functional::zipMap([](int64_t v1, int64_t v2) { return v1 * v2; },
+                            vectorOffsets, sizes);
+}
+
+SmallVector<int64_t, 4> makeSliceSizes(ArrayRef<int64_t> shape,
+                                       ArrayRef<int64_t> sizes,
+                                       ArrayRef<int64_t> elementOffsets) {
+  int64_t rank = shape.size();
+  SmallVector<int64_t, 4> sliceSizes(rank);
+  for (unsigned r = 0; r < rank; ++r)
+    sliceSizes[r] = std::min(sizes[r], shape[r] - elementOffsets[r]);
+  return sliceSizes;
+}
 
-Optional<SmallVector<int64_t, 4>> mlir::shapeRatio(ArrayRef<int64_t> superShape,
-                                                   ArrayRef<int64_t> subShape) {
+Optional<SmallVector<int64_t, 4>> shapeRatio(ArrayRef<int64_t> superShape,
   if (superShape.size() < subShape.size()) {
     return Optional<SmallVector<int64_t, 4>>();
   }
 
   // Starting from the end, compute the integer divisors.
   // Set the boolean `divides` if integral division is not possible.
   std::vector<int64_t> result;
   result.reserve(superShape.size());
Optional<SmallVector<int64_t, 4>> shapeRatio(ArrayRef<int64_t> superShape,
 
   assert(result.size() == superShape.size() &&
          "super to sub shape ratio is not of the same size as the super rank");
 
   // Reverse again to get it back in the proper order and return.
   return SmallVector<int64_t, 4>{result.rbegin(), result.rend()};
 }
 
-Optional<SmallVector<int64_t, 4>> mlir::shapeRatio(VectorType superVectorType,
-                                                   VectorType subVectorType) {
+Optional<SmallVector<int64_t, 4>> shapeRatio(VectorType superVectorType,
   assert(superVectorType.getElementType() == subVectorType.getElementType() &&
          "vector types must be of the same elemental type");
   return shapeRatio(superVectorType.getShape(), subVectorType.getShape());
 }
 
 /// Constructs a permutation map from memref indices to vector dimension.
 ///
 /// The implementation uses the knowledge of the mapping of enclosing loop to
static SetVector<Operation *> getParentsOfType(Operation *op) {
   return res;
 }
 
 /// Returns the enclosing AffineForOp, from closest to farthest.
 static SetVector<Operation *> getEnclosingforOps(Operation *op) {
   return getParentsOfType<AffineForOp>(op);
 }
 
-AffineMap mlir::makePermutationMap(
-    Operation *op, ArrayRef<Value> indices,
+AffineMap
+makePermutationMap(Operation *op, ArrayRef<Value> indices,
                    const DenseMap<Operation *, unsigned> &loopToVectorDim) {
   DenseMap<Operation *, unsigned> enclosingLoopToVectorDim;
   auto enclosingLoops = getEnclosingforOps(op);
   for (auto *forInst : enclosingLoops) {
     auto it = loopToVectorDim.find(forInst);
     if (it != loopToVectorDim.end()) {
       enclosingLoopToVectorDim.insert(*it);
     }
   }
-  return ::makePermutationMap(indices, enclosingLoopToVectorDim);
+  return makePermutationMap(indices, enclosingLoopToVectorDim);
 }
 
-bool mlir::matcher::operatesOnSuperVectorsOf(Operation &op,
-                                             VectorType subVectorType) {
+bool matcher::operatesOnSuperVectorsOf(Operation &op,
   // First, extract the vector type and distinguish between:
   //   a. ops that *must* lower a super-vector (i.e. vector.transfer_read,
   //      vector.transfer_write); and
   //   b. ops that *may* lower a super-vector (all other ops).
   // The ops that *may* lower a super-vector only do so if the super-vector to
   // sub-vector ratio exists. The ops that *must* lower a super-vector are
   // explicitly checked for this property.
   /// TODO(ntv): there should be a single function for all ops to do this so we
bool matcher::operatesOnSuperVectorsOf(Operation &op,
   // the vector type (but we would have to look at the compute and distinguish
   // between parallel, reduction and possibly other cases.
   if (!ratio.hasValue()) {
     return false;
   }
 
   return true;
 }
+
+} // namespace mlir