Backend changes

Author: justinrosner

mlir/include/mlir/Dialect/Rock/IR/RockOps.td

@@ -1141,11 +1141,13 @@ def Rock_GlobalPrefetchOp
 def Rock_GlobalLoadOp
     : Rock_Op<
           "global_load", [DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
-                          AllElementTypesMatch<["source", "result"]>]>,
+                          AllElementTypesMatch<["source", "result"]>,
+                          AttrSizedOperandSegments]>,
       Arguments<(
           ins Arg<MemRefOf<SupportedMemoryElems>, "source memory">:$source,
           I1:$valid, Variadic<Index>:$sourceCoord, UnitAttr:$needs64BitIdx,
-          UnitAttr:$canReadOffEnd)>,
+          UnitAttr:$canReadOffEnd, Optional<I64>:$pagePtr,
+          OptionalAttr<I64Attr>:$pageSize)>,
       Results<(outs AnyType:$result)> {
   let summary = "Load from global memory, applying bounds checks";
   let description = [{
@@ -1162,9 +1164,17 @@ def Rock_GlobalLoadOp
     If `canReadOffEnd` is present, then the `valid` input is ignored and the
     bounds-checked implementation of this operation is always used. This
     is used to simplify the implementation of certain utility kernels.
+
+    When `pagePtr` and `pageSize` are provided, this
+    operation loads from paged memory. The `source` memref is still used for
+    type information (element type, etc.), but the actual load uses `pagePtr`
+    as the base address. `sourceCoord[0]` is treated as the offset-in-page in
+    elements (a single flat index, not multi-dimensional coordinates).
   }];
   let assemblyFormat = [{
-    $source `[` $sourceCoord `]` `if` $valid attr-dict
+    $source `[` $sourceCoord `]` `if` $valid
+    (`paged` $pagePtr^)?
+    attr-dict
     `:` type($source) `->` type($result)
   }];
   let hasVerifier = 1;
@@ -1200,7 +1210,8 @@ def Rock_GlobalLoadToLDSOp
   }];
   let assemblyFormat = [{
     $source `[` $sourceCoord `]` 
-    `->` $dest `[` $destCoord `]`  `if` $valid attr-dict
+    `->` $dest `[` $destCoord `]`  `if` $valid
+    attr-dict
     `:` type($source) `->` type($dest)
   }];
   let hasVerifier = 1;
@@ -1302,14 +1313,21 @@ def Rock_ThreadwiseReadIntoOp
         [](::mlir::Type lhs, ::mlir::TypeRange rhs) -> bool {
           return llvm::all_of(rhs, [lhs](::mlir::Type t) { return t == lhs; });
         }
