// RUN: tf-opt %s -unfold-large-splat-constant | FileCheck %s

// CHECK-LABEL: @unfold_large_constant_splat
func.func @unfold_large_constant_splat() -> (tensor<10x10xf32>, tensor<1000x1000xf32>) {
  %0 = arith.constant dense<0.00000e+00> : tensor<10x10xf32>
  %1 = arith.constant dense<1.00000e+00> : tensor<1000x1000xf32>
  func.return %0, %1 : tensor<10x10xf32>, tensor<1000x1000xf32>

  // CHECK-DAG: %cst = arith.constant dense<0.000000e+00> : tensor<10x10xf32>

	BGE    index2to16   // Remaining string 2<=len<=16
	MOVD   $0xff00, R21 // Set up mask for upcoming loop
	MTVSRD R21, V25     // Move mask to Vreg
	VSPLTH $3, V25, V31 // Splat mask
	VSPLTH $0, V0, V1   // Splat 1st two bytes of sep
	VSPLTB $2, V0, V8   // Splat 3rd byte of sep

	// Loop to process 3 byte separator.
	// string[0:16] is in V2
	// string[2:18] is in V3
	// sep[0:2] splatted in V1
	// sec[3] splatted in v8

    const GraphdefToMlirOptions& import_options, mlir::MLIRContext* context);

// Similar as the above function, but replaces all constant tensors
// with randomly generated splat values.
absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>>
GraphdefToSplattedMlirTranslateFunction(
    llvm::StringRef input, const std::vector<std::string>& input_arrays,
    const std::vector<std::string>& input_dtypes,

#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"

namespace mlir {
namespace odml {
namespace {

#define DEBUG_TYPE "unfold-splat-constant-pass"

#define GEN_PASS_DEF_UNFOLDSPLATCONSTANTPASS
#include "tensorflow/compiler/mlir/lite/stablehlo/transforms/passes.h.inc"

// Undo the MHLO::BroadcastInDimOp folding pattern on splat tensor.
// TODO(b/295966255): Remove this pass after moving MHLO folders to a separate

// The threshold of constant bits to be unfolded (1Mb). If there is a splat
// constant with size equal or greater to this threshold, then it will be
// unfolded back to a regular `tfl.fill` operation.
constexpr int64_t kConstantSizeThresholdInBits = 1e+6;

// Pass which will replace large splat constant tensors to `tfl.Fill` op to
// reduce the size of the generated flatbuffer model size.

// RUN: odml-to-stablehlo-opt %s -unfold-splat-constant-pass -cse -verify-diagnostics | FileCheck %s

// CHECK-LABEL: @unfold_splat_constant_float
func.func @unfold_splat_constant_float() -> tensor<1x750xf32> {
  %cst = mhlo.constant dense<7.680000e+02> : tensor<1x750xf32>
  func.return %cst : tensor<1x750xf32>

  // CHECK-DAG: %0 = mhlo.constant dense<7.680000e+02> : tensor<f32>

// Returns a Constant op with a single value.
absl::StatusOr<arith::ConstantOp> CreateConstOpWithSingleValue(
    PatternRewriter* rewriter, Location loc, ShapedType shaped_type, int value);

// Returns a Constant op with a splat vector value.
absl::StatusOr<arith::ConstantOp> CreateConstOpWithVectorValue(
    PatternRewriter* rewriter, Location loc, ShapedType shaped_type, int value);

}  // namespace TFL
}  // namespace mlir

	LXVW4X (R9)(CONSTBASE), VS34
	LXVW4X (R10)(CONSTBASE), VS35

	// splat values from V17, V18 into V4-V11
	VSPLTW $0, V17, V4
	VSPLTW $1, V17, V5
	VSPLTW $2, V17, V6
	VSPLTW $3, V17, V7
	VSPLTW $0, V18, V8
	VSPLTW $1, V18, V9
	VSPLTW $2, V18, V10
	VSPLTW $3, V18, V11

	// VOR
	VOR V26, V26, V12

	// splat values from V19 -> V13, V14, V15
	VSPLTW $1, V19, V13
	VSPLTW $2, V19, V14

                                               bool is_intermediate = false,
                                               bool get_storage = false);

// Gets a constant splat for the given value of type. Requires value to be of
// type static shaped RankedTensorType. `unique_index` is used to get the unique
// value for the attribute.

  let dependentDialects = ["mlir::stablehlo::StablehloDialect"];
}

def UnfoldSplatConstantPass : Pass<"unfold-splat-constant-pass", "ModuleOp"> {
  let summary = "Replaces a splat constant tensor with a BroadcastInDim op.";
  let constructor = "mlir::odml::CreateUnfoldSplatConstantPass()";
}

def CompositeLoweringPass : Pass<"composite-lowering", "ModuleOp"> {