[MLIR][Affine] Fix assumption on int type in memref elt size method

[bondhugula]

Summary

Fix assumption on memref element type being int/float in memref elt size
related method and affine data copy generate.

Fixes https://github.com/llvm/llvm-project/issues/61310

[bondhugula]

mlir/include/mlir/Dialect/Affine/Analysis/Utils.h

382

@ftynse I'm not sure this should be exposed from here or if there's something else I should use or if this should be renamed to something specific.

[bondhugula]

mlir/include/mlir/Dialect/Affine/Analysis/Utils.h

382

For index-typed memrefs and for memrefs in general, at this stage, we don't have layout information. So this is really returning a tentative or a "best-effort" size, which is what the clients of this method (affine passes) desire. This is being exposed from MLIR Affine Analysis library since it has at least two users among affine passes/utilities. In general, the clients are expected to be aware that this size may not be the final size of the buffer which may be layout/backend dependent.

[bondhugula]

Rebase.

[ftynse]

mlir/include/mlir/Dialect/Affine/Analysis/Utils.h

382

If this is called from the pass, it may be possible to obtain the DataLayoutAnalysis in the pass from the analysis manager, and query it for the DataLayout object of the module. Given that, this should be able to compute the element size the same way the allocation lowering does.

Otherwise, I'd call it getMemRefIntOrFloatEltSizeInBytes, or something similar that indicates it being only suitable for ints, floats and vectors thereof.

[bondhugula]

Incorporate review suggestion.

[bondhugula]

mlir/include/mlir/Dialect/Affine/Analysis/Utils.h

382

Thanks. Renamed to getMemRefIntOrFloatEltSizeInBytes. This is used from a free-standing utility in Affine LoopUtils although it's eventually used from passes in the upstream repo. It's expected to be used from other utilities downstream.

[Lewuathe]

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

607

I'm not sure whether this approach is correct. But we assume the index type bit-width is set by kInternalStorageBitWidth, which is already used in the several parts.

https://reviews.llvm.org/D146019

Can we simply use this value for calculating the memref size as well?

[bondhugula]

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

607

Note that isIntOrFloat returns false for index types. So, we are returning nullopt (i.e., failing this method) when an provide an index typed memref is provided. It would be wrong to use a specific value since the width of index depends on the lowering target - it can be 32 bits, 64 bits or anything else.

[ftynse]

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

607

I'm not sure whether this approach is correct. But we assume the index type bit-width is set by kInternalStorageBitWidth, which is already used in the several parts.

As the name suggests, it is supposed to be used for internal storage. Please don't use it beyond that. Using the DataLayout is the right approach, and we do have platforms on which index is lowered to a type with a smaller width than the internal storage.

[bondhugula]

Rebase and add a comment.

[bondhugula]

MSAN issue/bug introduced by this fixed in https://reviews.llvm.org/D146698

mlir/include/mlir/Dialect/Affine/Analysis/Utils.h

struct MemRefRegion {
   /// order (major to minor / outermost to innermost). There may be additional
   /// variables since getMemRefRegion() is called with a specific loop depth,
   /// and thus the region is symbolic in the outer surrounding loops at that
   /// depth.
   // TODO: Replace this to exploit HyperRectangularSet.
   FlatAffineValueConstraints cst;
 };
 
-/// Returns the size of memref data in bytes if it's statically shaped,
-/// std::nullopt otherwise.
-std::optional<uint64_t> getMemRefSizeInBytes(MemRefType memRefType);
+/// Returns the size of a memref with element type int or float in bytes if it's
+/// statically shaped, std::nullopt otherwise.
+std::optional<uint64_t> getIntOrFloatMemRefSizeInBytes(MemRefType memRefType);
 
 /// Checks a load or store op for an out of bound access; returns failure if the
 /// access is out of bounds along any of the dimensions, success otherwise.
 /// Emits a diagnostic error (with location information) if emitError is true.
 template <typename LoadOrStoreOpPointer>
 LogicalResult boundCheckLoadOrStoreOp(LoadOrStoreOpPointer loadOrStoreOp,
                                       bool emitError = true);
 
 /// Returns the number of surrounding loops common to both A and B.
 unsigned getNumCommonSurroundingLoops(Operation &a, Operation &b);
 
 /// Gets the memory footprint of all data touched in the specified memory space
 /// in bytes; if the memory space is unspecified, considers all memory spaces.
 std::optional<int64_t> getMemoryFootprintBytes(AffineForOp forOp,
                                                int memorySpace = -1);
 
