Clean-up for TOSA pow folding

Address review comments:
* Add names to all function arguments in header
* `toBeBroadcasted` -> `toBeBroadcastedShape`
* Describe return values for broadcasting functions
* Add shortcut for the offset computation if no broadcasting is needed

Author: TinaAMD

mlir/include/mlir/Dialect/Tosa/Transforms/TosaFoldCommon.h

@@ -36,22 +36,23 @@ DenseElementsAttr applyElementWise(
 /// tensors. If the input tensors do not match \p targetType, broadcasting is
 /// applied.
 DenseElementsAttr applyElementWise(
-    const DenseElementsAttr &, const DenseElementsAttr &, TensorType targetType,
+    const DenseElementsAttr &first, const DenseElementsAttr &second,
+    TensorType targetType,
     const std::function<APFloat(const APFloat &, const APFloat &)> &toApply);
 
 /// Function that checks if \p toCheck is a dense TOSA constant float tensor.
 LogicalResult notifyIfNotConstantFloatTosaTensor(TypedValue<TensorType> toCheck,
                                                  TosaOp location,
-                                                 PatternRewriter &);
+                                                 PatternRewriter &rewriter);
 
 /// Function that checks if \p toCheck is a dense TOSA constant tensor.
 LogicalResult notifyIfNoTosaDenseConstantTensor(TypedValue<TensorType> toCheck,
                                                 TosaOp location,
-                                                PatternRewriter &);
+                                                PatternRewriter &rewriter);
 
 /// Function that checks if the type contained in \p toCheck is float.
 LogicalResult notifyIfNotFloat(TypedValue<TensorType> toCheck, TosaOp location,
-                               PatternRewriter &);
+                               PatternRewriter &rewriter);
 
 /// Compute the offset in \p shape which corresponds to the given \p index.
 OffsetType indexToOffset(DimensionType shape, DimensionType index);
@@ -60,18 +61,20 @@ OffsetType indexToOffset(DimensionType shape, DimensionType index);
 SmallVector<int64_t> offsetToIndex(DimensionType shape, OffsetType offset);
 
 /// Given an \p index into \p desiredShape, compute the corresponding index into
-/// \p toBeBroadcasted.
+/// \p toBeBroadcastedShape.
+/// \returns broadcasted index into \p toBeBroadcastedShape.
 SmallVector<int64_t> getBroadcastedIndex(DimensionType desiredShape,
-                                         DimensionType toBeBroadcasted,
+                                         DimensionType toBeBroadcastedShape,
                                          DimensionType index);
 /// Given an \p offset into \p desiredShape, compute the corresponding offset
-/// into \p toBeBroadcasted.
+/// into \p toBeBroadcastedShape.
+/// \returns broadcasted offset into \p toBeBroadcastedShape.
 OffsetType getBroadcastedOffset(DimensionType desiredShape,
-                                DimensionType toBeBroadcasted,
+                                DimensionType toBeBroadcastedShape,
                                 OffsetType offset);
 
 /// Function to compute the reciprocal.
-APFloat computeReciprocal(const APFloat &, Type);
+APFloat computeReciprocal(const APFloat &floatVal, Type floatTy);
 
 } // namespace tosa
 } // namespace mlir

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

@@ -12,9 +12,9 @@
 
 #include "mlir/Dialect/Tosa/Transforms/TosaFoldCommon.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
-#include <algorithm>
 #include <llvm/ADT/APFloat.h>
 #include <llvm/ADT/SmallVector.h>
+#include <algorithm>
 #include <mlir/IR/BuiltinAttributes.h>
 #include <mlir/IR/BuiltinTypes.h>
 #include <mlir/IR/Matchers.h>
@@ -155,13 +155,13 @@ SmallVector<int64_t> mlir::tosa::offsetToIndex(DimensionType shape,
 
 SmallVector<int64_t>
 mlir::tosa::getBroadcastedIndex(DimensionType desiredShape,
-                                DimensionType toBeBroadcasted,
+                                DimensionType toBeBroadcastedShape,
                                 DimensionType index) {
   SmallVector<int64_t> broadCasted;
   broadCasted.reserve(desiredShape.size());
   for (size_t i = 0; i < desiredShape.size(); i++) {
     auto toInsert = 0;
-    if (toBeBroadcasted[i] == desiredShape[i]) {
+    if (toBeBroadcastedShape[i] == desiredShape[i]) {
       toInsert = index[i];
     }
     broadCasted.push_back(toInsert);
@@ -170,12 +170,16 @@ mlir::tosa::getBroadcastedIndex(DimensionType desiredShape,
 }
 
 OffsetType mlir::tosa::getBroadcastedOffset(DimensionType desiredShape,
-                                            DimensionType toBeBroadcasted,
+                                            DimensionType toBeBroadcastedShape,
                                             OffsetType offset) {
+  // Simply return the offset if the shapes are equal.
+  if (desiredShape.equals(toBeBroadcastedShape)) {
+    return offset;
+  }
   auto indexInTarget = offsetToIndex(desiredShape, offset);
   auto indexBroadcasted =
-      getBroadcastedIndex(desiredShape, toBeBroadcasted, indexInTarget);
-  return indexToOffset(toBeBroadcasted, indexBroadcasted);
+      getBroadcastedIndex(desiredShape, toBeBroadcastedShape, indexInTarget);
+  return indexToOffset(toBeBroadcastedShape, indexBroadcasted);
 }
 
 APFloat mlir::tosa::computeReciprocal(const APFloat &floatVal, Type floatTy) {