Cose belle

2026-07-14 16:57:58 +02:00
parent d1a29ace3c
commit 51fdb830e5
38 changed files with 5365 additions and 914 deletions
@@ -5,7 +5,6 @@
 #include "llvm/ADT/SmallVector.h"

 #include <algorithm>
-#include <numeric>
 #include <optional>
 #include <type_traits>

@@ -122,14 +121,6 @@ static RankedTensorType getKeepdimsType(RankedTensorType inputType, Type element
  return RankedTensorType::get(shape, elementType, inputType.getEncoding());
 }

-static RankedTensorType getCompactKeptType(RankedTensorType inputType, Type elementType, ArrayRef<bool> reducedAxes) {
-  SmallVector<int64_t> shape;
-  for (auto [dim, isReduced] : llvm::zip_equal(inputType.getShape(), reducedAxes))
-    if (!isReduced)
-      shape.push_back(dim);
-  return RankedTensorType::get(shape, elementType, inputType.getEncoding());
-}
-
 static RankedTensorType getReducedSliceType(RankedTensorType inputType, ArrayRef<bool> reducedAxes) {
  SmallVector<int64_t> shape;
  shape.reserve(inputType.getRank());
@@ -228,59 +219,80 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
  return (*batchOp).getResult(0);
 }

-static Value buildKeepdimsFromLanePackedBatch(Value batchValue,
-                                              RankedTensorType keepdimsType,
-                                              RankedTensorType compactKeptType,
-                                              ArrayRef<bool> reducedAxes,
-                                              ConversionPatternRewriter& rewriter,
-                                              Location loc) {
-  auto batchType = cast<RankedTensorType>(batchValue.getType());
-  if (batchType == keepdimsType)
-    return batchValue;
+static FailureOr<Value> buildReduceMeanKeepdimsBlueprint(
+    Value batchValue, RankedTensorType keepdimsType,
+    ArrayRef<bool> reducedAxes, ConversionPatternRewriter& rewriter,
+    Location loc) {
+  auto batchType = dyn_cast<RankedTensorType>(batchValue.getType());
+  int64_t rank = keepdimsType.getRank();
+  if (!batchType || !batchType.hasStaticShape()
+      || !keepdimsType.hasStaticShape()
+      || static_cast<int64_t>(reducedAxes.size()) != rank
+      || batchType.getRank() != rank + 1
+      || batchType.getElementType() != keepdimsType.getElementType())
+    return failure();

-  SmallVector<ReassociationIndices> collapseToFlat {{}};
-  for (int64_t axis = 0; axis < batchType.getRank(); ++axis)
-    collapseToFlat.front().push_back(axis);
-
-  SmallVector<ReassociationIndices> expandFlatToCompact(1);
-  for (int64_t axis = 0; axis < compactKeptType.getRank(); ++axis)
-    expandFlatToCompact.front().push_back(axis);
-
-  SmallVector<ReassociationIndices> expandCompactToKeepdims;
-  ReassociationIndices pendingLeadingReducedAxes;
+  int64_t laneCount = 1;
+  SmallVector<int64_t> keptAxes;
+  SmallVector<int64_t> keptAxisStrides;
  for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
-    if (isReduced) {
-      if (expandCompactToKeepdims.empty())
-        pendingLeadingReducedAxes.push_back(axis);
-      else
-        expandCompactToKeepdims.back().push_back(axis);
-      continue;
-    }
-
-    expandCompactToKeepdims.emplace_back();
-    auto& group = expandCompactToKeepdims.back();
-    group.append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
-    pendingLeadingReducedAxes.clear();
-    group.push_back(axis);
+    int64_t dim = keepdimsType.getDimSize(axis);
+    if (dim <= 0 || (isReduced && dim != 1))
+      return failure();
+    if (!isReduced)
+      keptAxes.push_back(axis);
  }
-  if (!pendingLeadingReducedAxes.empty())
-    expandCompactToKeepdims.back().append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
+  keptAxisStrides.resize(keptAxes.size(), 1);
+  for (int64_t index = static_cast<int64_t>(keptAxes.size()) - 1;
+       index >= 0; --index) {
+    keptAxisStrides[index] = laneCount;
+    int64_t dim = keepdimsType.getDimSize(keptAxes[index]);
+    if (laneCount > std::numeric_limits<int64_t>::max() / dim)
+      return failure();
+    laneCount *= dim;
+  }
+  if (batchType.getDimSize(0) != laneCount
+      || llvm::any_of(batchType.getShape().drop_front(),
+                      [](int64_t dim) { return dim != 1; }))
+    return failure();

-  if (batchType.getNumElements() != batchType.getDimSize(0))
-    return {};
-  auto reshapeCompute =
-    createSpatCompute<1>(rewriter, loc, TypeRange {keepdimsType}, {}, ValueRange {batchValue}, [&](Value input) {
-      auto flatType = RankedTensorType::get({batchType.getNumElements()}, batchType.getElementType(), batchType.getEncoding());
-      Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, input, collapseToFlat);
-      Value compact = flat;
-      if (compactKeptType != flatType)
-        compact = tensor::ExpandShapeOp::create(rewriter, loc, compactKeptType, flat, expandFlatToCompact);
-      Value keepdims = compact;
-      if (keepdimsType != compactKeptType)
-        keepdims = tensor::ExpandShapeOp::create(rewriter, loc, keepdimsType, compact, expandCompactToKeepdims);
-      spatial::SpatYieldOp::create(rewriter, loc, keepdims);
-    });
-  return reshapeCompute.getResult(0);
+  SmallVector<int64_t> operandIndices(laneCount, 0);
+  SmallVector<int64_t> sourceSlots;
+  SmallVector<int64_t> sourceOffsets(laneCount, 0);
+  SmallVector<int64_t> fragmentOffsets;
+  sourceSlots.reserve(laneCount);
+  fragmentOffsets.reserve(laneCount * rank);
+  for (int64_t lane = 0; lane < laneCount; ++lane) {
+    sourceSlots.push_back(lane);
+    size_t keptAxisIndex = 0;
+    for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
+      if (isReduced) {
+        fragmentOffsets.push_back(0);
+        continue;
+      }
+      int64_t dim = keepdimsType.getDimSize(axis);
+      fragmentOffsets.push_back(
+        (lane / keptAxisStrides[keptAxisIndex]) % dim);
+      ++keptAxisIndex;
+    }
+  }
+  SmallVector<int64_t> fragmentSizes(fragmentOffsets.size(), 1);
+  SmallVector<int64_t> fragmentStrides(fragmentOffsets.size(), 1);
+  return spatial::SpatBlueprintOp::create(
+           rewriter, loc, keepdimsType, batchValue, ValueRange {},
+           rewriter.getStringAttr("nchw"),
+           rewriter.getStringAttr("fragmented"),
+           rewriter.getDenseI64ArrayAttr(fragmentOffsets),
+           rewriter.getDenseI64ArrayAttr(fragmentSizes),
+           rewriter.getStringAttr("reduce_mean_keepdims_fragments"),
+           rewriter.getStringAttr("fragment_assembly"),
+           rewriter.getDenseI64ArrayAttr(operandIndices),
+           rewriter.getDenseI64ArrayAttr(sourceSlots),
+           rewriter.getDenseI64ArrayAttr(sourceOffsets),
+           rewriter.getDenseI64ArrayAttr(fragmentStrides),
+           rewriter.getStringAttr("disjoint"),
+           rewriter.getStringAttr("complete"))
+    .getOutput();
 }

 static SmallVector<ReassociationIndices> buildCollapseReassociation(ArrayRef<bool> reducedAxes) {
@@ -357,26 +369,36 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ReduceMeanOp> {

    Location loc = reduceMeanOp.getLoc();
    RankedTensorType leafType = getAllOnesType(inputType, resultType.getElementType());
-    RankedTensorType compactKeptType = getCompactKeptType(inputType, resultType.getElementType(), reducedAxes);
    RankedTensorType keepdimsType = getKeepdimsType(inputType, resultType.getElementType(), reducedAxes);
    int64_t laneCount = 1;
-    for (int64_t dim : compactKeptType.getShape())
+    for (auto [dim, isReduced] : llvm::zip_equal(keepdimsType.getShape(), reducedAxes)) {
+      if (isReduced)
+        continue;
+      if (dim <= 0 || laneCount > std::numeric_limits<int32_t>::max() / dim)
+        return rewriter.notifyMatchFailure(
+          reduceMeanOp, "ReduceMean physical lane count is not representable");
      laneCount *= dim;
+    }
    RankedTensorType batchType = getLanePackedKeepdimsType(laneCount, leafType);

    auto lanePackedKeepdims =
      buildReduceMeanKeepdimsBatch(adaptor.getData(), reducedAxes, batchType, leafType, rewriter, loc);
    if (failed(lanePackedKeepdims))
      return failure();
-    Value reducedKeepdims =
-      buildKeepdimsFromLanePackedBatch(*lanePackedKeepdims, keepdimsType, compactKeptType, reducedAxes, rewriter, loc);
+    auto reducedKeepdims = buildReduceMeanKeepdimsBlueprint(
+      *lanePackedKeepdims, keepdimsType, reducedAxes, rewriter, loc);
+    if (failed(reducedKeepdims))
+      return rewriter.notifyMatchFailure(
+        reduceMeanOp,
+        "cannot build physical-fragment ReduceMean keepdims reconstruction");

    if (semantics->keepdims != 0) {
-      rewriter.replaceOp(reduceMeanOp, reducedKeepdims);
+      rewriter.replaceOp(reduceMeanOp, *reducedKeepdims);
      return success();
    }

-    Value reduced = squeezeReducedAxes(reducedKeepdims, resultType, reducedAxes, rewriter, loc);
+    Value reduced = squeezeReducedAxes(
+      *reducedKeepdims, resultType, reducedAxes, rewriter, loc);
    rewriter.replaceOp(reduceMeanOp, reduced);
    return success();
  }