+/// Returns the memref's element type's size in bytes where the elemental type
+/// is an int or float or a vector of such types.
+std::optional<int64_t> getMemRefIntOrFloatEltSizeInBytes(MemRefType memRefType);
+
 /// Simplify the integer set by simplifying the underlying affine expressions by
 /// flattening and some simple inference. Also, drop any duplicate constraints.
 /// Returns the simplified integer set. This method runs in time linear in the
 /// number of constraints.
 IntegerSet simplifyIntegerSet(IntegerSet set);
 
 /// Returns the innermost common loop depth for the set of operations in 'ops'.
 unsigned getInnermostCommonLoopDepth(

mlir/include/mlir/Dialect/Affine/LoopUtils.h

 /// inserted right before `end`. `copyOptions` provides various parameters, and
 /// the output argument `copyNests` is the set of all copy nests inserted, each
 /// represented by its root affine.for. Since we generate alloc's and dealloc's
 /// for all fast buffers (before and after the range of operations resp. or at a
 /// hoisted position), all of the fast memory capacity is assumed to be
 /// available for processing this block range. When 'filterMemRef' is specified,
 /// copies are only generated for the provided MemRef. Returns success if the
 /// explicit copying succeeded for all memrefs on which affine load/stores were
-/// encountered.
+/// encountered. For memrefs for whose element types a size in bytes can't be
+/// computed (`index` type), their capacity is not accounted for and the
+/// `fastMemCapacityBytes` copy option would be non-functional in such cases.
 LogicalResult affineDataCopyGenerate(Block::iterator begin, Block::iterator end,
                                      const AffineCopyOptions &copyOptions,
                                      std::optional<Value> filterMemRef,
                                      DenseSet<Operation *> &copyNests);
 
 /// A convenience version of affineDataCopyGenerate for all ops in the body of
 /// an AffineForOp.
 LogicalResult affineDataCopyGenerate(AffineForOp forOp,

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
   }
   cst.removeTrivialRedundancy();
 
   LLVM_DEBUG(llvm::dbgs() << "Memory region:\n");
   LLVM_DEBUG(cst.dump());
   return success();
 }
 
