//go:build msan

#include "go_asm.h"
#include "textflag.h"

#define RARG0 R0
#define RARG1 R1
#define RARG2 R2
#define FARG R3

// func runtime·domsanread(addr unsafe.Pointer, sz uintptr)
// Called from msanread.
TEXT	runtime·domsanread(SB), NOSPLIT, $0-16
	MOVD	addr+0(FP), RARG0
	MOVD	sz+8(FP), RARG1
	// void __msan_read_go(void *addr, uintptr_t sz);
	MOVD	$__msan_read_go(SB), FARG

// CHECK-LABEL: dequantize
func.func @dequantize(%arg0: tensor<2x3x!tf_type.qint8>, %min_range: tensor<f32>, %max_range: tensor<f32>) -> tensor<2x3xf32> {
  // CHECK-DAG: %[[HALF_RANGE:.*]] = "tf.Const"() <{value = dense<1.280000e+02> : tensor<f32>}>
  // CHECK-DAG: %[[C255:.*]] = "tf.Const"() <{value = dense<2.550000e+02> : tensor<f32>}>
  // CHECK-DAG: %[[CAST:.*]] = "tf.Cast"(%arg0) <{Truncate = false}>

  %0 = "tfl.maximum"(%arg0, %cst) : (tensor<2x3xbf16>, tensor<2x3xbf16>) -> tensor<2x3xbf16>
  func.return %0 : tensor<2x3xbf16>

  // CHECK: %[[RESULT:.*]] = "tfl.relu"(%arg0)
  // CHECK: return %[[RESULT]]
}

// CHECK-LABEL: fuseScalarAddIntoConv2dBf16
func.func @fuseScalarAddIntoConv2dBf16(%arg0: tensor<256x32x32x3xbf16>, %arg1: tensor<16x3x3x3xbf16>) -> tensor<256x8x7x16xbf16> {
  %cst = arith.constant dense<1.5> : tensor<bf16>

// CHECK-LABEL: func @testPwStreamResults
func.func @testPwStreamResults(%arg0: tensor<f32>, %arg1: tensor<f32>) {
  "tf.PwStreamResults"(%arg0, %arg1) {names = ["foo", "bar"]} : (tensor<f32>, tensor<f32>) -> ()
  return
}

// -----
// CHECK-LABEL: func @test_ifrt_call
func.func @test_ifrt_call(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) {

  // CHECK-LABEL: @serving_default
  // CHECK-SAME: %[[ARG0:[^:[:space:]]+]]
  // CHECK-SAME: %[[ARG1:[^:[:space:]]+]]
  func.func private @serving_default(%arg0: tensor<1x4xf32>, %arg1: tensor<4x3xf32>) -> tensor<1x3xf32> {

// CHECK-LABEL: TfAdd
func.func @TfAdd(%arg0: tensor<4xf64>, %arg1: tensor<4xf64>) -> tensor<4xf64> {
  %0 = "tfl.custom"(%arg0, %arg1) {
    custom_code = "FlexAdd",
    custom_option = #tfl<const_bytes : "0x03416464001412034164641A001A002A070A015412023002320000021B171414042801">
  } : (tensor<4xf64>, tensor<4xf64>) -> tensor<4xf64>

// CHECK: "tf.Add"(%arg0, %arg1) {T = f64}  : (tensor<4xf64>, tensor<4xf64>) -> tensor<4xf64>

// CHECK-NEXT: z.b
// CHECK-NEXT: z.c

// -----

// CHECK-LABEL: @terminator
func.func @terminator(%arg0: f32) -> (f32) attributes {ignore_side_effects_for_testing} {
  func.call @terminator(%arg0) : (f32) -> (f32)
  "x.a"(%arg0) : (f32) -> ()
  "y.a"(%arg0) : (f32) -> ()
  "z.a"(%arg0) : (f32) -> ()
  func.return %arg0 : f32
}
// CHECK: x.a
// CHECK: y.a
// CHECK: z.a
// CHECK: return
// CHECK-NOT: call
// CHECK: }

// RUN: tf-quant-opt %s -split-input-file -quant-preprocess-op -quant-prepare-quantize-drq | FileCheck %s

module {
  func.func @matmul(%arg0: tensor<1x2x2x3xf32>) -> (tensor<*xf32>) {
    %cst_0 = "tf.Const"() {value = dense<0.000000e+00> : tensor<2x1024xf32>} : () -> tensor<2x1024xf32>

  return %0 : tensor<1x3xf32>
}

// CHECK-LABEL: @_func
func.func private @_func(%arg0: tensor<1x3xf32>) -> (tensor<1x3xf32>) {
  return %arg0 : tensor<1x3xf32>
}
}


// -----
// SPMD: one input and no return
//
// CHECK-LABEL: func.func @serving_default(%arg0: tensor<1x3xf32>) {
// CHECK-NEXT:  "tf.IfrtCall"(%arg0)
// CHECK-SAME:       {program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []}

// CHECK-NOT: tf._noinline
func.func @inner_stateful_pcall_func(%arg0: tensor<i32>) -> tensor<i32> {
  %0 = "tf.StatefulPartitionedCall"(%arg0) {_xla_compile_device_type = "CPU", config = "", config_proto = "", device = "/device:CPU:0", executor_type = "", f = @func} : (tensor<i32>) -> (tensor<i32>)
  func.return %0 : tensor<i32>
}

func.func @func(%arg0: tensor<i32>) -> tensor<i32> {
  func.return %arg0 : tensor<i32>