Address comments

Author: adam-smnk

mlir/include/mlir/Conversion/LLVMCommon/Pattern.h

@@ -69,6 +69,15 @@ SmallVector<Value> decomposeValue(OpBuilder &builder, Location loc, Value src,
 /// function is used to combine multiple values into a single value.
 Value composeValue(OpBuilder &builder, Location loc, ValueRange src,
                    Type dstType);
+
+/// Performs the index computation to get to the element at `indices` of the
+/// memory pointed to by `memRefDesc`, using the layout map of `type`.
+/// The indices are linearized as:
+///   `base_offset + index_0 * stride_0 + ... + index_n * stride_n`.
+Value getStridedElementPtr(OpBuilder &builder, Location loc,
+                           const LLVMTypeConverter &converter, MemRefType type,
+                           Value memRefDesc, ValueRange indices,
+                           LLVM::GEPNoWrapFlags noWrapFlags = LLVM::GEPNoWrapFlags::none);
 } // namespace LLVM
 
 /// Base class for operation conversions targeting the LLVM IR dialect. It
@@ -107,8 +116,8 @@ class ConvertToLLVMPattern : public ConversionPattern {
   static Value createIndexAttrConstant(OpBuilder &builder, Location loc,
                                        Type resultType, int64_t value);
 
-  // This is a strided getElementPtr variant that linearizes subscripts as:
-  //   `base_offset + index_0 * stride_0 + ... + index_n * stride_n`.
+  /// Convenience wrapper for the corresponding helper utility.
+  /// This is a strided getElementPtr variant with linearized subscripts.
   Value getStridedElementPtr(
       ConversionPatternRewriter &rewriter, Location loc, MemRefType type,
       Value memRefDesc, ValueRange indices,

mlir/include/mlir/Dialect/AMX/AMX.td

@@ -25,8 +25,8 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef AMX_OPS
-#define AMX_OPS
+#ifndef AMX
+#define AMX
 
 include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 include "mlir/Dialect/AMX/AMXInterfaces.td"
@@ -371,4 +371,4 @@ def TileMulIOp : AMX_Op<"tile_muli", [Pure,
   let hasVerifier = 1;
 }
 
-#endif // AMX_OPS
+#endif // AMX

mlir/lib/Conversion/LLVMCommon/Pattern.cpp

@@ -62,38 +62,8 @@ Value ConvertToLLVMPattern::getStridedElementPtr(
     ConversionPatternRewriter &rewriter, Location loc, MemRefType type,
     Value memRefDesc, ValueRange indices,
     LLVM::GEPNoWrapFlags noWrapFlags) const {
-
-  auto [strides, offset] = type.getStridesAndOffset();
-
-  MemRefDescriptor memRefDescriptor(memRefDesc);
-  // Use a canonical representation of the start address so that later
-  // optimizations have a longer sequence of instructions to CSE.
-  // If we don't do that we would sprinkle the memref.offset in various
-  // position of the different address computations.
-  Value base =
-      memRefDescriptor.bufferPtr(rewriter, loc, *getTypeConverter(), type);
-
-  Type indexType = getIndexType();
-  Value index;
-  for (int i = 0, e = indices.size(); i < e; ++i) {
-    Value increment = indices[i];
-    if (strides[i] != 1) { // Skip if stride is 1.
-      Value stride =
-          ShapedType::isDynamic(strides[i])
-              ? memRefDescriptor.stride(rewriter, loc, i)
-              : createIndexAttrConstant(rewriter, loc, indexType, strides[i]);
-      increment = rewriter.create<LLVM::MulOp>(loc, increment, stride);
-    }
-    index =
-        index ? rewriter.create<LLVM::AddOp>(loc, index, increment) : increment;
-  }
-
-  Type elementPtrType = memRefDescriptor.getElementPtrType();
-  return index ? rewriter.create<LLVM::GEPOp>(
-                     loc, elementPtrType,
-                     getTypeConverter()->convertType(type.getElementType()),
-                     base, index, noWrapFlags)
-               : base;
+  return LLVM::getStridedElementPtr(rewriter, loc, *getTypeConverter(), type,
+                                    memRefDesc, indices, noWrapFlags);
 }
 
 // Check if the MemRefType `type` is supported by the lowering. We currently
@@ -513,3 +483,41 @@ Value mlir::LLVM::composeValue(OpBuilder &builder, Location loc, ValueRange src,
 
   return res;
 }
+
+Value mlir::LLVM::getStridedElementPtr(OpBuilder &builder, Location loc,
+                                       const LLVMTypeConverter &converter,
+                                       MemRefType type, Value memRefDesc,
+                                       ValueRange indices,
+                                       LLVM::GEPNoWrapFlags noWrapFlags) {
+  auto [strides, offset] = type.getStridesAndOffset();
+
+  MemRefDescriptor memRefDescriptor(memRefDesc);
+  // Use a canonical representation of the start address so that later
+  // optimizations have a longer sequence of instructions to CSE.
+  // If we don't do that we would sprinkle the memref.offset in various
+  // position of the different address computations.
+  Value base = memRefDescriptor.bufferPtr(builder, loc, converter, type);
+
+  Type indexType = converter.getIndexType();
+  Value index;
+  for (int i = 0, e = indices.size(); i < e; ++i) {
+    Value increment = indices[i];
+    if (strides[i] != 1) { // Skip if stride is 1.
+      Value stride =
+          ShapedType::isDynamic(strides[i])
+              ? memRefDescriptor.stride(builder, loc, i)
+              : builder.create<LLVM::ConstantOp>(
+                    loc, indexType, builder.getIndexAttr(strides[i]));
+      increment = builder.create<LLVM::MulOp>(loc, increment, stride);
+    }
+    index =
+        index ? builder.create<LLVM::AddOp>(loc, index, increment) : increment;
+  }
+
+  Type elementPtrType = memRefDescriptor.getElementPtrType();
+  return index ? builder.create<LLVM::GEPOp>(
+                     loc, elementPtrType,
+                     converter.convertType(type.getElementType()), base, index,
+                     noWrapFlags)
+               : base;
+}

mlir/lib/Dialect/AMX/IR/AMXDialect.cpp

@@ -64,46 +64,6 @@ static LogicalResult verifyMultShape(Operation *op, amx::TileType atp,
   return success();
 }
 
-/// Get pointer to a memref descriptor.
-/// Optionally, the base pointer can be offset using linearized index computed
-/// from the given indices.
-static Value getBufferPtr(Location loc, MemRefType type, Value buffer,
-                          ValueRange indices,
-                          const LLVMTypeConverter &typeConverter,
-                          RewriterBase &rewriter) {
-  auto [strides, offset] = type.getStridesAndOffset();
-
-  MemRefDescriptor memRefDescriptor(buffer);
-  Value base = memRefDescriptor.bufferPtr(rewriter, loc, typeConverter, type);
-
-  int numIndices = indices.size();
-  if (numIndices == 0)
-    return base;
-
-  assert(type.getRank() == numIndices &&
-         "expects number of indices equal to memref rank");
-  Value index;
-  Type indexType = typeConverter.getIndexType();
-  for (int i = 0; i < numIndices; ++i) {
-    Value increment = indices[i];
-    if (strides[i] != 1) { // Skip if stride is 1.
-      Value stride =
-          ShapedType::isDynamic(strides[i])
-              ? memRefDescriptor.stride(rewriter, loc, i)
-              : rewriter.create<LLVM::ConstantOp>(
-                    loc, indexType, rewriter.getIndexAttr(strides[i]));
-      increment = rewriter.create<LLVM::MulOp>(loc, increment, stride);
-    }
-    index =
-        index ? rewriter.create<LLVM::AddOp>(loc, index, increment) : increment;
-  }
-
-  Type elementPtrType = memRefDescriptor.getElementPtrType();
-  return rewriter.create<LLVM::GEPOp>(
-      loc, elementPtrType, typeConverter.convertType(type.getElementType()),
-      base, index);
-}
-
 /// Maps the 2-dim vector shape to the two 16-bit tile sizes. The first
 /// dimension directly translates into the number of rows of the tiles.
 /// The second dimensions needs to be scaled by the number of bytes.
@@ -122,7 +82,6 @@ static SmallVector<Value> getTileSizes(Location loc, amx::TileType tType,
 
 /// Maps the 2-dim memref shape to the 64-bit stride. Note that the buffer
 /// shape may "envelop" the actual tile shape, and may be dynamically sized.
-/// Returns failure if proper stride couldn't be found.
 static Value getStride(Location loc, MemRefType mType, Value base,
                        RewriterBase &rewriter) {
   assert(mType.getRank() >= 2 && "Invalid shape for AMX strides");
@@ -184,8 +143,8 @@ amx::TileLoadOp::getIntrinsicOperands(ArrayRef<Value> operands,
   SmallVector<Value> intrinsicOperands;
   intrinsicOperands.append(getTileSizes(loc, getTileType(), rewriter));
   intrinsicOperands.push_back(
-      getBufferPtr(loc, getMemRefType(), adaptor.getBase(),
-                   adaptor.getIndices(), typeConverter, rewriter));
+      LLVM::getStridedElementPtr(rewriter, loc, typeConverter, getMemRefType(),
+                                 adaptor.getBase(), adaptor.getIndices()));
   intrinsicOperands.push_back(
       getStride(loc, getMemRefType(), adaptor.getBase(), rewriter));
 
@@ -217,8 +176,8 @@ amx::TileStoreOp::getIntrinsicOperands(ArrayRef<Value> operands,
   SmallVector<Value> intrinsicOperands;
   intrinsicOperands.append(getTileSizes(loc, getTileType(), rewriter));
   intrinsicOperands.push_back(
-      getBufferPtr(loc, getMemRefType(), adaptor.getBase(),
-                   adaptor.getIndices(), typeConverter, rewriter));
+      LLVM::getStridedElementPtr(rewriter, loc, typeConverter, getMemRefType(),
+                                 adaptor.getBase(), adaptor.getIndices()));
   intrinsicOperands.push_back(
       getStride(loc, getMemRefType(), adaptor.getBase(), rewriter));
   intrinsicOperands.push_back(adaptor.getVal());