[mlir][tosa] Constant folding for reciprocal

[TinaAMD]

Summary

Add constant fold for tosa.reciprocal, which can be applied if the input is a dense constant tensor. The reciprocal is computed for every element and the result is a tensor with the same dimensions as the input tensor.

As the input tensor might require a lot of memory and the folding might double the required memory, a heuristic decides when to actually apply the folding. Currently, the operation will be replaced only if the input constant is a splat (i.e. requires little memory) or has in single user (similar to the already existing fold for constant transposes). This keeps the additionally required space low.

We implemented folds for various operators and would like to contribute them successively. Those folds are for:

• tosa.add
• tosa.cast
• tosa.clamp
• tosa.mul
• tosa.pow
• tosa.rsqrt

[eric-k256]

This seems like a reasonable fold to me, and a good heuristic to start with. Adding @jpienaar in case he sees anything I've missed.

[TinaAMD]

Ping @jpienaar

[rsuderman]

I can land now

[rsuderman]

One change before landing. Everything is broken up between 3-4 files of <100 lines in length. It doesn't seem worth fracturing this up to this degree. Could you have TosaFolders.cpp and merge them all together, including the transpose pass)

[TinaAMD]

Thanks for the feedback!

I merged the changes into a single new file TosaFolders.cpp, it also includes the already pre-existing transpose fold.

[TinaAMD]

If you are happy with the revision, could you commit the changes for me? As this is my first contribution, I don't have commit access yet.
Would sending you my email address by Discord work for you (username tinaamd there as well)?

I would suggest parts of the description as commit message:
Add constant fold for tosa.reciprocal, which can be applied if the input is a dense constant tensor. The reciprocal is computed for every element and the result is a tensor with the same dimensions as the input tensor.

As the input tensor might require a lot of memory and the folding might double the required memory, a heuristic decides when to actually apply the folding. Currently, the operation will be replaced only if the input constant is a splat (i.e. requires little memory) or has in single user (similar to the already existing fold for constant transposes). This keeps the additionally required space low.

[mgehre-amd]

In D150578#4470338, @TinaAMD wrote:

If you are happy with the revision, could you commit the changes for me? As this is my first contribution, I don't have commit access yet.
Would sending you my email address by Discord work for you (username tinaamd there as well)?

I would suggest parts of the description as commit message:
Add constant fold for tosa.reciprocal, which can be applied if the input is a dense constant tensor. The reciprocal is computed for every element and the result is a tensor with the same dimensions as the input tensor.

As the input tensor might require a lot of memory and the folding might double the required memory, a heuristic decides when to actually apply the folding. Currently, the operation will be replaced only if the input constant is a splat (i.e. requires little memory) or has in single user (similar to the already existing fold for constant transposes). This keeps the additionally required space low.

Hey, I'm happy to commit this patch for you.

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

 //===-- Passes.h - TOSA optimization pass declarations ----------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 
 // Expose Rewrite Functions that decompose TOSA Ops into further TOSA Ops.
 // The rewrites can be selectively added to a conversion pass.
 void populateTosaDecomposeConv2D(MLIRContext *ctx, RewritePatternSet &patterns);
 void populateTosaDecomposeTransposeConv(MLIRContext *ctx,
                                         RewritePatternSet &patterns);
 void populateTosaDecomposeDepthwise(MLIRContext *ctx,
                                     RewritePatternSet &patterns);
+void populateTosaFoldConstantReciprocalPatterns(MLIRContext *ctx,
+                                                RewritePatternSet &patterns);
 void populateTosaFoldConstantTransposePatterns(MLIRContext *ctx,
                                                RewritePatternSet &patterns);
 
 std::unique_ptr<Pass> createTosaLayerwiseConstantFoldPass();
 std::unique_ptr<Pass> createTosaInferShapesPass();
 std::unique_ptr<Pass> createTosaMakeBroadcastablePass();
 std::unique_ptr<Pass> createTosaTestQuantUtilAPIPass();
 std::unique_ptr<Pass> createTosaOptionalDecompositions();

mlir/include/mlir/Dialect/Tosa/Utils/FoldUtils.h

