Cose belle

src/PIM/Conversion/ONNXToSpatial/CompileTime.cpp

@@ -9,10 +9,12 @@
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 
+#include <cstring>
 #include <utility>
 
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -21,24 +23,7 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
-namespace {
-
-static bool hasStaticUnitStrides(tensor::ExtractSliceOp extractSliceOp) {
-  return llvm::all_of(extractSliceOp.getStaticStrides(), [](int64_t stride) { return stride == 1; });
-}
-
-static bool hasConstantIndices(tensor::ExtractOp extractOp) {
-  return llvm::all_of(extractOp.getIndices(), [](Value index) { return matchConstantIndexValue(index).has_value(); });
-}
-
-static bool isStaticTensorResult(Operation* op) {
-  return llvm::all_of(op->getResultTypes(), [](Type type) {
-    auto shapedType = dyn_cast<ShapedType>(type);
-    return shapedType && shapedType.hasStaticShape();
-  });
-}
-
-static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr denseAttr, ArrayRef<int64_t> perms) {
+FailureOr<DenseElementsAttr> transposeDenseElementsAttr(DenseElementsAttr denseAttr, ArrayRef<int64_t> perms) {
   auto tensorType = dyn_cast<RankedTensorType>(denseAttr.getType());
   if (!tensorType)
     return failure();
@@ -59,7 +44,45 @@ static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr den
 
   auto transposedType = RankedTensorType::get(transposedShape, tensorType.getElementType(), tensorType.getEncoding());
   if (denseAttr.isSplat())
-    return DenseElementsAttr::get(transposedType, denseAttr.getSplatValue<Attribute>());
+    return DenseElementsAttr::getFromRawBuffer(transposedType, denseAttr.getRawData());
+
+  const unsigned elementBitWidth = tensorType.getElementTypeBitWidth();
+  const ArrayRef<char> inputData = denseAttr.getRawData();
+  if (elementBitWidth % 8 == 0) {
+    const size_t elementBytes = elementBitWidth / 8;
+    const size_t expectedBytes = denseAttr.getNumElements() * elementBytes;
+    if (inputData.size() == expectedBytes) {
+      SmallVector<char> transposedData(expectedBytes);
+      if (rank == 2 && perms[0] == 1 && perms[1] == 0) {
+        const int64_t rows = tensorType.getDimSize(0);
+        const int64_t columns = tensorType.getDimSize(1);
+        for (int64_t row = 0; row < rows; ++row)
+          for (int64_t column = 0; column < columns; ++column)
+            std::memcpy(transposedData.data() + (column * rows + row) * elementBytes,
+                        inputData.data() + (row * columns + column) * elementBytes,
+                        elementBytes);
+        return DenseElementsAttr::getFromRawBuffer(transposedType, transposedData);
+      }
+
+      SmallVector<int64_t> originalStrides = computeRowMajorStrides(tensorType.getShape());
+      SmallVector<int64_t> transposedStrides = computeRowMajorStrides(transposedShape);
+      SmallVector<int64_t> originalIndices(rank);
+      for (int64_t linearIndex = 0; linearIndex < tensorType.getNumElements(); ++linearIndex) {
+        int64_t remaining = linearIndex;
+        for (int64_t dim = 0; dim < rank; ++dim) {
+          originalIndices[dim] = originalStrides.empty() ? 0 : remaining / originalStrides[dim];
+          remaining = originalStrides.empty() ? 0 : remaining % originalStrides[dim];
+        }
+        int64_t transposedLinearIndex = 0;
+        for (int64_t dim = 0; dim < rank; ++dim)
+          transposedLinearIndex += originalIndices[perms[dim]] * transposedStrides[dim];
+        std::memcpy(transposedData.data() + transposedLinearIndex * elementBytes,
+                    inputData.data() + linearIndex * elementBytes,
+                    elementBytes);
+      }
+      return DenseElementsAttr::getFromRawBuffer(transposedType, transposedData);
+    }
+  }
 
   SmallVector<Attribute> originalValues(denseAttr.getValues<Attribute>());
   SmallVector<Attribute> transposedValues(originalValues.size());
@@ -84,16 +107,30 @@ static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr den
   return DenseElementsAttr::get(transposedType, transposedValues);
 }
 
+namespace {
+
+static bool hasStaticUnitStrides(tensor::ExtractSliceOp extractSliceOp) {
+  return llvm::all_of(extractSliceOp.getStaticStrides(), [](int64_t stride) { return stride == 1; });
+}
+
+static bool hasConstantIndices(tensor::ExtractOp extractOp) {
+  return llvm::all_of(extractOp.getIndices(), [](Value index) { return matchConstantIndexValue(index).has_value(); });
+}
+
+static bool isStaticTensorResult(Operation* op) {
+  return llvm::all_of(op->getResultTypes(), [](Type type) {
+    auto shapedType = dyn_cast<ShapedType>(type);
+    return shapedType && shapedType.hasStaticShape();
+  });
+}
+
 static FailureOr<DenseElementsAttr> reshapeDenseElements(DenseElementsAttr denseAttr, RankedTensorType resultType) {
   auto sourceType = dyn_cast<RankedTensorType>(denseAttr.getType());
-  if (!sourceType || !resultType || sourceType.getNumElements() != resultType.getNumElements())
+  if (!sourceType || !resultType || sourceType.getNumElements() != resultType.getNumElements()
+      || sourceType.getElementType() != resultType.getElementType())
     return failure();
 
-  if (denseAttr.isSplat())
-    return DenseElementsAttr::get(resultType, denseAttr.getSplatValue<Attribute>());
-
-  SmallVector<Attribute> values(denseAttr.getValues<Attribute>());
-  return DenseElementsAttr::get(resultType, values);
+  return DenseElementsAttr::getFromRawBuffer(resultType, denseAttr.getRawData());
 }
 
 static FailureOr<DenseElementsAttr> extractSliceDenseElements(DenseElementsAttr denseAttr,
@@ -161,7 +198,7 @@ static DenseElementsAttr getHostConstantDenseElementsAttrImpl(Value value, llvm:
     perm.reserve(transposeOp.getPermAttr().size());
     for (IntegerAttr attr : transposeOp.getPermAttr().getAsRange<IntegerAttr>())
       perm.push_back(attr.getInt());
-    auto transposedAttr = transposeDenseElements(inputAttr, perm);
+    auto transposedAttr = transposeDenseElementsAttr(inputAttr, perm);
     return succeeded(transposedAttr) ? *transposedAttr : nullptr;
   }
 
@@ -171,7 +208,7 @@ static DenseElementsAttr getHostConstantDenseElementsAttrImpl(Value value, llvm:
       return nullptr;
 
     SmallVector<int64_t> perm(transposeOp.getPermutation().begin(), transposeOp.getPermutation().end());
-    auto transposedAttr = transposeDenseElements(inputAttr, perm);
+    auto transposedAttr = transposeDenseElementsAttr(inputAttr, perm);
     return succeeded(transposedAttr) ? *transposedAttr : nullptr;
   }
 
@@ -219,6 +256,9 @@ getCompileTimeSourceImpl(Operation* op, llvm::SmallPtrSetImpl<Operation*>& visit
 
   chainLength += 1;
 
+  if (!isShapingOnlyOp(op))
+    return std::nullopt;
+
   if (auto extractOp = dyn_cast<tensor::ExtractOp>(op))
     return hasConstantIndices(extractOp)
            ? getCompileTimeSourceImpl(extractOp.getTensor().getDefiningOp(), visited, chainLength)

src/PIM/Conversion/ONNXToSpatial/CompileTime.hpp

@@ -4,6 +4,8 @@
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Value.h"
 
+#include "llvm/ADT/ArrayRef.h"
+
 namespace onnx_mlir {
 
 struct CompileTimeSource {
@@ -19,4 +21,7 @@ bool isCompileTimeOp(mlir::Operation* op);
 
 mlir::DenseElementsAttr getHostConstDenseElementsAttr(mlir::Value value);
 
+mlir::FailureOr<mlir::DenseElementsAttr> transposeDenseElementsAttr(
+  mlir::DenseElementsAttr denseAttr, llvm::ArrayRef<int64_t> permutation);
+
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -20,6 +20,7 @@
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -330,17 +331,25 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
     RewritePatternSet helperPatterns(ctx);
     populateGemmPatterns(helperPatterns, ctx);
     populateTransposePatterns(helperPatterns, ctx);
-    if (failed(applyPartialConversion(moduleOp, helperTarget, std::move(helperPatterns)))) {
-      moduleOp.emitError("failed to lower helper ONNX ops emitted by selected Spatial plan lowering");
-      signalPassFailure();
-      return;
+    FrozenRewritePatternSet frozenHelperPatterns(
+      std::move(helperPatterns));
+    SmallVector<Operation*> topLevelHelperOps;
+    funcOp.walk([&](Operation* op) {
+      if (isa<spatial::SpatGraphCompute,
+              spatial::SpatGraphComputeBatch>(op))
+        return WalkResult::skip();
+      if (isa<ONNXGemmOp, ONNXTransposeOp>(op))
+        topLevelHelperOps.push_back(op);
+      return WalkResult::advance();
+    });
+    for (Operation *helper : topLevelHelperOps) {
+      if (failed(applyPartialConversion(
+            helper, helperTarget, frozenHelperPatterns))) {
+        moduleOp.emitError("failed to lower helper ONNX ops emitted by selected Spatial plan lowering");
+        signalPassFailure();
+        return;
+      }
     }
-    FrozenRewritePatternSet nestedHelperPatterns([&] {
-      RewritePatternSet patterns(ctx);
-      populateGemmPatterns(patterns, ctx);
-      populateTransposePatterns(patterns, ctx);
-      return patterns;
-    }());
     ConversionTarget nestedHelperTarget(*ctx);
     nestedHelperTarget.addLegalDialect<spatial::SpatialDialect,
                                        tensor::TensorDialect,
@@ -356,7 +365,8 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         computeLikeOps.push_back(op);
     });
     for (Operation* op : computeLikeOps) {
-      if (failed(applyFullConversion(op, nestedHelperTarget, nestedHelperPatterns))) {
+      if (failed(applyFullConversion(
+            op, nestedHelperTarget, frozenHelperPatterns))) {
         op->emitOpError("failed to lower nested helper ONNX ops emitted by selected Spatial plan lowering");
         signalPassFailure();
         return;
@@ -392,6 +402,12 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
       signalPassFailure();
     } else {
       dumpModule(moduleOp, "spatial1_graph");
+      spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
+      if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Spatial1)
+          && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) {
+        signalPassFailure();
+        return;
+      }
     }
 
     if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp

@@ -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();
   }

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Transpose.cpp

@@ -5,7 +5,7 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -52,35 +52,12 @@ static FailureOr<Value> materializeTransposedConstant(Value input,
     return failure();
   }
 
-  if (denseAttr.isSplat())
-    return getOrCreateConstant(rewriter,
-                               rewriter.getInsertionBlock()->getParentOp(),
-                               DenseElementsAttr::get(resultType, denseAttr.getSplatValue<Attribute>()),
-                               resultType);
-
-  SmallVector<Attribute> inputValues(denseAttr.getValues<Attribute>());
-  SmallVector<Attribute> resultValues(inputValues.size());
-  SmallVector<int64_t> inputStrides = computeRowMajorStrides(inputType.getShape());
-  SmallVector<int64_t> resultStrides = computeRowMajorStrides(resultType.getShape());
-  SmallVector<int64_t> inputIndices(inputType.getRank(), 0);
-
-  for (auto [linearIndex, value] : llvm::enumerate(inputValues)) {
-    int64_t remaining = static_cast<int64_t>(linearIndex);
-    for (int64_t dim = 0; dim < inputType.getRank(); ++dim) {
-      inputIndices[dim] = inputStrides.empty() ? 0 : remaining / inputStrides[dim];
-      remaining = inputStrides.empty() ? 0 : remaining % inputStrides[dim];
-    }
-
-    int64_t resultLinearIndex = 0;
-    for (int64_t dim = 0; dim < resultType.getRank(); ++dim)
-      resultLinearIndex += inputIndices[permutation[dim]] * resultStrides[dim];
-
-    resultValues[resultLinearIndex] = value;
-  }
-
+  auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation);
+  if (failed(transposedAttr) || transposedAttr->getType() != resultType)
+    return failure();
   return getOrCreateConstant(rewriter,
                              rewriter.getInsertionBlock()->getParentOp(),
-                             DenseElementsAttr::get(resultType, resultValues),
+                             *transposedAttr,
                              resultType);
 }