// `stablehlo.dot_general` is legalizable to `tfl.fully_connected`.
  const bool is_per_axis_quantizable =
      IsDotGeneralFullyConnected(dot_general_op).value();
  if (!is_per_axis_quantizable) return std::nullopt;
  return filter_rank - 1;
}

bool ContainsConvOrDot(StringRef str) {
  return str.contains("_conv") || str.contains("_dot_general");
}

#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project

namespace mlir {
namespace odml {
// Converts mhlo.dot_general to tfl.BatchMatMul. Reshape and Transpose ops will
// be inserted to convert to well-formed matrix multiply; i.e., mhlo.dot_general
// -> tfl.batch_matmul(mhlo.transpose(mhlo.reshape(operand)), ...).
// Note:
// 1) Reshape/transpose are inserted because tfl.BatchMatMul requires

See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// This file implements logic for legalizing mhlo.dot_general to
// tflite.batch_matmul.

#include <cstddef>
#include <cstdint>
#include <functional>
#include <numeric>

#include "llvm/ADT/APFloat.h"
#include "llvm/ADT/APInt.h"
#include "llvm/ADT/ArrayRef.h"

      // Add is there to prevent from dot_generals from being DCEed.
      %6 = stablehlo.add %1, %5 : tensor<1x3xf32>
      return %6 : tensor<1x3xf32>
    }

    // Callee of non-quantized op.
    func.func private @composite_dot_general_fn_1(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> {

def IsNull : Constraint<CPred<"!$0">>;

// This pattern optimizes:
//   conv/dot_general + a + b -> conv/dot_general + (a + b)
//   conv/dot_general - a - b -> conv/dot_general - (a + b)
//   conv/dot_general + a - b -> conv/dot_general + (a - b)
//   conv/dot_general - a + b -> conv/dot_general - (a - b)
foreach OpsTuple = [
    [CHLO_BroadcastAddOp, CHLO_BroadcastAddOp, CHLO_BroadcastAddOp],

                                       "unsupported dot_general layout");
  }

  // (B = batch dims, C = contracting dims, Y/Z = other dims)
  //
  // This pattern converts:
  //   %before = "mhlo.reshape"(%lhs : BxY1xC) : BxY2xC
  //   %dot = "mhlo.dot_general"(%before, %rhs : BxCxZ) : BxY2xZ
  //   %after = "mhlo.reshape"(%dot) : BxY1xZ
  // to:
  //   %dot : "mhlo.dot_general"(%lhs : BxY1xC, %rhs : BxCxZ) : BxY1xZ

// CHECK-LABEL: private @composite_dot_general_with_bias_dynamic_fn_1
// CHECK: %[[DOT_GENERAL:.*]] = stablehlo.dot_general %arg0, %arg1
// CHECK: %[[SHAPE_OF_0:.*]] = shape.shape_of %[[DOT_GENERAL]]
// CHECK: %[[DYNAMIC_BROADCAST_IN_DIM_0:.*]] = stablehlo.dynamic_broadcast_in_dim %arg2, %[[SHAPE_OF_0]]
// CHECK: %[[ADD:.*]] = stablehlo.add %[[DOT_GENERAL]], %[[DYNAMIC_BROADCAST_IN_DIM_0]]
// CHECK: return %[[ADD]] : tensor<?x10xf32>
// CHECK: }

  %0 = "mhlo.dot_general"(%arg0, %arg3) {
    dot_dimension_numbers = #mhlo.dot<
      lhs_contracting_dimensions = [2],
      rhs_contracting_dimensions = [0]
  >} : (tensor<1x1x512xf32>, tensor<512x13xf32>) -> tensor<1x1x13xf32>
  %1 = "mhlo.dot_general"(%arg1, %arg3) {
    dot_dimension_numbers = #mhlo.dot<
      lhs_contracting_dimensions = [2],

// DISABLE-ALL-DOT-GENERAL-LABEL: private @composite_dot_general_fn_1
// DISABLE-ALL-DOT-GENERAL-SAME: tf_quant.composite_function
// DISABLE-ALL-DOT-GENERAL: %[[DOT_GENERAL:.+]] = stablehlo.dot_general %arg0, %arg1
// DISABLE-ALL-DOT-GENERAL: return %[[DOT_GENERAL:.+]] : tensor<1x1x64xf32>
// DISABLE-ALL-DOT-GENERAL: }

// -----

        "@llvm-project//mlir:Support",
        "@llvm-project//mlir:TransformUtils",
        "@local_xla//xla/mlir_hlo",
    ],
)

cc_library(
    name = "dot_general",
    srcs = [
        "dot_general.cc",
    ],
    hdrs = ["dot_general.h"],
    deps = [
        "//tensorflow/compiler/mlir/lite:tensorflow_lite",
        "@llvm-project//llvm:Support",
        "@llvm-project//mlir:FuncDialect",