-      }]>]>,
-      Arguments<(ins Arg<MemRefOf<NativeMemoryOpTypes>, "source view">:$source,
+      }]>]> {
+
+  let arguments =
+      (ins Arg<MemRefOf<NativeMemoryOpTypes>, "source view">:$source,
           Arg<MemRefOf<NativeMemoryOpTypes>, "destination registers">:$dest,
           Variadic<VectorOfNonZeroRankOf<[I1]>>:$dynamicValidities,
           TransformMapArrayAttr:$extraViews, Variadic<Index>:$extraIndices,
           UnitAttr:$forceUnroll, UnitAttr:$useIndexDiffs,
-          OptionalAttr<Rock_LDSTransposeConfigAttr>:$ldsTransposeConfig)>,
-      Results<(outs Optional<VectorOfNonZeroRankOf<[I1]>>:$validityRecord)> {
+          OptionalAttr<Rock_LDSTransposeConfigAttr>:$ldsTransposeConfig,
+          Optional<MemRefOf<[I64]>>:$ldsPagePtrs,
+          Optional<Index>:$firstPageIndex, OptionalAttr<IndexAttr>:$pageSize,
+          OptionalAttr<IndexAttr>:$numPagesPerBatch);
+
+  let results = (outs Optional<VectorOfNonZeroRankOf<[I1]>>:$validityRecord);
+
   let summary = "Read values from transformed source into destination";
 
   let description = [{
@@ -1341,8 +1359,8 @@ def Rock_ThreadwiseReadIntoOp
     L is the length of `%dest`, V is the maximum vectorization computed
     for MAPS.
 
-    The input to extraViews: (the transforms on %source) must have the form
-    (extraIdx0, ... , extraIdxN, iteration_number)
+    The input to extraViews (the transforms on %source) must have the form
+    (extraIdx0, ... , extraIdxN, iteration_number).
 
     Primarily, extraIndices would be used to pass in tid and bid. This would need
     to have a matching view in [extraViews]source. The extraIndices could be used
@@ -1368,6 +1386,14 @@ def Rock_ThreadwiseReadIntoOp
 
     This operation is also used during fusion to represent loads from additional
     arguments like bias tensors.
+
+    When `ldsPagePtrs`, `firstPageIndex`, and `pageSize` are provided, this
+    operation reads from paged memory where data is scattered across
+    non-contiguous pages accessed via a page table. The page pointers are
+    pre-loaded into LDS by the calling code to avoid redundant global memory
+    accesses across threads. During lowering, the transform maps compute the
+    flat position for each element, determining the page index and offset
+    within page.
   }];
 
   let builders = [
@@ -1377,28 +1403,38 @@ def Rock_ThreadwiseReadIntoOp
                     "ValueRange":$extraIndices, "bool":$forceUnroll,
                     "bool":$useIndexDiffs),
                 [{
-        build($_builder, $_state, TypeRange{}, source, dest, ValueRange{}, extraViews, extraIndices, forceUnroll, useIndexDiffs, /*ldsTransposeConfig=*/nullptr);
+        build($_builder, $_state, TypeRange{}, source, dest, ValueRange{}, extraViews, extraIndices, forceUnroll, useIndexDiffs, /*ldsTransposeConfig=*/nullptr, /*ldsPagePtrs=*/Value{}, /*firstPageIndex=*/Value{}, /*pageSize=*/nullptr, /*numPagesPerBatch=*/nullptr);
       }]>,
       // Builder with explicit ldsTransposeConfig support
       OpBuilder<(ins "Value":$source, "Value":$dest, "ArrayAttr":$extraViews,
                     "ValueRange":$extraIndices, "bool":$forceUnroll,
                     "bool":$useIndexDiffs,
                     "LDSTransposeConfigAttr":$ldsTransposeConfig),
                 [{
-        build($_builder, $_state, TypeRange{}, source, dest, ValueRange{}, extraViews, extraIndices, forceUnroll, useIndexDiffs, ldsTransposeConfig);
+        build($_builder, $_state, TypeRange{}, source, dest, ValueRange{}, extraViews, extraIndices, forceUnroll, useIndexDiffs, ldsTransposeConfig, /*ldsPagePtrs=*/Value{}, /*firstPageIndex=*/Value{}, /*pageSize=*/nullptr, /*numPagesPerBatch=*/nullptr);
       }]>,
       // Builder with validityRecord but without ldsTransposeConfig
       OpBuilder<(ins "TypeRange":$validityRecord, "Value":$source,
                     "Value":$dest, "ValueRange":$dynamicValidities,
                     "ArrayAttr":$extraViews, "ValueRange":$extraIndices,
                     "bool":$forceUnroll, "bool":$useIndexDiffs),
                 [{
-        build($_builder, $_state, validityRecord, source, dest, dynamicValidities, extraViews, extraIndices, forceUnroll, useIndexDiffs, /*ldsTransposeConfig=*/nullptr);
+        build($_builder, $_state, validityRecord, source, dest, dynamicValidities, extraViews, extraIndices, forceUnroll, useIndexDiffs, /*ldsTransposeConfig=*/nullptr, /*ldsPagePtrs=*/Value{}, /*firstPageIndex=*/Value{}, /*pageSize=*/nullptr, /*numPagesPerBatch=*/nullptr);
+      }]>,
+      // Builder with validityRecord and ldsTransposeConfig
+      OpBuilder<(ins "TypeRange":$validityRecord, "Value":$source,
+                    "Value":$dest, "ValueRange":$dynamicValidities,
+                    "ArrayAttr":$extraViews, "ValueRange":$extraIndices,
+                    "bool":$forceUnroll, "bool":$useIndexDiffs,
+                    "LDSTransposeConfigAttr":$ldsTransposeConfig),
+                [{
+        build($_builder, $_state, validityRecord, source, dest, dynamicValidities, extraViews, extraIndices, forceUnroll, useIndexDiffs, ldsTransposeConfig, /*ldsPagePtrs=*/Value{}, /*firstPageIndex=*/Value{}, /*pageSize=*/nullptr, /*numPagesPerBatch=*/nullptr);
       }]>];
 
   let assemblyFormat = [{
     attr-dict $extraViews `(` $source `)` (`[` $extraIndices^ `]`)? `->` $dest
     (`if` ` ` `[` $dynamicValidities^ `]`)?
+    (`paged` $ldsPagePtrs^ `:` type($ldsPagePtrs) `[` $firstPageIndex `]`)?
     `:` type($source) `->` type($dest) (`,` type($validityRecord)^)?
   }];
   let hasVerifier = 1;