+//===- FoldUtils.h - Helper Functions for Folds -----------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Helper functions useful for various different TOSA constant folds.
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_DIALECT_TOSA_UTILS_FOLD_UTILS_H
+#define MLIR_DIALECT_TOSA_UTILS_FOLD_UTILS_H
+
+#include <functional>
+#include <mlir/IR/BuiltinAttributes.h>
+
+namespace mlir {
+namespace tosa {
+
+/// Rounding mode to be used on floating point operations that require rounding.
+static constexpr llvm::RoundingMode tosaRoundingMode =
+    llvm::APFloat::rmNearestTiesToEven;
+
+/// Apply the given transformation \p toApply to every element of the tensor to
+/// be transformed \p toTransform.
+///
+/// Elements of \p toTransform are extracted as \p SrcValueType.
+///
+/// \returns A tensor with the same size as \p toTransform, containing
+/// \p TargetValueType values of type \p TargetType.
+template <class SrcValType, class TargetValType, class TargetType>
+DenseElementsAttr applyElementWise(
+    const DenseElementsAttr &toTransform,
+    const std::function<TargetValType(const SrcValType &, TargetType)> &toApply,
+    TargetType targetType);
+
+} // namespace tosa
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TOSA_UTILS_FOLD_UTILS_H

mlir/lib/Dialect/Tosa/CMakeLists.txt

 add_mlir_dialect_library(MLIRTosaDialect
   IR/TosaOps.cpp
   IR/TosaCanonicalizations.cpp
   Utils/ConversionUtils.cpp
+  Utils/FoldUtils.cpp
   Utils/QuantUtils.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tosa
 
   DEPENDS
   MLIRTosaAttributesIncGen
   MLIRTosaOpsIncGen

mlir/lib/Dialect/Tosa/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRTosaTransforms
   TosaDecomposeTransposeConv.cpp
   TosaDecomposeConv2D.cpp
   TosaDecomposeDepthwise.cpp
+  TosaFoldCommon.cpp
+  TosaFoldConstantReciprocal.cpp
   TosaFoldConstantTranspose.cpp
   TosaInferShapes.cpp
   TosaLayerwiseConstantFoldPass.cpp
   TosaMakeBroadcastable.cpp
   TosaOptionalDecompositions.cpp
   TosaValidation.cpp
 
   ADDITIONAL_HEADER_DIRS

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

+//===- TosaFoldCommon.h - Helper Functions for Folds ------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Helper functions useful for various different TOSA constant folds.
+//
+//===----------------------------------------------------------------------===//
+#ifndef MLIR_DIALECT_TOSA_TRANSFORMS_TOSA_FOLD_COMMON_H
+#define MLIR_DIALECT_TOSA_TRANSFORMS_TOSA_FOLD_COMMON_H
+
+#include <llvm/ADT/APFloat.h>
+#include <mlir/Dialect/Tosa/IR/TosaOps.h>
+#include <mlir/IR/PatternMatch.h>
+
+namespace mlir {
+namespace tosa {
+
+/// Function that checks if \p toCheck is a dense TOSA constant float tensor.
+LogicalResult notifyIfNotConstantFloatTosaTensor(TypedValue<TensorType> toCheck,
+                                                 TosaOp location,
+                                                 PatternRewriter &rewriter);
+
+/// Function that checks if \p toCheck is a dense TOSA constant tensor.
+LogicalResult notifyIfNoTosaDenseConstantTensor(TypedValue<TensorType> toCheck,
+                                                TosaOp location,
+                                                PatternRewriter &rewriter);
+
+/// Function that checks if the type contained in \p toCheck is float.
+LogicalResult notifyIfNotFloat(TypedValue<TensorType> toCheck, TosaOp location,
+                               PatternRewriter &rewriter);
+
+/// Heuristic to decide when to replace a unary operation on a constant with the
+/// folded value.
+/// Folding operations on constants can lead to an increased memory usage
+/// whenever the input cannot be replaced but a new constant is inserted. Hence,
+/// this will currently only suggest folding when the memory impact is
+/// negligible.
+/// Takes the \p unaryOp and the constant input \p values.
+/// \returns Whether folding should be applied.
+bool constantUnaryOpShouldBeFolded(TosaOp unaryOp, DenseElementsAttr values);
+
+} // namespace tosa
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TOSA_TRANSFORMS_TOSA_FOLD_COMMON_H

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

+//===- TosaFoldCommon.cpp -------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Helper functions useful for various different TOSA constant folds.
+//
+//===----------------------------------------------------------------------===//
+
+#include "TosaFoldCommon.h"
+#include <mlir/Dialect/Tosa/IR/TosaOps.h>
+#include <mlir/IR/BuiltinAttributes.h>
+#include <mlir/IR/BuiltinTypes.h>
+#include <mlir/IR/Matchers.h>
+#include <mlir/Support/LogicalResult.h>
+
+using namespace mlir;
+using namespace mlir::tosa;
+
+LogicalResult
+mlir::tosa::notifyIfNotConstantFloatTosaTensor(TypedValue<TensorType> toCheck,
+                                               TosaOp location,
+                                               PatternRewriter &rewriter) {
+  auto floatCheck = notifyIfNotFloat(toCheck, location, rewriter);
+  if (failed(floatCheck)) {
+    return floatCheck;
+  }
+  return notifyIfNoTosaDenseConstantTensor(toCheck, location, rewriter);
+}
+
+LogicalResult
+mlir::tosa::notifyIfNoTosaDenseConstantTensor(TypedValue<TensorType> toCheck,
+                                              TosaOp location,
+                                              PatternRewriter &rewriter) {
+  // Check whether the tensor is constant and dense
+  // TODO We currently ensure the tensor is dense by using the correct type for
+  // the bind_value, however we do not actually need this value. It would be
+  // nicer to only have a check here.
+  DenseElementsAttr tmp;
+  if (!matchPattern(toCheck, m_Constant(&tmp))) {
+    return rewriter.notifyMatchFailure(location,
+                                       "Non-const or non-dense input tensor");
+  }
+
+  // Make sure it actually is a TOSA constant (the match allows for other
+  // constants as well)
+  if (isa<ConstOp>(toCheck.getDefiningOp())) {
+    return success();
+  }
+
+  return rewriter.notifyMatchFailure(location,
+                                     "The reciprocal can only be folded if "
+                                     "it operates on a TOSA constant");
+}
+
+LogicalResult mlir::tosa::notifyIfNotFloat(TypedValue<TensorType> toCheck,
+                                           TosaOp location,
+                                           PatternRewriter &rewriter) {
+  if (isa<FloatType>(toCheck.getType().getElementType())) {
+    return success();
+  }
+  return rewriter.notifyMatchFailure(location,
+                                     "Unexpected input tensor type: the "
+                                     "TOSA spec only allows floats");
+}
+
+bool mlir::tosa::constantUnaryOpShouldBeFolded(TosaOp unaryOp,
+                                               DenseElementsAttr values) {
+  assert(unaryOp->getNumOperands() == 1);
+  auto inputOp = unaryOp->getOperand(0);
+
+  // If the input is a splat, we don't care for the number of users
+  if (isa<SplatElementsAttr>(values)) {
+    return true;
+  }
+
+  // If this is the only use of the tensor it should be replaced as no
+  // additional memory is required
+  return inputOp.hasOneUse();
+}

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

+//===- TosaFoldConstantReciprocal.cpp -------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Fold TOSA reciprocal operation on constant data
+//
+//===----------------------------------------------------------------------===//
+
+#include "TosaFoldCommon.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/Dialect/Tosa/Transforms/Passes.h"
+#include "mlir/Dialect/Tosa/Utils/FoldUtils.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Matchers.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/FloatingPointMode.h"
+#include "llvm/ADT/SmallVector.h"
+
+using namespace mlir;
+using namespace mlir::tosa;
+
+namespace {
+
+struct TosaFoldConstantReciprocal : public OpRewritePattern<ReciprocalOp> {
+
+  using OpRewritePattern::OpRewritePattern;
+
+  static APFloat computeReciprocal(const APFloat &floatVal, FloatType floatTy) {
+    auto recipAttr = FloatAttr::get(floatTy, 1.0);
+    APFloat recip = recipAttr.getValue();
+    recip.divide(floatVal, tosaRoundingMode);
+
+    return recip;
+  }
+
+  LogicalResult matchAndRewrite(ReciprocalOp recip,
+                                PatternRewriter &rewriter) const override {
+    auto inputTensor = recip.getInput1();
+
+    // Check that we can apply folding
+    auto preCondCheck =
+        notifyIfNotConstantFloatTosaTensor(inputTensor, recip, rewriter);
+    if (failed(preCondCheck)) {
+      return preCondCheck;
+    }
+
+    // Extract the tensor values
+    DenseElementsAttr inputValues;
+    matchPattern(inputTensor, m_Constant(&inputValues));
+
+    // Check whether this should be folded.
+    if (!constantUnaryOpShouldBeFolded(recip, inputValues)) {
+      return rewriter.notifyMatchFailure(
+          recip, "Currently, reciprocals will only be folded if the input "
+                 "tensor has a single user");
+    }
+
+    // Create a new tensor with the updated values
+    auto newTensor = applyElementWise<APFloat, APFloat, FloatType>(
+        inputValues, &computeReciprocal,
+        cast<FloatType>(inputValues.getElementType()));
+
+    // Replace the use of the reciprocal with the transformed tensor
+    rewriter.replaceOpWithNewOp<ConstOp>(recip, newTensor.getType(), newTensor);
+    return success();
+  }
+};
+
+} // namespace
+
+void mlir::tosa::populateTosaFoldConstantReciprocalPatterns(
+    MLIRContext *ctx, RewritePatternSet &patterns) {
+  patterns.add<TosaFoldConstantReciprocal>(ctx);
+}

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

 //===- TosaLayerwiseConstantFoldPass.cpp ----------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //

 struct TosaLayerwiseConstantFoldPass
     : public tosa::impl::TosaLayerwiseConstantFoldPassBase<
           TosaLayerwiseConstantFoldPass> {
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
     auto func = getOperation();
 
+    mlir::tosa::populateTosaFoldConstantReciprocalPatterns(ctx, patterns);
     mlir::tosa::populateTosaFoldConstantTransposePatterns(ctx, patterns);
     populateTosaOpsCanonicalizationPatterns(ctx, patterns);
 
     if (applyPatternsAndFoldGreedily(func, std::move(patterns)).failed())
       signalPassFailure();
   }
 };
 
 } // namespace
 
 std::unique_ptr<Pass> mlir::tosa::createTosaLayerwiseConstantFoldPass() {
   return std::make_unique<TosaLayerwiseConstantFoldPass>();
 }

mlir/lib/Dialect/Tosa/Utils/FoldUtils.cpp

+//===- FoldUtils.cpp ------------------------------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Helper functions useful for various different TOSA constant folds.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Tosa/Utils/FoldUtils.h"
+
+#include <llvm/ADT/APFloat.h>
+#include <llvm/ADT/SmallVector.h>
+#include <mlir/IR/BuiltinAttributes.h>
+#include <mlir/IR/BuiltinTypes.h>
+
+using namespace mlir;
+using namespace mlir::tosa;
+
+template <class SrcValType, class TargetValType, class TargetType>
+DenseElementsAttr mlir::tosa::applyElementWise(
+    const DenseElementsAttr &toTransform,
+    const std::function<TargetValType(const SrcValType &, TargetType)> &toApply,
+    TargetType targetType) {
+  SmallVector<TargetValType> transformedValues;
+  // We already know the amount of values we will insert, reserve space for
+  // all of them to avoid dynamic resizing
+  transformedValues.reserve(toTransform.getNumElements());
+  for (auto val : toTransform.getValues<SrcValType>()) {
+    auto transformedVal = toApply(val, targetType);
+    transformedValues.push_back(transformedVal);
+  }
+
+  // Make sure that the output tensor has the expected output type
+  auto inShape = toTransform.getType();
+  auto outTy = inShape.cloneWith({}, targetType);
+
+  return DenseElementsAttr::get(outTy, transformedValues);
+}
+
+template DenseElementsAttr
+mlir::tosa::applyElementWise<APFloat, APFloat, FloatType>(
+    const DenseElementsAttr &toTransform,
+    const std::function<APFloat(const APFloat &, FloatType)> &toApply,
+    FloatType targetType);

mlir/test/Dialect/Tosa/constant-reciprocal-fold.mlir

+// RUN: mlir-opt --split-input-file --tosa-layerwise-constant-fold %s | FileCheck %s
+
+// CHECK-LABEL: @reciprocal_fold_single_valued
+func.func @reciprocal_fold_single_valued() -> tensor<f32> {
+  // CHECK: [[RES:]] ={{.*}}tosa.const{{.*}}2.5{{0*}}e-01{{.*}}tensor<f32>
+  // CHECK-NOT: tosa.reciprocal
+  // CHECK: return [[RES]]
+  %0 = "tosa.const"() {value = dense<4.0> : tensor<f32>} : () -> tensor<f32>
+  %1 = "tosa.reciprocal"(%0) : (tensor<f32>) -> tensor<f32>
+  return %1 : tensor<f32>
+}
+
+// CHECK-LABEL: @reciprocal_fold_splat
+func.func @reciprocal_fold_splat() -> tensor<12x7xf32> {
+  // CHECK: [[RES:]] ={{.*}}tosa.const{{.*}}2.5{{0*}}e-01{{.*}}tensor<12x7xf32>
+  // CHECK-NOT: tosa.reciprocal
+  // CHECK: return [[RES]]
+  %0 = "tosa.const"() {value = dense<4.0> : tensor<12x7xf32>} : () -> tensor<12x7xf32>
+  %1 = "tosa.reciprocal"(%0) : (tensor<12x7xf32>) -> tensor<12x7xf32>
+  return %1 : tensor<12x7xf32>
+}
+
+// CHECK-LABEL: @reciprocal_div_zero
+func.func @reciprocal_div_zero() -> tensor<f32> {
+  // 0x7F800000 is the value for +infinity
+  // CHECK: [[RES:]] ={{.*}}tosa.const{{.*}}0x7F800000
+  // CHECK-NOT: tosa.reciprocal
+  // CHECK: return [[RES]]
+  %0 = "tosa.const"() {value = dense<0.0> : tensor<f32>} : () -> tensor<f32>
+  %1 = "tosa.reciprocal"(%0) : (tensor<f32>) -> tensor<f32>
+  return %1 : tensor<f32>
+}
+
+// CHECK-LABEL: @reciprocal_div_neg_zero
+func.func @reciprocal_div_neg_zero() -> tensor<f32> {
+  // 0xFF800000 is the value for -infinity
+  // CHECK: [[RES:]] ={{.*}}tosa.const{{.*}}0xFF800000
+  // CHECK-NOT: tosa.reciprocal
+  // CHECK: return [[RES]]
+  %0 = "tosa.const"() {value = dense<-0.0> : tensor<f32>} : () -> tensor<f32>
+  %1 = "tosa.reciprocal"(%0) : (tensor<f32>) -> tensor<f32>
+  return %1 : tensor<f32>
+}
+
+// CHECK-LABEL: @reciprocal_div_nan
+func.func @reciprocal_div_nan() -> tensor<f32> {
+  // 0x7FC00000 is the value for NAN
+  // CHECK: [[RES:]] ={{.*}}tosa.const{{.*}}0x7FC00000
+  // CHECK-NOT: tosa.reciprocal
+  // CHECK: return [[RES]]
+  %0 = "tosa.const"() {value = dense<0x7FC00000> : tensor<f32>} : () -> tensor<f32>
+  %1 = "tosa.reciprocal"(%0) : (tensor<f32>) -> tensor<f32>
+  return %1 : tensor<f32>
+}
+
+// CHECK-LABEL: @reciprocal_div_infinity
+func.func @reciprocal_div_infinity() -> tensor<f32> {
+  // CHECK: [[RES:]] ={{.*}}tosa.const{{.*}}<0.{{0*}}e+00>
+  // CHECK-NOT: tosa.reciprocal
+  // CHECK: return [[RES]]
+  %0 = "tosa.const"() {value = dense<0x7F800000> : tensor<f32>} : () -> tensor<f32>
+  %1 = "tosa.reciprocal"(%0) : (tensor<f32>) -> tensor<f32>
+  return %1 : tensor<f32>
+}
+
+// CHECK-LABEL: @reciprocal_div_neg_infinity
+func.func @reciprocal_div_neg_infinity() -> tensor<f32> {
+  // CHECK: [[RES:]] ={{.*}}tosa.const{{.*}}<-0.{{0*}}e+00>
+  // CHECK-NOT: tosa.reciprocal
+  // CHECK: return [[RES]]
+  %0 = "tosa.const"() {value = dense<0xFF800000> : tensor<f32>} : () -> tensor<f32>
+  %1 = "tosa.reciprocal"(%0) : (tensor<f32>) -> tensor<f32>
+  return %1 : tensor<f32>
+}
+
+// CHECK-LABEL: @reciprocal_div_underflow
+func.func @reciprocal_div_underflow() -> tensor<2xf16> {
+  // CHECK: [[RES:]] ={{.*}}tosa.const{{.*}}-0.{{0*}}e+00, 0.{{0*}}e+00
+  // CHECK-NOT: tosa.reciprocal
+  // CHECK: return [[RES]]
+  %0 = "tosa.const"() {value = dense<[-6.0e+15, 6.0e+15]> : tensor<2xf16>} : () -> tensor<2xf16>
+  %1 = "tosa.reciprocal"(%0) : (tensor<2xf16>) -> tensor<2xf16>
+  return %1 : tensor<2xf16>
+}
+
+// CHECK-LABEL: @reciprocal_div_overflow
+func.func @reciprocal_div_overflow() -> tensor<2xf16> {
+  // CHECK: [[RES:]] ={{.*}}tosa.const{{.*}}0x7C00, 0xFC00
+  // CHECK-NOT: tosa.reciprocal
+  // CHECK: return [[RES]]
+  %0 = "tosa.const"() {value = dense<[0.0000001, -0.0000001]> : tensor<2xf16>} : () -> tensor<2xf16>
+  %1 = "tosa.reciprocal"(%0) : (tensor<2xf16>) -> tensor<2xf16>
+  return %1 : tensor<2xf16>
+}
+
+// CHECK-LABEL: @reciprocal_no_fold
+// The folding optimization works only intra-procedurally, so we won't be able
+// to fold anything here
+func.func @reciprocal_no_fold(%arg0: tensor<?x?xf32>) -> tensor<?x?xf32> {
+  // CHECK: tosa.reciprocal
+  // CHECK-NEXT: return
+  %0 = "tosa.reciprocal"(%arg0) : (tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+// CHECK-LABEL: @reciprocal_fold
+func.func @reciprocal_fold() -> tensor<4x6xf32> {
+  // CHECK: [[RES:]] ={{.*}}tosa.const
+  // CHECK-SAME{LITERAL}: [[5.68828249, 11.4416485, 1.6880486, 0.680272102, -0.875350117, 0.342313349],
+  // CHECK-SAME{LITERAL}:  [-4.81231928, 0.698080301, 0.65432179, -82.6446304, -4.33651352, -0.747551739],
+  // CHECK-SAME{LITERAL}:  [-12.4378109, 13.140605, 1.89501607, 0.885582745, 4.08830738, 1.4396776],
+  // CHECK-SAME{LITERAL}:  [2.02880907, -1.53280187, 0.552730501, 7.15819644, 0.64495325, -0.973709881]]
+  // CHECK-NOT: tosa.reciprocal
+  // CHECK: return [[RES]]
+  %0 = "tosa.const"() { value = dense<[
+                        [ 0.1758,  0.0874,  0.5924,  1.4700, -1.1424,  2.9213],
+                        [-0.2078,  1.4325,  1.5283, -0.0121, -0.2306, -1.3377],
+                        [-0.0804,  0.0761,  0.5277,  1.1292,  0.2446,  0.6946],
+                        [ 0.4929, -0.6524,  1.8092,  0.1397,  1.5505, -1.0270]]>
+                        : tensor<4x6xf32>
+                      } : () -> tensor<4x6xf32>
+  %1 = "tosa.reciprocal"(%0) : (tensor<4x6xf32>) -> tensor<4x6xf32>
+  return %1 : tensor<4x6xf32>
+}
+
+// CHECK-LABEL: @reciprocal_of_const_sparse
+// Sparse tensors are currently not supported
+func.func @reciprocal_of_const_sparse() -> tensor<32xbf16> {
+  // CHECK: tosa.const
+  // CHECK: tosa.reciprocal
+    %0 = "tosa.const"() { value = sparse<
+          [[0], [3], [11], [17], [20], [23], [25], [30], [31]],
+          [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]>
+          : tensor<32xbf16> } : () -> tensor<32xbf16>
+    %1 = "tosa.reciprocal"(%0) : (tensor<32xbf16>) -> tensor<32xbf16>
+    return %1 : tensor<32xbf16>
+}