[LinalgToXeGPU] Support squeezable any-D memrefs

dchigarev · dchigarev · commit 9dd135edb6bb · 2024-11-29T13:55:20.000Z
Signed-off-by: dchigarev &lt;dmitry.chigarev@intel.com&gt;
diff --git a/include/gc/Transforms/Utils/ValueUtils.h b/include/gc/Transforms/Utils/ValueUtils.h
@@ -53,6 +53,27 @@ Value flattenMemref(PatternRewriter &rewriter, Location loc, Value srcMemref);
 // Return true if the memref has shared memory space.
 bool hasSharedMemSpace(mlir::Value memref);
 
+// Go through all parent 'memref.subview' ops for the given `memref`
+// and return the folded offsets of all subviews and the root memref.
+std::tuple<SmallVector<Value>, Value>
+computeSubviewOffsets(PatternRewriter &rewriter, Location loc, Value memref);
+
+// Return the strides of the memref
+SmallVector<OpFoldResult> getMemrefStrides(PatternRewriter &rewriter,
+                                           Location loc, Value memref);
+
+// Squeeze the leading dimensions of a given memref up to 'maxDims'.
+FailureOr<Value> squeezeMemref(PatternRewriter &rewriter, Location loc,
+                               Value memref, size_t maxDims = 2);
+
+// Squeeze the leading dimensions of memref operands of a given 'linalgOp'.
+LogicalResult maybeSqueezeDims(PatternRewriter &rewriter,
+                               linalg::LinalgOp linalgOp, size_t maxDims = 2);
+
+// Return if a memref with the given shape can be squeezed to the shape of
+// 'maxDims'. Only leading dimensions are considered squeezable.
+bool canSqueezeDims(llvm::ArrayRef<int64_t> shape, size_t maxDims = 2);
+
 } // namespace utils
 } // namespace mlir
 
diff --git a/lib/gc/Transforms/GPU/LinalgToXeGPU.cpp b/lib/gc/Transforms/GPU/LinalgToXeGPU.cpp
@@ -62,28 +62,6 @@ static Value createFullMask(PatternRewriter &rewriter, Location loc,
   return res.getResult();
 }
 
-// Extracts the offsets from a subview operation as values.
-// The differense from mlir::getMixedOffsets is that this function
-// returns the offsets as mlir::Value that can already be used as an argument
-// for other mlir::Operations.
-static SmallVector<Value> extractOffsetsAsValues(PatternRewriter &rewriter,
-                                                 Location loc,
-                                                 memref::SubViewOp subview) {
-  SmallVector<Value> offsetValues;
-  auto staticOffsets = subview.getStaticOffsets();
-  auto dynamicOffsets = subview.getOffsets();
-  size_t dynIdx = 0;
-  for (size_t i = 0; i < staticOffsets.size(); i++) {
-    if (staticOffsets[i] == ShapedType::kDynamic)
-      offsetValues.push_back(dynamicOffsets[dynIdx++]);
-    else
-      offsetValues.push_back(
-          rewriter.create<arith::ConstantIndexOp>(loc, staticOffsets[i]));
-  }
-
-  return offsetValues;
-}
-
 // Max number of elements to load/store from SLM
 constexpr int64_t maxSLMTileSize = 32;
 
@@ -214,7 +192,8 @@ static LogicalResult isValidMemrefOperand(linalg::LinalgOp linalgOp,
         linalgOp, "Expect memref operand for XeGPU lowering");
   }
 