-static unsigned getMemRefEltSizeInBytes(MemRefType memRefType) {
+std::optional<int64_t>
+mlir::getMemRefIntOrFloatEltSizeInBytes(MemRefType memRefType) {
   auto elementType = memRefType.getElementType();
 
   unsigned sizeInBits;
   if (elementType.isIntOrFloat()) {
     sizeInBits = elementType.getIntOrFloatBitWidth();
-  } else {
-    auto vectorType = elementType.cast<VectorType>();
+  } else if (auto vectorType = elementType.dyn_cast<VectorType>()) {
+    if (vectorType.getElementType().isIntOrFloat())
       sizeInBits =
           vectorType.getElementTypeBitWidth() * vectorType.getNumElements();
+    else
+      return std::nullopt;
+  } else {
+    return std::nullopt;
   }
   return llvm::divideCeil(sizeInBits, 8);
 }
 
 // Returns the size of the region.
 std::optional<int64_t> MemRefRegion::getRegionSize() {
   auto memRefType = memref.getType().cast<MemRefType>();
 
std::optional<int64_t> MemRefRegion::getRegionSize() {
   SmallVector<Value, 4> bufIndices;
 
   // Compute the extents of the buffer.
   std::optional<int64_t> numElements = getConstantBoundingSizeAndShape();
   if (!numElements) {
     LLVM_DEBUG(llvm::dbgs() << "Dynamic shapes not yet supported\n");
     return std::nullopt;
   }
-  return getMemRefEltSizeInBytes(memRefType) * *numElements;
+  auto eltSize = getMemRefIntOrFloatEltSizeInBytes(memRefType);
+  if (!eltSize)
+    return std::nullopt;
+  return *eltSize * *numElements;
 }
 
 /// Returns the size of memref data in bytes if it's statically shaped,
 /// std::nullopt otherwise.  If the element of the memref has vector type, takes
 /// into account size of the vector as well.
 //  TODO: improve/complete this when we have target data.
-std::optional<uint64_t> mlir::getMemRefSizeInBytes(MemRefType memRefType) {
+std::optional<uint64_t>
+mlir::getIntOrFloatMemRefSizeInBytes(MemRefType memRefType) {
   if (!memRefType.hasStaticShape())
     return std::nullopt;
   auto elementType = memRefType.getElementType();
   if (!elementType.isIntOrFloat() && !elementType.isa<VectorType>())
     return std::nullopt;
 
-  uint64_t sizeInBytes = getMemRefEltSizeInBytes(memRefType);
+  auto sizeInBytes = getMemRefIntOrFloatEltSizeInBytes(memRefType);
+  if (!sizeInBytes)
+    return std::nullopt;
   for (unsigned i = 0, e = memRefType.getRank(); i < e; i++) {
-    sizeInBytes = sizeInBytes * memRefType.getDimSize(i);
+    sizeInBytes = *sizeInBytes * memRefType.getDimSize(i);
   }
   return sizeInBytes;
 }
 
 template <typename LoadOrStoreOp>
 LogicalResult mlir::boundCheckLoadOrStoreOp(LoadOrStoreOp loadOrStoreOp,
                                             bool emitError) {
   static_assert(llvm::is_one_of<LoadOrStoreOp, AffineReadOpInterface,

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

 // This can increase the loop depth at which we can fuse a slice, since we are
 // pushing loop carried dependence to a greater depth in the loop nest.
 static void sinkSequentialLoops(MemRefDependenceGraph::Node *node) {
   assert(isa<AffineForOp>(node->op));
   AffineForOp newRootForOp = sinkSequentialLoops(cast<AffineForOp>(node->op));
   node->op = newRootForOp;
 }
 
-//  TODO: improve/complete this when we have target data.
-static unsigned getMemRefEltSizeInBytes(MemRefType memRefType) {
-  auto elementType = memRefType.getElementType();
-
-  unsigned sizeInBits;
-  if (elementType.isIntOrFloat()) {
-    sizeInBits = elementType.getIntOrFloatBitWidth();
-  } else {
-    auto vectorType = elementType.cast<VectorType>();
-    sizeInBits =
-        vectorType.getElementTypeBitWidth() * vectorType.getNumElements();
-  }
-  return llvm::divideCeil(sizeInBits, 8);
-}
-
 // Creates and returns a private (single-user) memref for fused loop rooted
 // at 'forOp', with (potentially reduced) memref size based on the
 // MemRefRegion written to by 'srcStoreOpInst' at depth 'dstLoopDepth'.
 // TODO: consider refactoring the common code from generateDma and
 // this one.
 static Value createPrivateMemRef(AffineForOp forOp, Operation *srcStoreOpInst,
                                  unsigned dstLoopDepth,
                                  std::optional<unsigned> fastMemorySpace,
for (unsigned d = 0; d < rank; ++d) {
     assert(lbDivisors[d] > 0);
     offset =
         (offset + lbs[d][cst->getNumCols() - 1 - rank]).floorDiv(lbDivisors[d]);
     offsets.push_back(offset);
   }
 
   // Create 'newMemRefType' using 'newShape' from MemRefRegion accessed
   // by 'srcStoreOpInst'.
-  uint64_t bufSize = getMemRefEltSizeInBytes(oldMemRefType) * *numElements;
+  auto eltSize = getMemRefIntOrFloatEltSizeInBytes(oldMemRefType);
+  assert(eltSize && "memrefs with size elt types expected");
+  uint64_t bufSize = *eltSize * *numElements;
   unsigned newMemSpace;
   if (bufSize <= localBufSizeThreshold && fastMemorySpace.has_value()) {
     newMemSpace = *fastMemorySpace;
   } else {
     newMemSpace = oldMemRefType.getMemorySpaceAsInt();
   }
   auto newMemRefType = MemRefType::get(newShape, oldMemRefType.getElementType(),
                                        {}, newMemSpace);

mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp

if (!existingBuf) {
 
     // Create the fast memory space buffer just before the 'affine.for'
     // operation.
     fastMemRef =
         prologue.create<memref::AllocOp>(loc, fastMemRefType).getResult();
     // Record it.
     fastBufferMap[memref] = fastMemRef;
     // fastMemRefType is a constant shaped memref.
-    *sizeInBytes = *getMemRefSizeInBytes(fastMemRefType);
+    auto maySizeInBytes = getIntOrFloatMemRefSizeInBytes(fastMemRefType);
+    // We don't account for things of unknown size.
+    if (!maySizeInBytes)
+      maySizeInBytes = 0;
+
     LLVM_DEBUG(emitRemarkForBlock(*block)
                << "Creating fast buffer of type " << fastMemRefType
                << " and size " << llvm::divideCeil(*sizeInBytes, 1024)
                << " KiB\n");
   } else {
     // Reuse the one already created.
     fastMemRef = fastBufferMap[memref];
     *sizeInBytes = 0;

mlir/test/Dialect/Affine/affine-data-copy.mlir

func.func @affine_parallel(%85:memref<2x5x4x2xi64>) {
   // CHECK-NEXT:    affine.for %{{.*}} = 0 to 4
   // CHECK-NEXT:      affine.for %{{.*}} = 0 to 2
 
   // CHECK:     affine.for
   // CHECK-NEXT:  affine.parallel
   // CHECK-NEXT:    affine.parallel
   return
 }
+
+// CHECK-LABEL: func @index_elt_type
+func.func @index_elt_type(%arg0: memref<1x2x4x8xindex>) {
+  affine.for %arg1 = 0 to 1 {
+    affine.for %arg2 = 0 to 2 {
+      affine.for %arg3 = 0 to 4 {
+        affine.for %arg4 = 0 to 8 {
+          affine.store %arg4, %arg0[%arg1, %arg2, %arg3, %arg4] : memref<1x2x4x8xindex>
+        }
+      }
+    }
+  }
+
+  // CHECK:     affine.for %{{.*}} = 0 to 1
+  // CHECK-NEXT:  affine.for %{{.*}} = 0 to 2
+  // CHECK-NEXT:    affine.for %{{.*}} = 0 to 4
+  // CHECK-NEXT:      affine.for %{{.*}} = 0 to 8
+
+  // CHECK:     affine.for %{{.*}} = 0 to 2
+  // CHECK-NEXT:  affine.for %{{.*}} = 0 to 4
+  // CHECK-NEXT:    affine.for %{{.*}} = 0 to 8
+  return
+}