[mlir][tosa] Switch TosaFoldConstantTranspose to use ElementsAttr.

[jpienaar]

Summary

Also avoid redoing index calculation.

mlir/lib/Dialect/Tosa/Transforms/TosaFoldConstantTranspose.cpp

struct TosaFoldConstantTranspose : public OpRewritePattern<tosa::TransposeOp> {
 
   LogicalResult matchAndRewrite(tosa::TransposeOp op,
                                 PatternRewriter &rewriter) const override {
     auto outputType = op.getType().cast<ShapedType>();
     // TOSA supports quantized types.
     if (!outputType.getElementType().isIntOrIndexOrFloat())
       return failure();
 
-    DenseElementsAttr inputValues;
+    ElementsAttr inputValues;
     if (!matchPattern(op.getInput1(), m_Constant(&inputValues)))
       return failure();
     // Make sure the input is a constant that has a single user.
     if (!llvm::hasSingleElement(op.getInput1().getDefiningOp()->getUsers()))
       return failure();
 
     DenseIntElementsAttr permAttr;
     if (!matchPattern(op.getPerms(), m_Constant(&permAttr)))
LogicalResult matchAndRewrite(tosa::TransposeOp op,
     SmallVector<Attribute, 4> outputValues;
     outputValues.resize(numElements);
 
     // Transpose the input constant. Because we don't know its rank in advance,
     // we need to loop over the range [0, element count) and delinearize the
     // index.
     auto attrValues = inputValues.getValues<Attribute>();
     ArrayRef<int64_t> outputShape = outputType.getShape();
-    for (int srcLinearIndex = 0; srcLinearIndex < numElements;
-         ++srcLinearIndex) {
+    for (const auto &it : llvm::enumerate(attrValues)) {
       SmallVector<uint64_t, 6> srcIndices(inputType.getRank(), 0);
-      int totalCount = srcLinearIndex;
+      int totalCount = it.index();
       for (int dim = inputType.getRank() - 1; dim >= 0; --dim) {
         srcIndices[dim] = totalCount % inputShape[dim];
         totalCount /= inputShape[dim];
       }
 
       SmallVector<uint64_t, 6> dstIndices(outputType.getRank(), 0);
       for (int dim = outputType.getRank() - 1; dim >= 0; --dim)
         dstIndices[dim] = srcIndices[permValues[dim]];
 
       uint64_t dstLinearIndex = dstIndices.front();
       for (int dim = 1; dim < outputType.getRank(); ++dim)
         dstLinearIndex = dstLinearIndex * outputShape[dim] + dstIndices[dim];
 
-      outputValues[dstLinearIndex] = attrValues[srcIndices];
+      outputValues[dstLinearIndex] = it.value();
     }
 
     rewriter.replaceOpWithNewOp<tosa::ConstOp>(
         op, outputType, DenseElementsAttr::get(outputType, outputValues));
     return success();
   }
 };
 
 } // namespace
 
 void mlir::tosa::populateTosaFoldConstantTransposePatterns(
     MLIRContext *ctx, RewritePatternSet &patterns) {
   patterns.add<TosaFoldConstantTranspose>(ctx);
 }