Increase constant fold policy size threshold.

To accommodate some of the failing tests, increase the size threshold of constant fold policy from 128 KB to 256 KB.

PiperOrigin-RevId: 320853763
Change-Id: I68f905a0236cc0923c28bd9bf182f35beb72f36b

Author: Prakalp Srivastava

tensorflow/compiler/mlir/tensorflow/tests/constant-fold.mlir

@@ -443,7 +443,7 @@ func @DontRemoveTrivialMul(%arg0: tensor<1x6x8x1xf32>) -> tensor<1x6x8x1xf32> {
   // CHECK: return %[[RESULT]] : tensor<1x6x8x1xf32>
 }
 
-// Do not fold if total result size is large (>128 KB) and more than 2 times
+// Do not fold if total result size is large (>256 KB) and more than 2 times
 // the size of operands.
 
 // LINT.IfChange(folding-policy-test)

tensorflow/compiler/mlir/tensorflow/transforms/constant_fold.cc

@@ -40,7 +40,7 @@ namespace TF {
 // LINT.IfChange(folding-policy)
 static bool ShouldBeFolded(Operation* inst) {
   constexpr int kSizeFactor = 2;
-  constexpr int64_t kSizeThreshold = (1 << 20);  // 128 KB
+  constexpr int64_t kSizeThreshold = (1 << 21);  // 256 KB
   bool has_unknown_shape = false;
   auto get_size = [&](TypeRange types) {
     int64_t size = 0;