[mlir][MemRef] Move narrow type emulation common methods to MemRefUtils.

[hanchung]

Summary

It also unifies the computation of StridedLayoutAttr. If the stride is
static known value, we can just use it.

[hanchung]

This addresses the refactoring part about my review comment in https://reviews.llvm.org/D154178

[mravishankar]

Thanks for the refactoring.

[hanchung]

Add a MemRefUtils target to avoid cyclic deps

[hanchung]

fix bazel deps

[nicolasvasilache]

Great, thanks for surfacing this cleanly, this looks useful!

[nicolasvasilache]

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

50

Didn't look deeply at this method yet but have you tried using/extending utils from mlir/include/mlir/Dialect/Utils/IndexingUtils.h ?

Given that I am seeing a lot of linearize / delinearize business in there I thought I would surface.

Bonus points to refactoring this enough that it could become a new helper directly in IndexingUtils.h, but this may be a bridge too far, haven't thought about this deeply enough yet.

mlir/include/mlir/Dialect/MemRef/Utils/MemRefUtils.h

 // either by passes, optimization sequences, or in turn by other transformation
 // utilities.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H
 #define MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H
 
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+
 namespace mlir {
 
 class MemRefType;
 
 namespace memref {
 
 /// Returns true, if the memref type has static shapes and represents a
 /// contiguous chunk of memory.
 bool isStaticShapeAndContiguousRowMajor(MemRefType type);
 
+/// Returns the flattened 1-D memref and linearized offset for narrow type
+/// emulation.
+///
+/// The emulation only works on 1D memref types. To make this work on N-D
+/// memref, we need to linearize the offset.
+///
+/// For example, to emulate i4 to i8, the following op:
+///
+/// %0 = memref.load %arg0[%v0, %v1] :
+///                  memref<?x?xi4, strided<[?, ?], offset: ?>>
+///
+/// can be replaced with
+///
+/// %b, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %0
+///
+/// %linearized_offset = %v0 * %stride#0 + %v1 * %stride#1
+/// %linearized_size = %size0 * %size1
+/// %scaled_linear_offset = %linearized_offset / 8 * 4
+/// %scaled_base_offset = %offset / 8 * 4
+///
+/// %linearized = memref.reinterpret_cast %b, offset = [%scaled_base_offset],
+///                      sizes = [%linearized_size], strides = [%stride#1]
+///
+/// %new_load = memref.load %linearized[%scaled_linear_offset] :
+///                         memref<?xi8, strided<[?], offset: ?>>
+std::pair<Value, Value>
+getLinearizeMemRefAndOffset(Location loc, MemRefType sourceType, int srcBits,
+                            int dstBits, SmallVector<Value> indices,
+                            memref::ExtractStridedMetadataOp stridedMetadata,
+                            OpBuilder &builder);
+
 } // namespace memref
 } // namespace mlir
 
 #endif // MLIR_DIALECT_MEMREF_UTILS_MEMREFUTILS_H

mlir/lib/Dialect/Bufferization/IR/BufferizationOps.cpp

 //===----------------------------------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Matchers.h"
 #include <optional>
 
 using namespace mlir;
 using namespace mlir::bufferization;
 

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 //===----------------------------------------------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
-#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"

mlir/lib/Dialect/MemRef/Transforms/EmulateNarrowType.cpp

 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Transforms/NarrowTypeEmulationConverter.h"
 #include "mlir/Dialect/Arith/Transforms/Passes.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Transforms/Passes.h"
 #include "mlir/Dialect/MemRef/Transforms/Transforms.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MathExtras.h"
 #include <cassert>
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Utility functions
 //===----------------------------------------------------------------------===//
 
