// CHECK-NEXT:  %0 = "tf.IfrtCall"(%arg1, %arg0)
// CHECK-SAME:       {program_id = [[PROGRAM_ID:.*]] : i64, variable_arg_indices = []
// CHECK-SAME:       (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    %1 = "tf.Identity"(%arg1) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
// CHECK-NEXT:    %2 = "tf.IfrtCall"(%1, %arg0)
// CHECK-SAME:       {program_id = [[PROGRAM_ID]] : i64, variable_arg_indices = []

// CHECK-NEXT:    "tf.MatMul"(%arg0, [[TENSOR]]) : (tensor<1x3xf32>, tensor<3x1xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    "tf.IfrtCall"(%arg0, [[ARRAYKEY]]) <{program_id = 6515870160938153680 : i64, variable_arg_indices = [1 : i32]}> {__tpu_compile_metadata_text = "retvals { sharding { } }"} : (tensor<1x3xf32>, tensor<!tf_type.string>) -> tensor<1x1xf32>
// CHECK-NEXT:    return
//

  return
}

// -----
// CHECK-LABEL: func @test_ifrt_call
func.func @test_ifrt_call(%arg0: tensor<?xf32>, %arg1: tensor<?xf32>) {
  %result = "tf.IfrtCall"(%arg0, %arg1) <{program_id = 1234 : i64, variable_arg_indices = [0 : i32, 1 : i32], variable_names = ["a", "b"]}> : (tensor<?xf32>, tensor<?xf32>) -> (tensor<1x1xf32>)
  func.return
}

// -----

// CHECK-NEXT:   [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"([[HANDLE2]])
// CHECK-SAME:       used_by_host = false 
// CHECK-NEXT:   [[RES:%.*]] = "tf.IfrtCall"([[KEY]], %arg0) <{program_id = 6515870160938153680 : i64, variable_arg_indices = [0 : i32]}>
// CHECK-SAME:    : (tensor<!tf_type.string>, tensor<1x3xf32>) -> tensor<1x1xf32>
// CHECK-NEXT:    return [[RES]] : tensor<1x1xf32>
//
module {

    During lowering from a `tf_device.cluster_func` op to a `tf.IfrtCall` op,
    the region owned by the former will be outlined to a function with a
    `tfrt_ifrt_serving.program_id` attribute. After that, the runtime ensures
    that the outlined function is compiled into an executable and is available
    for lookup from `IfrtCall` TF ops.

mlir::LogicalResult IfrtCallOp::verify() {
  auto func = getOperation()->getParentOfType<mlir::func::FuncOp>();
  if (func != nullptr && func->hasAttr("tfrt_ifrt_serving.program_id")) {
    return emitOpError() << "cannot be nested inside an IFRT program";
  }

  for (mlir::Value arg : getArgs()) {
    if (mlir::isa<mlir::TF::ResourceType>(

   // CHECK-SAME: Y
   // CHECK: OpA
   // CHECK: OpB
   // CHECK: OpC
   "tf.OpA"() : () -> ()
   %status_x, %program_x = "tf._TPUCompileMlir"() { metadata = "X", mlir_module = "..." } : () -> (tensor<!tf_type.string>, tensor<!tf_type.string>)
   "tf.OpB"() : () -> ()
   %status_y, %program_y = "tf._TPUCompileMlir"() { metadata = "Y", mlir_module = "..." } : () -> (tensor<!tf_type.string>, tensor<!tf_type.string>)
   "tf.OpC"() : () -> ()

 */

package org.gradle.kotlin.dsl.execution

import org.gradle.internal.hash.HashCode
import org.gradle.kotlin.dsl.support.KotlinCompilerOptions
import java.lang.ref.WeakReference


class ProgramId(
    val templateId: String,
    val sourceHash: HashCode,
    parentClassLoader: ClassLoader,
    private val accessorsClassPathHash: HashCode? = null,
    private val classPathHash: HashCode? = null,

    interface Host {

        fun cachedClassFor(
            programId: ProgramId
        ): CompiledScript?

        fun cache(
            specializedProgram: CompiledScript,
            programId: ProgramId
        )

        fun cachedDirFor(
            scriptHost: KotlinScriptHost<*>,
            programId: ProgramId,
            compilationClassPath: ClassPath,
            accessorsClassPath: ClassPath,

            accessorsClassPath: ClassPath
        )

        fun compileSecondStageOf(
            program: StagedProgram,
            scriptHost: KotlinScriptHost<*>,
            programId: ProgramId,
            programKind: ProgramKind,
            programTarget: ProgramTarget,
            accessorsClassPath: ClassPath
        ): CompiledScript

        fun handleScriptException(
            exception: Throwable,