@@ -1634,7 +1670,9 @@ def Rock_BlockwiseLoadTileOp
           Rock_BlockwiseMatrixParamsAttr:$matrixParamsB, UnitAttr:$isA,
           Variadic<Index>:$sourceIndices,
           OptionalAttr<Rock_GemmFeaturesAttr>:$features, I32Attr:$blockSize,
-          RockAccelTuningParamAttrInterface:$params)> {
+          RockAccelTuningParamAttrInterface:$params,
+          Optional<MemRefOf<[I64]>>:$pageTable,
+          OptionalAttr<IndexAttr>:$pageSize)> {
   let summary =
       "Blockwise load tile from global memory to LDS and/or registers";
   let description = [{
@@ -1651,10 +1689,20 @@ def Rock_BlockwiseLoadTileOp
     `isA` determines if we are loading an A matrix or B matrix. `G`, `M` and `N` are the GEMM sizes.
     `elementTypeA` and `elementTypeB` are used to construct AccelEmitter. They are data types for the Matrix A & B of the GEMMs. 
     `elementLoadType` is the element type of the global buffer this BlockwiseLoadTileOp is trying to load. 
-    `elementType` is the elementType in the registers.  Note that it may differ from `elementLoadType` because of input fusions.  
+    `elementType` is the elementType in the registers.  Note that it may differ from `elementLoadType` because of input fusions.
+
+    For paged attention (when K/V tensors are stored in non-contiguous pages),
+    the following optional parameters enable indirect memory access:
+    - `pageTable`: The i64 page table memref containing base pointers for each page.
+      Shape is [batch, numPages, 1] where each entry is a pointer to a page.
+    - `pageSize`: Number of elements per page, used to compute page boundaries.
+    When these are provided, the lowering will compute page indices from logical
+    positions and load page pointers to resolve physical addresses.
   }];
   let assemblyFormat = [{
-    $source (`[` $sourceIndices^ `]`)? (`LDS` `->` $destLDS^)? (`->` $destRegisters^)? attr-dict
+    $source (`[` $sourceIndices^ `]`)? (`LDS` `->` $destLDS^)? (`->` $destRegisters^)?
+    (`paged` $pageTable^ `:` type($pageTable))?
+    attr-dict
     `:` type($source) (`LDS` `->` type($destLDS)^)? (`->` type($destRegisters)^)?
   }];
 

mlir/include/mlir/Dialect/Rock/utility/transformMapUtils.h

@@ -304,6 +304,11 @@ SmallVector<StringRef> getStringRefsFor(ArrayRef<SmallString<8>> strings);
 // type of the block argument, otherwise it returns failure.
 FailureOr<Type> getInputFusionElementType(Value transformed);
 
+// Compute the flat position in the virtual paged space by evaluating
+// the source transforms on the given coordinates.
+FailureOr<Value> computeFlatPosition(OpBuilder &b, Location loc, Value source,
+                                     ValueRange indices);
+
 } // end namespace rock
 } // end namespace mlir
 #endif

mlir/lib/Dialect/Rock/IR/RockDialect.cpp