-/// The emulation only works on 1D memref types.
-/// To make this work on N-D memref, we need to linearize the offset.
-///
-/// For example, to emulate i4 to i8, the following op:
-///
-/// %0 = memref.load %arg0[%v0, %v1] :
-///                  memref<?x?xi4, strided<[?, ?], offset: ?>>
-///
-/// can be replaced with
-///
-/// %b, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %0
-///
-/// %linearized_offset = %v0 * %stride#0 + %v1 * %stride#1
-/// %linearized_size = %size0 * %size1
-/// %scaled_linear_offset = %linearized_offset / 8 * 4
-/// %scaled_base_offset = %offset / 8 * 4
-///
-/// %linearized = memref.reinterpret_cast %b, offset = [%scaled_base_offset],
-///                      sizes = [%linearized_size], strides = [%stride#1]
-///
-/// %new_load = memref.load %linearized[%scaled_linear_offset] :
-///                         memref<?xi8, strided<[?], offset: ?>>
-
-static Value
-linearizeMemrefLoad(Location loc, MemRefType sourceType, int srcBits,
-                    int dstBits, SmallVector<Value> indices,
-                    memref::ExtractStridedMetadataOp stridedMetadata,
-                    OpBuilder &builder) {
-  auto srcElementType = sourceType.getElementType();
-  unsigned sourceRank = indices.size();
-
-  Value baseBuffer = stridedMetadata.getBaseBuffer();
-  SmallVector<Value> baseSizes = stridedMetadata.getSizes();
-  SmallVector<Value> baseStrides = stridedMetadata.getStrides();
-  Value baseOffset = stridedMetadata.getOffset();
-  assert(indices.size() == baseStrides.size());
-
-  // Create the affine symbols and values for linearization.
-  SmallVector<AffineExpr> symbols(2 * sourceRank + 2);
-  bindSymbolsList(builder.getContext(), MutableArrayRef{symbols});
-  symbols[0] = builder.getAffineSymbolExpr(0);
-  AffineExpr addMulMap = symbols.front();
-  AffineExpr mulMap = symbols.front();
-
-  SmallVector<OpFoldResult> offsetValues(2 * sourceRank + 2);
-  offsetValues[0] = builder.getIndexAttr(0);
-  SmallVector<OpFoldResult> sizeValues(sourceRank + 1);
-  sizeValues[0] = builder.getIndexAttr(1);
-
-  for (unsigned i = 0; i < sourceRank; ++i) {
-    unsigned offsetIdx = 2 * i + 1;
-    addMulMap = addMulMap + symbols[offsetIdx] * symbols[offsetIdx + 1];
-    offsetValues[offsetIdx] = indices[i];
-    offsetValues[offsetIdx + 1] = baseStrides[i];
-
-    unsigned sizeIdx = i + 1;
-    mulMap = mulMap * symbols[sizeIdx];
-    sizeValues[sizeIdx] = baseSizes[i];
-  }
-
-  // Adjust linearizedOffset by the scale factor (dstBits / srcBits).
-  OpFoldResult scaler = builder.getIndexAttr(dstBits / srcBits);
-  AffineExpr scaledAddMulMap = addMulMap.floorDiv(symbols.back());
-  offsetValues.back() = scaler;
-
-  OpFoldResult linearizedOffset = affine::makeComposedFoldedAffineApply(
-      builder, loc, scaledAddMulMap, offsetValues);
-  OpFoldResult linearizedSize =
-      affine::makeComposedFoldedAffineApply(builder, loc, mulMap, sizeValues);
-
-  // Adjust baseOffset by the scale factor (dstBits / srcBits).
-  AffineExpr s0, s1;
-  bindSymbols(builder.getContext(), s0, s1);
-  OpFoldResult adjustBaseOffset = affine::makeComposedFoldedAffineApply(
-      builder, loc, s0.floorDiv(s1), {baseOffset, scaler});
-
-  // Flatten n-D MemRef to 1-D MemRef.
-  auto layoutAttr = StridedLayoutAttr::get(
-      sourceType.getContext(), ShapedType::kDynamic, {ShapedType::kDynamic});
-  int64_t staticShape = sourceType.hasStaticShape()
-                            ? sourceType.getNumElements()
-                            : ShapedType::kDynamic;
-  auto flattenMemrefType = MemRefType::get(
-      staticShape, srcElementType, layoutAttr, sourceType.getMemorySpace());
-
-  auto reinterpret = builder.create<memref::ReinterpretCastOp>(
-      loc, flattenMemrefType, baseBuffer,
-      getValueOrCreateConstantIndexOp(builder, loc, adjustBaseOffset),
-      getValueOrCreateConstantIndexOp(builder, loc, linearizedSize),
-      baseStrides.back());
-
-  return builder.create<memref::LoadOp>(
-      loc, srcElementType, reinterpret.getResult(),
-      getValueOrCreateConstantIndexOp(builder, loc, linearizedOffset));
-}
-
 /// When data is loaded/stored in `targetBits` granularity, but is used in
 /// `sourceBits` granularity (`sourceBits` < `targetBits`), the `targetBits` is
 /// treated as an array of elements of width `sourceBits`.
 /// Return the bit offset of the value at position `srcIdx`. For example, if
 /// `sourceBits` equals to 4 and `targetBits` equals to 8, the x-th element is
 /// located at (x % 2) * 4. Because there are two elements in one i8, and one
 /// element has 4 bits.
 static Value getOffsetForBitwidth(Location loc, Value srcIdx, int sourceBits,
matchAndRewrite(memref::LoadOp op, OpAdaptor adaptor,
 
     Value newLoad, lastIdx;
     if (sourceRank == 0) {
       newLoad = rewriter.create<memref::LoadOp>(
           loc, srcElementType, adaptor.getMemref(), adaptor.getIndices());
 
       lastIdx = stridedMetadata.getOffset();
     } else {
-      newLoad = linearizeMemrefLoad(loc, sourceType, srcBits, dstBits, indices,
+      auto [reinterpret, linearizedOffset] =
+          memref::getLinearizeMemRefAndOffset(loc, sourceType, srcBits, dstBits,
+                                              adaptor.getIndices(),
                                               stridedMetadata, rewriter);
 
+      newLoad = rewriter.create<memref::LoadOp>(loc, srcElementType,
+                                                reinterpret, linearizedOffset);
+
       lastIdx = adaptor.getIndices().back();
     }
 
     // Get the offset and shift the bits to the rightmost.
     // Note, currently only the big-endian is supported.
     auto castLastIdx =
         rewriter.create<arith::IndexCastUIOp>(loc, srcElementType, lastIdx);
 

mlir/lib/Dialect/MemRef/Utils/CMakeLists.txt

 add_mlir_dialect_library(MLIRMemRefUtils
   MemRefUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/inlude/mlir/Dialect/MemRef/Utils
 
   LINK_LIBS PUBLIC
   MLIRIR
+  MLIRAffineDialect
+  MLIRArithUtils
 )

mlir/lib/Dialect/MemRef/Utils/MemRefUtils.cpp

 //===- MemRefUtils.cpp - Utilities to support the MemRef dialect ----------===//
 //
 // 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 utilities for the MemRef dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 
 namespace mlir {
 namespace memref {
 
 bool isStaticShapeAndContiguousRowMajor(MemRefType type) {
   if (!type.hasStaticShape())
     return false;
bool isStaticShapeAndContiguousRowMajor(MemRefType type) {
   while (curDim >= 0 && type.getDimSize(curDim) == 1) {
     --curDim;
   }
 
   // All dims are unit-strided or size-1.
   return curDim < 0;
 }
 
+std::pair<Value, Value>
+getLinearizeMemRefAndOffset(Location loc, MemRefType sourceType, int srcBits,
+                            int dstBits, SmallVector<Value> indices,
+                            memref::ExtractStridedMetadataOp stridedMetadata,
+                            OpBuilder &builder) {
+  auto srcElementType = sourceType.getElementType();
+  unsigned sourceRank = indices.size();
+
+  Value baseBuffer = stridedMetadata.getBaseBuffer();
+  SmallVector<Value> baseSizes = stridedMetadata.getSizes();
+  SmallVector<Value> baseStrides = stridedMetadata.getStrides();
+  Value baseOffset = stridedMetadata.getOffset();
+  assert(indices.size() == baseStrides.size());
+
+  // Create the affine symbols and values for linearization.
+  SmallVector<AffineExpr> symbols(2 * sourceRank + 2);
+  bindSymbolsList(builder.getContext(), MutableArrayRef{symbols});
+  symbols[0] = builder.getAffineSymbolExpr(0);
+  AffineExpr addMulMap = symbols.front();
+  AffineExpr mulMap = symbols.front();
+
+  SmallVector<OpFoldResult> offsetValues(2 * sourceRank + 2);
+  offsetValues[0] = builder.getIndexAttr(0);
+  SmallVector<OpFoldResult> sizeValues(sourceRank + 1);
+  sizeValues[0] = builder.getIndexAttr(1);
+
+  for (unsigned i = 0; i < sourceRank; ++i) {
+    unsigned offsetIdx = 2 * i + 1;
+    addMulMap = addMulMap + symbols[offsetIdx] * symbols[offsetIdx + 1];
+    offsetValues[offsetIdx] = indices[i];
+    offsetValues[offsetIdx + 1] = baseStrides[i];
+
+    unsigned sizeIdx = i + 1;
+    mulMap = mulMap * symbols[sizeIdx];
+    sizeValues[sizeIdx] = baseSizes[i];
+  }
+
+  // Adjust linearizedOffset by the scale factor (dstBits / srcBits).
+  OpFoldResult scaler = builder.getIndexAttr(dstBits / srcBits);
+  AffineExpr scaledAddMulMap = addMulMap.floorDiv(symbols.back());
+  offsetValues.back() = scaler;
+
+  OpFoldResult linearizedOffset = affine::makeComposedFoldedAffineApply(
+      builder, loc, scaledAddMulMap, offsetValues);
+  OpFoldResult linearizedSize =
+      affine::makeComposedFoldedAffineApply(builder, loc, mulMap, sizeValues);
+
+  // Adjust baseOffset by the scale factor (dstBits / srcBits).
+  AffineExpr s0, s1;
+  bindSymbols(builder.getContext(), s0, s1);
+  OpFoldResult adjustBaseOffset = affine::makeComposedFoldedAffineApply(
+      builder, loc, s0.floorDiv(s1), {baseOffset, scaler});
+
+  // Flatten n-D MemRef to 1-D MemRef.
+  std::optional<int64_t> stride =
+      getConstantIntValue(stridedMetadata.getConstifiedMixedStrides().back());
+  auto layoutAttr =
+      StridedLayoutAttr::get(sourceType.getContext(), ShapedType::kDynamic,
+                             {stride ? stride.value() : ShapedType::kDynamic});
+  int64_t staticShape = sourceType.hasStaticShape()
+                            ? sourceType.getNumElements()
+                            : ShapedType::kDynamic;
+  auto flattenMemrefType = MemRefType::get(
+      staticShape, srcElementType, layoutAttr, sourceType.getMemorySpace());
+
+  auto reinterpret = builder.create<memref::ReinterpretCastOp>(
+      loc, flattenMemrefType, baseBuffer,
+      getValueOrCreateConstantIndexOp(builder, loc, adjustBaseOffset),
+      getValueOrCreateConstantIndexOp(builder, loc, linearizedSize),
+      baseStrides.back());
+
+  return std::make_pair(reinterpret, getValueOrCreateConstantIndexOp(
+                                         builder, loc, linearizedOffset));
+}
+
 } // namespace memref
 } // namespace mlir

mlir/lib/Dialect/Vector/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRVectorTransforms
   MLIRArithDialect
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
   MLIRDialectUtils
   MLIRGPUDialect
   MLIRIR
   MLIRLinalgDialect
   MLIRMemRefDialect
+  MLIRMemRefUtils
   MLIRSCFDialect
   MLIRSideEffectInterfaces
   MLIRTensorDialect
   MLIRTransforms
   MLIRVectorDialect
   MLIRVectorInterfaces
   MLIRVectorUtils
   )

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

 //===- VectorEmulateNarrowType.cpp - Narrow type emulation ----*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Arith/Transforms/NarrowTypeEmulationConverter.h"
 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/MathExtras.h"
 #include <cassert>
 
 using namespace mlir;
 
-//===----------------------------------------------------------------------===//
-// Utility functions
-//===----------------------------------------------------------------------===//
-
-/// The emulation only works on 1D memref types.
-/// To make this work on N-D memref, we need to linearize the offset.
-///
-/// For example, to emulate i4 to i8, the following op:
-///
-/// %0 = memref.load %arg0[%v0, %v1] :
-///                  memref<?x?xi4, strided<[?, ?], offset: ?>>
-///
-/// can be replaced with
-///
-/// %b, %offset, %sizes:2, %strides:2 = memref.extract_strided_metadata %0
-///
-/// %linearized_offset = %v0 * %stride#0 + %v1 * %stride#1
-/// %linearized_size = %size0 * %size1
-/// %scaled_linear_offset = %linearized_offset / 8 * 4
-/// %scaled_base_offset = %offset / 8 * 4
-///
-/// %linearized = memref.reinterpret_cast %b, offset = [%scaled_base_offset],
-///                      sizes = [%linearized_size], strides = [%stride#1]
-///
-/// %new_load = vector.load %linearized[%scaled_linear_offset] :
-///                         memref<?xi8, strided<[?], offset: ?>>
-
-static Value
-linearizeVectorLoad(Location loc, MemRefType sourceType, int srcBits,
-                    int dstBits, SmallVector<Value> indices,
-                    memref::ExtractStridedMetadataOp stridedMetadata,
-                    int numElements, OpBuilder &builder) {
-  auto srcElementType = sourceType.getElementType();
-  unsigned sourceRank = indices.size();
-
-  Value baseBuffer = stridedMetadata.getBaseBuffer();
-  SmallVector<Value> baseSizes = stridedMetadata.getSizes();
-  SmallVector<Value> baseStrides = stridedMetadata.getStrides();
-  Value baseOffset = stridedMetadata.getOffset();
-  assert(indices.size() == baseStrides.size());
-
-  // Create the affine symbols and values for linearization.
-  SmallVector<AffineExpr> symbols(2 * sourceRank + 2);
-  bindSymbolsList(builder.getContext(), MutableArrayRef{symbols});
-  symbols[0] = builder.getAffineSymbolExpr(0);
-  AffineExpr addMulMap = symbols.front();
-  AffineExpr mulMap = symbols.front();
-
-  SmallVector<OpFoldResult> offsetValues(2 * sourceRank + 2);
-  offsetValues[0] = builder.getIndexAttr(0);
-  SmallVector<OpFoldResult> sizeValues(sourceRank + 1);
-  sizeValues[0] = builder.getIndexAttr(1);
-
-  for (unsigned i = 0; i < sourceRank; ++i) {
-    unsigned offsetIdx = 2 * i + 1;
-    addMulMap = addMulMap + symbols[offsetIdx] * symbols[offsetIdx + 1];
-    offsetValues[offsetIdx] = indices[i];
-    offsetValues[offsetIdx + 1] = baseStrides[i];
-
-    unsigned sizeIdx = i + 1;
-    mulMap = mulMap * symbols[sizeIdx];
-    sizeValues[sizeIdx] = baseSizes[i];
-  }
-
-  // Adjust linearizedOffset by the scale factor (dstBits / srcBits).
-  OpFoldResult scaler = builder.getIndexAttr(dstBits / srcBits);
-  AffineExpr scaledAddMulMap = addMulMap.floorDiv(symbols.back());
-  offsetValues.back() = scaler;
-
-  OpFoldResult linearizedOffset = affine::makeComposedFoldedAffineApply(
-      builder, loc, scaledAddMulMap, offsetValues);
-  OpFoldResult linearizedSize =
-      affine::makeComposedFoldedAffineApply(builder, loc, mulMap, sizeValues);
-
-  // Adjust baseOffset by the scale factor (dstBits / srcBits).
-  AffineExpr s0, s1;
-  bindSymbols(builder.getContext(), s0, s1);
-  OpFoldResult adjustBaseOffset = affine::makeComposedFoldedAffineApply(
-      builder, loc, s0.floorDiv(s1), {baseOffset, scaler});
-
-  // Flatten n-D MemRef to 1-D MemRef.
-  auto layoutAttr = StridedLayoutAttr::get(sourceType.getContext(),
-                                           ShapedType::kDynamic, {1});
-  int64_t staticShape = sourceType.hasStaticShape()
-                            ? sourceType.getNumElements()
-                            : ShapedType::kDynamic;
-  auto flattenMemrefType = MemRefType::get(
-      staticShape, srcElementType, layoutAttr, sourceType.getMemorySpace());
-
-  auto reinterpret = builder.create<memref::ReinterpretCastOp>(
-      loc, flattenMemrefType, baseBuffer,
-      getValueOrCreateConstantIndexOp(builder, loc, adjustBaseOffset),
-      getValueOrCreateConstantIndexOp(builder, loc, linearizedSize),
-      baseStrides.back());
-
-  return builder.create<vector::LoadOp>(
-      loc, VectorType::get(numElements, srcElementType),
-      reinterpret.getResult(),
-      getValueOrCreateConstantIndexOp(builder, loc, linearizedOffset));
-}
-
 namespace {
 
 //===----------------------------------------------------------------------===//
 // ConvertVectorLoad
 //===----------------------------------------------------------------------===//
 
 struct ConvertVectorLoad final : OpConversionPattern<vector::LoadOp> {
   using OpConversionPattern::OpConversionPattern;
matchAndRewrite(vector::LoadOp op, OpAdaptor adaptor,
 
     auto origElements = op.getVectorType().getNumElements();
     if (origElements % scale != 0)
       return failure();
 
     auto stridedMetadata = rewriter.create<memref::ExtractStridedMetadataOp>(
         loc, adaptor.getBase());
 
-    auto numElements = int(std::ceil(double(origElements) / scale));
-    auto newLoad = linearizeVectorLoad(loc, sourceType, srcBits, dstBits,
-                                       adaptor.getIndices(), stridedMetadata,
-                                       numElements, rewriter);
+    auto [reinterpret, linearizedOffset] = memref::getLinearizeMemRefAndOffset(
+        loc, sourceType, srcBits, dstBits, adaptor.getIndices(),
+        stridedMetadata, rewriter);
+
+    auto srcElementType = sourceType.getElementType();
+    auto numElements =
+        static_cast<int>(std::ceil(static_cast<double>(origElements) / scale));
+    auto newLoad = rewriter.create<vector::LoadOp>(
+        loc, VectorType::get(numElements, srcElementType), reinterpret,
+        linearizedOffset);
 
     numElements *= scale;
     auto castType = VectorType::get(numElements, oldElementType);
     auto bitCast = rewriter.create<vector::BitCastOp>(loc, castType, newLoad);
 
     rewriter.replaceOp(op, bitCast->getResult(0));
     return success();
   }

mlir/test/Dialect/MemRef/emulate-narrow-type-diff-load-compute.mlir

 // -----
 
 // CHECK-LABEL: func @memref_load_i4
 // CHECK-SAME:    (%[[ARG:.*]]: index)
 // CHECK-NEXT:    %[[M:.*]]  = memref.alloc() : memref<4xi8>
 // CHECK-NEXT:    %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]], %[[STRIDES:.*]] = memref.extract_strided_metadata %[[M]] : memref<4xi8> -> memref<i8>, index, index, index
 // CHECK-NEXT:    %[[INDEX:.*]] = affine.apply #[[$MAP0]]()[%[[ARG]], %[[STRIDES]]]
 // CHECK-NEXT:    %[[AOFF:.*]] = affine.apply #[[$MAP1]]()[%[[OFFSET]]]
-// CHECK-NEXT:    %[[CAST:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[AOFF]]], sizes: [%[[SIZES]]], strides: [%[[STRIDES]]] : memref<i8> to memref<4xi8, strided<[?], offset: ?>>
-// CHECK-NEXT:    %[[LOAD:.*]] = memref.load %[[CAST]][%[[INDEX]]] : memref<4xi8, strided<[?], offset: ?>>
+// CHECK-NEXT:    %[[CAST:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[AOFF]]], sizes: [%[[SIZES]]], strides: [%[[STRIDES]]] : memref<i8> to memref<4xi8, strided<[1], offset: ?>>
+// CHECK-NEXT:    %[[LOAD:.*]] = memref.load %[[CAST]][%[[INDEX]]] : memref<4xi8, strided<[1], offset: ?>>
 // CHECK-NEXT:    %[[I:.*]] = arith.index_castui %[[ARG]] : index to i8
 // CHECK-NEXT:    %[[C2:.*]] = arith.constant 2 : i8
 // CHECK-NEXT:    %[[C4:.*]] = arith.constant 4 : i8
 // CHECK-NEXT:    %[[REM:.*]] = arith.remui %[[I]], %[[C2]] : i8
 // CHECK-NEXT:    %[[STEP:.*]] = arith.muli %[[REM]], %[[C4]] : i8
 // CHECK-NEXT:    %[[SHIFT:.*]] = arith.shrsi %[[LOAD]], %[[STEP]] : i8
 // CHECK-NEXT:    %[[RES:.*]] = arith.trunci %[[SHIFT]] : i8 to i4
 // CHECK-NEXT:    return
 func.func @memref_load_i4(%arg0: index) {
     %0 = memref.alloc() : memref<4xi4>
     %1 = memref.load %0[%arg0] : memref<4xi4>
     return
 }
 
 // -----
 
 // CHECK-LABEL: func @memref_load_i4_rank2
 // CHECK-SAME:    (%[[ARG:.*]]: memref<4x128xi8>, %[[ARG0:.*]]: index, %[[ARG1:.*]]: index)
 // CHECK-NEXT:    memref.assume_alignment %[[ARG]], 64 : memref<4x128xi8>
 // CHECK-NEXT:    %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[ARG]] : memref<4x128xi8> -> memref<i8>, index, index, index, index, index
 // CHECK-NEXT:    %[[INDEX:.*]] = affine.apply #[[$MAP2]]()[%[[ARG0]], %[[STRIDES]]#0, %[[ARG1]], %[[STRIDES]]#1]
 // CHECK-NEXT:    %[[LSIZE:.*]] = affine.apply #[[$MAP3]]()[%[[SIZES]]#0, %[[SIZES]]#1]
 // CHECK-NEXT:    %[[AOFF:.*]] = affine.apply #[[$MAP1]]()[%[[OFFSET]]]
-// CHECK-NEXT:    %[[CAST:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[AOFF]]], sizes: [%[[LSIZE]]], strides: [%[[STRIDES]]#1] : memref<i8> to memref<512xi8, strided<[?], offset: ?>>
-// CHECK-NEXT:    %[[LOAD:.*]] = memref.load %[[CAST]][%[[INDEX]]] : memref<512xi8, strided<[?], offset: ?>>
+// CHECK-NEXT:    %[[CAST:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[AOFF]]], sizes: [%[[LSIZE]]], strides: [%[[STRIDES]]#1] : memref<i8> to memref<512xi8, strided<[1], offset: ?>>
+// CHECK-NEXT:    %[[LOAD:.*]] = memref.load %[[CAST]][%[[INDEX]]] : memref<512xi8, strided<[1], offset: ?>>
 // CHECK-NEXT:    %[[I:.*]] = arith.index_castui %[[ARG1]] : index to i8
 // CHECK-NEXT:    %[[C2:.*]] = arith.constant 2 : i8
 // CHECK-NEXT:    %[[C4:.*]] = arith.constant 4 : i8
 // CHECK-NEXT:    %[[REM:.*]] = arith.remui %[[I]], %[[C2]] : i8
 // CHECK-NEXT:    %[[STEP:.*]] = arith.muli %[[REM]], %[[C4]] : i8
 // CHECK-NEXT:    %[[SHIFT:.*]] = arith.shrsi %[[LOAD]], %[[STEP]] : i8
 // CHECK-NEXT:    %[[RES:.*]] = arith.trunci %[[SHIFT]] : i8 to i4
 // CHECK-NEXT:    return
 func.func @memref_load_i4_rank2(%0: memref<4x128xi4>, %arg0: index, %arg1: index) {
     memref.assume_alignment %0, 64 : memref<4x128xi4>
     %1 = memref.load %0[%arg0,%arg1] : memref<4x128xi4>
     return
 }

mlir/test/Dialect/MemRef/emulate-narrow-type-same-load-compute.mlir

 // -----
 
 // CHECK-LABEL: func @memref_load_i4
 // CHECK-SAME:    (%[[ARG:.*]]: index)
 // CHECK-NEXT:    %[[M:.*]]  = memref.alloc() : memref<4xi8>
 // CHECK-NEXT:    %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]], %[[STRIDES:.*]] = memref.extract_strided_metadata %[[M]] : memref<4xi8> -> memref<i8>, index, index, index
 // CHECK-NEXT:    %[[INDEX:.*]] = affine.apply #[[$MAP0]]()[%[[ARG]], %[[STRIDES]]]
 // CHECK-NEXT:    %[[AOFF:.*]] = affine.apply #[[$MAP1]]()[%[[OFFSET]]]
-// CHECK-NEXT:    %[[CAST:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[AOFF]]], sizes: [%[[SIZES]]], strides: [%[[STRIDES]]] : memref<i8> to memref<4xi8, strided<[?], offset: ?>>
-// CHECK-NEXT:    %[[LOAD:.*]] = memref.load %[[CAST]][%[[INDEX]]] : memref<4xi8, strided<[?], offset: ?>>
+// CHECK-NEXT:    %[[CAST:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[AOFF]]], sizes: [%[[SIZES]]], strides: [%[[STRIDES]]] : memref<i8> to memref<4xi8, strided<[1], offset: ?>>
+// CHECK-NEXT:    %[[LOAD:.*]] = memref.load %[[CAST]][%[[INDEX]]] : memref<4xi8, strided<[1], offset: ?>>
 // CHECK-NEXT:    %[[I:.*]] = arith.index_castui %[[ARG]] : index to i8
 // CHECK-NEXT:    %[[C2:.*]] = arith.constant 2 : i8
 // CHECK-NEXT:    %[[C4:.*]] = arith.constant 4 : i8
 // CHECK-NEXT:    %[[REM:.*]] = arith.remui %[[I]], %[[C2]] : i8
 // CHECK-NEXT:    %[[STEP:.*]] = arith.muli %[[REM]], %[[C4]] : i8
 // CHECK-NEXT:    %[[SHIFT:.*]] = arith.shrsi %[[LOAD]], %[[STEP]] : i8
 // CHECK-NEXT:    %[[MASK:.*]] = arith.constant 15 : i8
 // CHECK-NEXT:    %[[RES:.*]] = arith.andi %[[SHIFT]], %[[MASK]] : i8
 // CHECK-NEXT:    return
 func.func @memref_load_i4(%arg0: index) {
     %0 = memref.alloc() : memref<4xi4>
     %1 = memref.load %0[%arg0] : memref<4xi4>
     return
 }
 
 // -----
 
 // CHECK-LABEL: func @memref_load_i4_rank2
 // CHECK-SAME:    (%[[ARG:.*]]: memref<4x128xi8>, %[[ARG0:.*]]: index, %[[ARG1:.*]]: index)
 // CHECK-NEXT:    memref.assume_alignment %[[ARG]], 64 : memref<4x128xi8>
 // CHECK-NEXT:    %[[BASE:.*]], %[[OFFSET:.*]], %[[SIZES:.*]]:2, %[[STRIDES:.*]]:2 = memref.extract_strided_metadata %[[ARG]] : memref<4x128xi8> -> memref<i8>, index, index, index, index, index
 // CHECK-NEXT:    %[[INDEX:.*]] = affine.apply #[[$MAP2]]()[%[[ARG0]], %[[STRIDES]]#0, %[[ARG1]], %[[STRIDES]]#1]
 // CHECK-NEXT:    %[[LSIZE:.*]] = affine.apply #[[$MAP3]]()[%[[SIZES]]#0, %[[SIZES]]#1]
 // CHECK-NEXT:    %[[AOFF:.*]] = affine.apply #[[$MAP1]]()[%[[OFFSET]]]
-// CHECK-NEXT:    %[[CAST:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[AOFF]]], sizes: [%[[LSIZE]]], strides: [%[[STRIDES]]#1] : memref<i8> to memref<512xi8, strided<[?], offset: ?>>
-// CHECK-NEXT:    %[[LOAD:.*]] = memref.load %[[CAST]][%[[INDEX]]] : memref<512xi8, strided<[?], offset: ?>>
+// CHECK-NEXT:    %[[CAST:.*]] = memref.reinterpret_cast %[[BASE]] to offset: [%[[AOFF]]], sizes: [%[[LSIZE]]], strides: [%[[STRIDES]]#1] : memref<i8> to memref<512xi8, strided<[1], offset: ?>>
+// CHECK-NEXT:    %[[LOAD:.*]] = memref.load %[[CAST]][%[[INDEX]]] : memref<512xi8, strided<[1], offset: ?>>
 // CHECK-NEXT:    %[[I:.*]] = arith.index_castui %[[ARG1]] : index to i8
 // CHECK-NEXT:    %[[C2:.*]] = arith.constant 2 : i8
 // CHECK-NEXT:    %[[C4:.*]] = arith.constant 4 : i8
 // CHECK-NEXT:    %[[REM:.*]] = arith.remui %[[I]], %[[C2]] : i8
 // CHECK-NEXT:    %[[STEP:.*]] = arith.muli %[[REM]], %[[C4]] : i8
 // CHECK-NEXT:    %[[SHIFT:.*]] = arith.shrsi %[[LOAD]], %[[STEP]] : i8
 // CHECK-NEXT:    %[[MASK:.*]] = arith.constant 15 : i8
 // CHECK-NEXT:    %[[RES:.*]] = arith.andi %[[SHIFT]], %[[MASK]] : i8
 // CHECK-NEXT:    return
 func.func @memref_load_i4_rank2(%0: memref<4x128xi4>, %arg0: index, %arg1: index) {
     memref.assume_alignment %0, 64 : memref<4x128xi4>
     %1 = memref.load %0[%arg0,%arg1] : memref<4x128xi4>
     return
 }

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":ArithUtils",
         ":BufferizationDialect",
         ":BufferizationTransforms",
         ":DialectUtils",
         ":FuncDialect",
         ":IR",
         ":LinalgDialect",
         ":MemRefDialect",
+        ":MemRefUtils",
         ":Pass",
         ":SCFDialect",
         ":SideEffectInterfaces",
         ":Support",
         ":TensorDialect",
         ":Transforms",
         ":VectorDialect",
         ":VectorEnumsIncGen",
deps = [
         ":ArithDialect",
         ":ConversionPassIncGen",
         ":DataLayoutInterfaces",
         ":FuncDialect",
         ":IR",
         ":LLVMCommonConversion",
         ":LLVMDialect",
         ":MemRefDialect",
+        ":MemRefUtils",
         ":Pass",
         ":Support",
         ":Transforms",
         "//llvm:Support",
     ],
 )
 
 cc_library(

 )
 
 cc_library(
     name = "MemRefDialect",
     srcs = glob(
         [
             "lib/Dialect/MemRef/IR/*.cpp",
             "lib/Dialect/MemRef/IR/*.h",
-            "lib/Dialect/MemRef/Utils/*.cpp",
         ],
     ),
     hdrs = [
         "include/mlir/Dialect/MemRef/IR/MemRef.h",
         "include/mlir/Dialect/MemRef/IR/MemRefMemorySlot.h",
         "include/mlir/Dialect/MemRef/IR/ValueBoundsOpInterfaceImpl.h",
-        "include/mlir/Dialect/MemRef/Utils/MemRefUtils.h",
     ],
     includes = ["include"],
     deps = [
         ":ArithDialect",
         ":ArithUtils",
         ":CastInterfaces",
         ":ComplexDialect",
         ":ControlFlowInterfaces",
         ":CopyOpInterface",
         ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":MemRefBaseIncGen",
         ":MemRefOpsIncGen",
         ":MemorySlotInterfaces",
         ":ShapedOpInterfaces",
         ":Support",
         ":ValueBoundsOpInterface",
         ":ViewLikeInterface",
         "//llvm:Support",
         "//llvm:TargetParser",
     ],
 )
 
+cc_library(
+    name = "MemRefUtils",
+    srcs = glob(
+        [
+            "lib/Dialect/MemRef/Utils/*.cpp",
+        ],
+    ),
+    hdrs = [
+        "include/mlir/Dialect/MemRef/Utils/MemRefUtils.h",
+    ],
+    includes = ["include"],
+    deps = [
+        ":AffineDialect",
+        ":ArithUtils",
+        ":MemRefDialect",
+    ],
+)
+
 gentbl_cc_library(
     name = "MemRefPassIncGen",
     strip_include_prefix = "include",
     tbl_outs = [
         (
             [
                 "-gen-pass-decls",
                 "-name=MemRef",
deps = [
         ":DialectUtils",
         ":FuncDialect",
         ":GPUDialect",
         ":IR",
         ":InferTypeOpInterface",
         ":LoopLikeInterface",
         ":MemRefDialect",
         ":MemRefPassIncGen",
+        ":MemRefUtils",
         ":NVGPUDialect",
         ":Pass",
         ":RuntimeVerifiableOpInterface",
         ":Support",
         ":TensorDialect",
         ":Transforms",
         ":ValueBoundsOpInterface",
         ":VectorDialect",
deps = [
         ":BufferizationDialect",
         ":BufferizationEnumsIncGen",
         ":BufferizationPassIncGen",
         ":ControlFlowInterfaces",
         ":FuncDialect",
         ":IR",
         ":LoopLikeInterface",
         ":MemRefDialect",
+        ":MemRefUtils",
         ":Pass",
         ":SideEffectInterfaces",
         ":TensorDialect",
         ":Transforms",
         ":ViewLikeInterface",
         "//llvm:Support",
     ],
 )