-  if (type.getShape().size() > maxDims) {
+  if (type.getShape().size() > maxDims &&
+      !utils::canSqueezeDims(type.getShape(), maxDims)) {
     return rewriter.notifyMatchFailure(
         linalgOp, "Too high dimensionality for XeGPU operations");
   }
@@ -856,43 +835,31 @@ static SmallVector<Value> createSLMDescTiles(PatternRewriter &rewriter,
   auto srcType = cast<MemRefType>(src.getType());
   assert(srcType.getRank() == 2 && "Expected a 2D memref");
 
-  SmallVector<int64_t> memrefStrides;
-  Value blockOffset;
-
   // 'imex::ConvertGPUXToSPIRVPass' doesn't allow 'memref.subview' ops in the
   // GPU kernel. We have to merge the subview offsets into the descriptor
   // offset.
-  if (auto subView = dyn_cast<memref::SubViewOp>(src.getDefiningOp())) {
-    auto offsets = extractOffsetsAsValues(rewriter, loc, subView);
-    assert(offsets.size() == 2 && "Expected 2D subview offsets");
-
-    auto xIntOffs = offsets[0];
-    auto yIntOffs = offsets[1];
-
-    // compute 'blockOffset' (beginning of the subview block in the original
-    // flat memref)
-    auto rowStride =
-        cast<MemRefType>(subView.getOperand(0).getType()).getShape()[1];
-    auto rowStrideValue =
-        rewriter.create<arith::ConstantIndexOp>(loc, rowStride);
-
-    auto rowBlockOffset =
-        rewriter.create<arith::MulIOp>(loc, xIntOffs, rowStrideValue)
-            .getResult();
-    blockOffset = rewriter.create<arith::AddIOp>(loc, rowBlockOffset, yIntOffs)
-                      .getResult();
+  auto [offsets, rootMemref] = utils::computeSubviewOffsets(rewriter, loc, src);
+  auto rootStridesFold = utils::getMemrefStrides(rewriter, loc, rootMemref);
+  auto rootStrides =
+      getValueOrCreateConstantIndexOp(rewriter, loc, rootStridesFold);
 
-    memrefStrides = {rowStride, 1};
-    src = subView.getOperand(0);
-  } else {
-    // If the source is not a subview, then the blockOffset is 0
-    blockOffset = rewriter.create<arith::ConstantIndexOp>(loc, 0);
-    memrefStrides = {srcType.getShape()[1], 1};
+  assert(rootStrides.size() == offsets.size() &&
+         "Expected same number of strides and offsets");
+
+  // blockOffset = sum(rootStrides[i] * offsets[i])
+  Value blockOffset = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+  for (size_t i = 0; i < rootStrides.size(); i++) {
+    auto mul = rewriter.create<arith::MulIOp>(loc, rootStrides[i], offsets[i]);
+    blockOffset = rewriter.create<arith::AddIOp>(loc, blockOffset, mul);
   }
 
-  // Scatter descriptors only work with 1D memrefs
-  src = utils::flattenMemref(rewriter, loc, src);
+  auto memrefStridesFold = utils::getMemrefStrides(rewriter, loc, src);
+  auto [memrefStrides, memrefStridesDynamic] =
+      decomposeMixedValues(memrefStridesFold);
+  assert(memrefStridesDynamic.size() == 0 &&
+         "Expected all values to be resolved");
 
+  src = utils::flattenMemref(rewriter, loc, rootMemref);
   return createScatterDescriptorTiles(
       rewriter, loc, /*flatMemref=*/src, /*loadShape2D=*/loadShape,
       /*tileSize2D=*/descTile, /*memrefStrides=*/memrefStrides,
@@ -1839,6 +1806,11 @@ struct ConvertGemmLikeToXeGPU : public OpRewritePattern<LinalgOpTy> {
     if (failed(isOutputValid))
       return isOutputValid;
 
+    if (failed(mlir::utils::maybeSqueezeDims(rewriter, gemmLikeOp))) {
+      return rewriter.notifyMatchFailure(
+          gemmLikeOp, "Failed to squeeze dimensions of GEMM-like operation");
+    }
+
     // Ensure that reduction dimension tiling also works for smaller
     // workloads.
     auto aType = cast<ShapedType>(gemmLikeOp.getDpsInputs()[0].getType());
@@ -1894,6 +1866,12 @@ struct ConvertNamedEltwiseToXeGPU : public OpRewritePattern<LinalgOpTy> {
     if (failed(isOutputValid))
       return isOutputValid;
 
+    if (failed(utils::maybeSqueezeDims(rewriter, eltwiseOp))) {
+      return rewriter.notifyMatchFailure(
+          eltwiseOp,
+          "Could not squeeze dimensions of the elementwise operation");
+    }
+
     return createEltwiseKernel(eltwiseOp, rewriter);
   }
 
@@ -1988,6 +1966,12 @@ struct ConvertMemoryFillToXeGPU : public OpRewritePattern<LinalgOpTy> {
     if (failed(isOutputValid))
       return isOutputValid;
 
+    if (failed(utils::maybeSqueezeDims(rewriter, linalgOp))) {
+      return rewriter.notifyMatchFailure(
+          linalgOp,
+          "Could not squeeze dimensions of the memory fill operation");
+    }
+
     return createMemoryFillKernel(linalgOp, rewriter);
   }
 
diff --git a/lib/gc/Transforms/Utils/ValueUtils.cpp b/lib/gc/Transforms/Utils/ValueUtils.cpp
@@ -6,10 +6,14 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include <numeric>
+
+#include "mlir/Dialect/Affine/ViewLikeInterfaceUtils.h"
 #include "mlir/Dialect/GPU/IR/GPUDialect.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Matchers.h"
@@ -155,9 +159,10 @@ Value flattenMemref(PatternRewriter &rewriter, Location loc, Value srcMemref) {
   auto srcType = cast<MemRefType>(srcMemref.getType());
 
   assert(srcType && "Expected a memref type");
-  assert(srcType.getRank() == 2 && "Expected a 2D memref");
 
-  int64_t flatSize = srcType.getShape()[0] * srcType.getShape()[1];
+  auto shapeNd = srcType.getShape();
+  int64_t flatSize =
+      std::accumulate(shapeNd.begin(), shapeNd.end(), 1, std::multiplies<>());
 
   Value offset = rewriter.create<arith::ConstantIndexOp>(loc, 0);
   Value size = rewriter.create<arith::ConstantIndexOp>(loc, flatSize);
@@ -193,5 +198,128 @@ bool hasSharedMemSpace(mlir::Value memref) {
   return false;
 }
 
+std::tuple<SmallVector<Value>, Value>
+computeSubviewOffsets(PatternRewriter &rewriter, Location loc, Value memref) {
+  auto fillVal = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+  auto origShape = dyn_cast<MemRefType>(memref.getType()).getShape();
+
+  SmallVector<Value> resolvedOffsets(origShape.size(), fillVal);
+
+  while (auto subViewOp = memref.getDefiningOp<memref::SubViewOp>()) {
+    auto currentOffsets = getAsOpFoldResult(resolvedOffsets);
+    resolvedOffsets.clear();
+
+    affine::resolveIndicesIntoOpWithOffsetsAndStrides(
+        rewriter, memref.getLoc(), subViewOp.getMixedOffsets(),
+        subViewOp.getMixedStrides(), subViewOp.getDroppedDims(), currentOffsets,
+        resolvedOffsets);
+    memref = subViewOp.getOperand(0);
+  }
+
+  return std::make_tuple(resolvedOffsets, memref);
+}
+
+SmallVector<OpFoldResult> getMemrefStrides(PatternRewriter &rewriter,
+                                           Location loc, Value memref) {
+  auto type = dyn_cast<MemRefType>(memref.getType());
+
+  auto stridedLayout = dyn_cast<StridedLayoutAttr>(type.getLayout());
+  if (stridedLayout) {
+    auto strides = stridedLayout.getStrides();
+    return getMixedValues(strides, {}, rewriter);
+  }
+
+  auto sizes = getMixedValues(type.getShape(), {}, rewriter);
+  auto strides = memref::computeStridesIRBlock(loc, rewriter, sizes);
+  return strides;
+}
+
+FailureOr<Value> squeezeMemref(PatternRewriter &rewriter, Location loc,
+                               Value memref, size_t maxDims = 2) {
+  auto type = dyn_cast<MemRefType>(memref.getType());
+  auto shape = type.getShape();
+
+  if (shape.size() <= maxDims)
+    return memref;
+
+  for (size_t i = 0; i < shape.size() - maxDims; i++)
+    if (shape[i] != 1)
+      return failure();
+
+  auto offsets =
+      getMixedValues(SmallVector<int64_t>(shape.size(), 0), {}, rewriter);
+  auto sizes = getMixedValues(shape, {}, rewriter);
+  auto staticStrides = utils::getStaticStrides(memref).value();
+  auto strides =
+      getMixedValues(SmallVector<int64_t>(shape.size(), 1), {}, rewriter);
+
+  SmallVector<int64_t> newShape(shape.begin() + shape.size() - maxDims,
+                                shape.end());
+  SmallVector<int64_t> newStrides(
+      staticStrides.begin() + shape.size() - maxDims, staticStrides.end());
+
+  int64_t newOffset = 0;
+  if (auto memrefLayout = dyn_cast<StridedLayoutAttr>(type.getLayout()))
+    newOffset = memrefLayout.getOffset();
+
+  auto newLayout = StridedLayoutAttr::get(
+      rewriter.getContext(), /*offset=*/newOffset, /*strides=*/newStrides);
+  MemRefType newMemRefType = MemRefType::get(newShape, type.getElementType(),
+                                             newLayout, type.getMemorySpace());
+
+  auto squeezedSubview =
+      rewriter
+          .create<memref::SubViewOp>(loc, newMemRefType, memref, offsets, sizes,
+                                     strides)
+          .getResult();
+  return squeezedSubview;
+}
+
+LogicalResult maybeSqueezeDims(PatternRewriter &rewriter,
+                               linalg::LinalgOp linalgOp, size_t maxDims) {
+  SmallVector<std::pair<size_t, Value>> newOperands;
+  auto operands = linalgOp->getOperands();
+  auto loc = linalgOp.getLoc();
+
+  for (size_t i = 0; i < operands.size(); i++) {
+    auto operand = operands[i];
+    auto type = dyn_cast<MemRefType>(operand.getType());
+    if (!type) {
+      // maybe should 'continue' here instead and skip non-memref operands?
+      // TODO: replace this with 'continue' if such case would appear someday
+      return rewriter.notifyMatchFailure(
+          linalgOp, "Expect memref operand for XeGPU lowering");
+    }
+
+    if (type.getShape().size() <= maxDims)
+      continue;
+
+    auto res = squeezeMemref(rewriter, loc, operand, maxDims);
+    if (failed(res)) {
+      return rewriter.notifyMatchFailure(
+          linalgOp, "Can't squeeze memref to the desired number of dimensions");
+    }
+
+    auto flatSubview = res.value();
+    newOperands.emplace_back(i, flatSubview);
+  }
+
+  for (auto [i, operand] : newOperands)
+    linalgOp->setOperand(i, operand);
+
+  return success();
+}
+
+bool canSqueezeDims(llvm::ArrayRef<int64_t> shape, size_t maxDims) {
+  if (shape.size() <= maxDims)
+    return true;
+
+  for (size_t i = 0; i < shape.size() - maxDims; i++)
+    if (shape[i] != 1)
+      return false;
+
+  return true;
+}
+
 } // namespace utils
 } // namespace mlir