@@ -2009,8 +2009,31 @@ static LogicalResult verifyGlobalLoadAndPrefetch(LoadOrPrefetch op) {
   MemRefType sourceType = op.getSource().getType();
   size_t nDims = sourceType.getRank();
 
-  if (op.getSourceCoord().size() != nDims)
+  // Check if this op has paging attributes
+  bool isPaged = false;
+  if constexpr (std::is_same_v<LoadOrPrefetch, GlobalLoadOp>) {
+    isPaged = op.getPagePtr() != nullptr;
+
+    // Verify paging attributes consistency
+    bool hasPagePtr = op.getPagePtr() != nullptr;
+    bool hasPageSize = op.getPageSize().has_value();
+    if (hasPagePtr != hasPageSize) {
+      return op.emitOpError(
+          "pagePtr and pageSize must both be set or both be unset");
+    }
+
+    if (hasPageSize && op.getPageSize().value() <= 0) {
+      return op.emitOpError("pageSize must be positive");
+    }
+  }
+
+  // For paged loads, we expect exactly 1 coordinate (flat offset-in-page)
+  if (isPaged && op.getSourceCoord().size() != 1) {
+    return op.emitOpError("Expected 1 coordinate for paged load");
+  } else if (op.getSourceCoord().size() != nDims) {
     return op.emitOpError("Expected " + Twine(nDims) + " coordinates");
+  }
+
   Attribute memSpaceAttr = sourceType.getMemorySpace();
   auto gpuMemSpaceAttr = dyn_cast_or_null<gpu::AddressSpaceAttr>(memSpaceAttr);
   if (memSpaceAttr && (!gpuMemSpaceAttr ||
@@ -2328,6 +2351,11 @@ ThreadwiseReadIntoOp::cloneWithExtraIndices(OpBuilder &builder,
 void ThreadwiseReadIntoOp::getEffects(
     SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
   getCommonEffects(*this, effects);
+  // If this is a paged load, we also read from the LDS page pointer buffer
+  if (getLdsPagePtrs()) {
+    effects.emplace_back(MemoryEffects::Read::get(),
+                         &getLdsPagePtrsMutable()[0]);
+  }
 }
 
 LogicalResult ThreadwiseReadIntoOp::verify() {
@@ -2393,6 +2421,37 @@ LogicalResult ThreadwiseReadIntoOp::verify() {
           "in register-to-register reads produced by input fusion");
     }
   }
+
+  // Verify paged attention attributes consistency
+  bool hasLdsPagePtrs = getLdsPagePtrs() != nullptr;
+  bool hasFirstPageIndex = getFirstPageIndex() != nullptr;
+  bool hasPageSize = getPageSize().has_value();
+
+  if (hasLdsPagePtrs != hasFirstPageIndex || hasLdsPagePtrs != hasPageSize) {
+    return emitOpError(
+        "ldsPagePtrs, firstPageIndex, and pageSize must all be "
+        "set together for paged attention, or none should be set");
+  }
+
+  if (hasPageSize && getPageSize().value().getSExtValue() <= 0) {
+    return emitOpError("pageSize must be positive");
+  }
+
+  if (getNumPagesPerBatch().has_value() &&
+      getNumPagesPerBatch().value().getSExtValue() <= 0) {
+    return emitOpError("numPagesPerBatch must be positive");
+  }
+
+  if (hasLdsPagePtrs) {
+    MemRefType ldsPagePtrsType = cast<MemRefType>(getLdsPagePtrs().getType());
+    if (ldsPagePtrsType.getRank() != 1) {
+      return emitOpError("ldsPagePtrs must be a 1D memref");
+    }
+    if (!ldsPagePtrsType.getElementType().isInteger(64)) {
+      return emitOpError("ldsPagePtrs must have i64 element type");
+    }
+  }
+
   return success();
 }
 
@@ -2568,6 +2627,34 @@ LogicalResult BlockwiseLoadTileOp::verify() {
     return emitOpError("destRegisters must be set unless loadType is "
                        "Default/DirectToLDSDefault");
 
+  // Verify paged attention attributes consistency
+  bool hasPageTable = getPageTable() != nullptr;
+  bool hasPageSize = getPageSize().has_value();
+
+  if (hasPageTable != hasPageSize) {
+    return emitOpError(
+        "pageTable and pageSize must both be set or both be unset");
+  }
+
+  if (hasPageSize && getPageSize().value().getSExtValue() <= 0) {
+    return emitOpError("pageSize must be positive");
+  }
+
+  if (hasPageTable) {
+    MemRefType pageTableType = cast<MemRefType>(getPageTable().getType());
+    if (pageTableType.getRank() != 3) {
+      return emitOpError(
+          "pageTable must be a 3D memref with shape [batch, numPages, 1]");
+    }
+    if (pageTableType.getShape()[2] != 1) {
+      return emitOpError(
+          "pageTable last dimension must be 1 (shape [batch, numPages, 1])");
+    }
+    if (!pageTableType.getElementType().isInteger(64)) {
+      return emitOpError("pageTable must have i64 element type");
+    }
+  }
+
   return success();
 }