chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tools/BUILD

@@ -0,0 +1,58 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "translate_cl_options",
+    srcs = [
+        "tf_mlir_translate_cl.cc",
+    ],
+    hdrs = [
+        "tf_mlir_translate_cl.h",
+    ],
+    deps = [
+        "@llvm-project//llvm:Support",
+    ],
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "translate_registration",
+    srcs = [
+        "tf_mlir_translate_registration.cc",
+    ],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tensorflow/translate:mlir_roundtrip_flags",
+        "//tensorflow/compiler/mlir/tensorflow/translate/tools:file_tf_mlir_translate",
+        "//tensorflow/compiler/mlir/tf2xla/api/v2:tf_executor_to_graph",
+        "//tensorflow/compiler/mlir/tools:translate_cl_options",
+        "//tensorflow/compiler/tf2xla:xla_compiler",
+        "//tensorflow/compiler/tf2xla/kernels:xla_ops",
+        "//tensorflow/core:core_cpu_base",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:protos_all_cc",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:TranslateLib",
+        "@tsl//tsl/platform:protobuf",
+        "@xla//xla/client:client_library",
+        "@xla//xla/client:compile_only_client",
+        "@xla//xla/service/cpu:cpu_compiler",
+        "@xla//xla/service/cpu:cpu_transfer_manager",
+        "@xla//xla/stream_executor/host:host_platform",
+        "@xla//xla/stream_executor/host:host_platform_id",
+    ],
+    alwayslink = 1,
+)
+
+cc_library(
+    name = "safe_cast",
+    hdrs = ["safe_cast.h"],
+)

tensorflow/compiler/mlir/tools/kernel_gen/BUILD

@@ -0,0 +1,267 @@
+load(
+    "@local_config_rocm//rocm:build_defs.bzl",
+    "if_rocm_is_configured",
+)
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load(
+    "@xla//xla/stream_executor:build_defs.bzl",
+    "if_gpu_is_configured",
+)
+load(
+    "@xla//xla/tsl/platform/default:cuda_build_defs.bzl",
+    "if_cuda_is_configured",
+)
+load(
+    "//tensorflow:tensorflow.bzl",
+    "check_deps",
+    "tf_cc_binary",
+)
+load(
+    "//tensorflow/core/platform:build_config.bzl",
+    "tf_proto_library",
+)
+load(
+    "//tensorflow/core/platform:build_config_root.bzl",
+    "if_llvm_aarch32_available",
+    "if_llvm_aarch64_available",
+    "if_llvm_powerpc_available",
+    "if_llvm_system_z_available",
+    "if_llvm_x86_available",
+)
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [":friends"],
+    licenses = ["notice"],
+)
+
+package_group(
+    name = "friends",
+    packages = [
+        "//tensorflow/compiler/...",
+        "//tensorflow/core/kernels/mlir_generated/...",
+    ],
+)
+
+cc_library(
+    name = "kernel_creator",
+    srcs = ["kernel_creator.cc"],
+    hdrs = ["kernel_creator.h"],
+    copts = if_cuda_is_configured(["-DGOOGLE_CUDA=1"]) + if_rocm_is_configured(["-DTENSORFLOW_USE_ROCM=1"]),
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow:dump_mlir_util",
+        "//tensorflow/compiler/mlir/tools/kernel_gen/transforms:bufferize",
+        "//tensorflow/compiler/mlir/tools/kernel_gen/transforms:gpu_passes",  # fixdeps: keep
+        "//tensorflow/compiler/mlir/tools/kernel_gen/transforms:passes",
+        "//tensorflow/core:lib",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AffineToStandard",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:ArithTransforms",
+        "@llvm-project//mlir:BufferizationTransforms",
+        "@llvm-project//mlir:BuiltinToLLVMIRTranslation",
+        "@llvm-project//mlir:ComplexDialect",
+        "@llvm-project//mlir:ComplexToStandard",
+        "@llvm-project//mlir:ControlFlowDialect",
+        "@llvm-project//mlir:DataLayoutInterfaces",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:GPUDialect",
+        "@llvm-project//mlir:GPUToGPURuntimeTransforms",
+        "@llvm-project//mlir:GPUToLLVMIRTranslation",
+        "@llvm-project//mlir:GPUToNVVMTransforms",
+        "@llvm-project//mlir:GPUTransforms",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LLVMIRTransforms",
+        "@llvm-project//mlir:LLVMToLLVMIRTranslation",
+        "@llvm-project//mlir:LinalgTransforms",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:MemRefTransforms",
+        "@llvm-project//mlir:NVVMToLLVMIRTranslation",
+        "@llvm-project//mlir:NVVMTransforms",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:ROCDLToLLVMIRTranslation",
+        "@llvm-project//mlir:ReconcileUnrealizedCasts",
+        "@llvm-project//mlir:SCFToControlFlow",
+        "@llvm-project//mlir:SCFToGPU",
+        "@llvm-project//mlir:SCFTransforms",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:ShapeToStandard",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@stablehlo//:chlo_ops",
+        "@stablehlo//:linalg_passes",
+        "@stablehlo//:stablehlo_passes",
+        "@xla//xla/mlir_hlo",
+        "@xla//xla/mlir_hlo:all_passes",  # fixdeps: keep
+        "@xla//xla/mlir_hlo:mhlo_passes",
+        "@xla//xla/mlir_hlo:transforms_gpu_passes",
+        "@xla//xla/mlir_hlo:transforms_passes",
+    ],
+)
+
+check_deps(
+    name = "hlo_to_kernel_check_deps",
+    disallowed_deps = ["@xla//xla:literal"],
+    deps = [":hlo_to_kernel"],
+)
+
+tf_cc_binary(
+    name = "hlo_to_kernel",
+    srcs = ["hlo_to_kernel.cc"],
+    visibility = [
+        "//tensorflow/compiler/mlir/tools/kernel_gen/tests/hlo_to_kernel:__pkg__",
+        "//tensorflow/core/kernels/mlir_generated:__pkg__",
+    ],
+    deps = [
+        ":kernel_creator",
+        "//tensorflow/compiler/mlir:init_mlir",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/core:lib",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//llvm:Analysis",
+        "@llvm-project//llvm:CodeGen",
+        "@llvm-project//llvm:Core",
+        "@llvm-project//llvm:MC",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:Target",
+        "@llvm-project//llvm:TargetParser",
+        "@llvm-project//mlir:BufferizationInterfaces",
+        "@llvm-project//mlir:ExecutionEngineUtils",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LLVMToLLVMIRTranslation",
+        "@llvm-project//mlir:MemRefTransforms",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:ToLLVMIRTranslation",
+        "@xla//xla/service/llvm_ir:llvm_command_line_options",
+    ] + if_llvm_aarch32_available([
+        "@llvm-project//llvm:ARMCodeGen",  # fixdeps: keep
+    ]) + if_llvm_aarch64_available([
+        "@llvm-project//llvm:AArch64CodeGen",  # fixdeps: keep
+    ]) + if_llvm_powerpc_available([
+        "@llvm-project//llvm:PowerPCCodeGen",  # fixdeps: keep
+    ]) + if_llvm_system_z_available([
+        "@llvm-project//llvm:SystemZCodeGen",  # fixdeps: keep
+    ]) + if_llvm_x86_available([
+        "@llvm-project//llvm:X86CodeGen",  # fixdeps: keep
+        "@llvm-project//llvm:X86Disassembler",  # fixdeps: keep
+    ]),
+)
+
+tf_cc_binary(
+    name = "kernel-gen-opt",
+    srcs = ["tools/kernel-gen-opt/kernel-gen-opt.cc"],
+    visibility = ["//tensorflow/compiler/mlir/tools/kernel_gen/tests:__subpackages__"],
+    deps = [
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tools/kernel_gen/ir:tf_framework_ops",
+        "//tensorflow/compiler/mlir/tools/kernel_gen/transforms:gpu_passes",  # fixdeps: keep
+        "//tensorflow/compiler/mlir/tools/kernel_gen/transforms:passes",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AllPassesAndDialects",
+        "@llvm-project//mlir:MlirOptLib",
+        "@llvm-project//mlir:RegisterAllDialects",
+        "@llvm-project//mlir:RegisterAllPasses",
+        "@stablehlo//:register",
+        "@xla//xla/mlir_hlo:all_passes",
+        "@xla//xla/mlir_hlo:hlo_dialect_registration",
+    ],
+)
+
+exports_files(["tf_framework_c_interface.h"])
+
+cc_library(
+    name = "tf_framework_c_interface",
+    srcs = ["tf_framework_c_interface.cc"],
+    hdrs = ["tf_framework_c_interface.h"],
+    copts = if_cuda_is_configured(["-DGOOGLE_CUDA=1"]) + if_rocm_is_configured(["-DTENSORFLOW_USE_ROCM=1"]),
+    deps = [
+        ":compile_cache_item_proto_cc",
+        ":kernel_creator",
+        ":tf_gpu_runtime_wrappers",
+        ":tf_jit_cache",
+        "//tensorflow/compiler/mlir/tools/kernel_gen/ir:tf_framework_ops",
+        "//tensorflow/core:framework",
+        "//tensorflow/core:lib",
+        "//tensorflow/core/platform:refcount",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/strings",
+        "@com_google_absl//absl/strings:string_view",
+        "@llvm-project//llvm:JITLink",
+        "@llvm-project//llvm:OrcShared",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ExecutionEngine",
+        "@llvm-project//mlir:ExecutionEngineUtils",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MemRefTransforms",
+        "@llvm-project//mlir:Parser",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:mlir_c_runner_utils",
+        "@llvm-project//mlir:mlir_runner_utils",
+        "@xla//xla/stream_executor:stream",
+    ],
+)
+
+cc_library(
+    name = "tf_jit_cache",
+    srcs = ["tf_jit_cache.cc"],
+    hdrs = ["tf_jit_cache.h"],
+    deps = [
+        "//tensorflow/core:framework",
+        "//tensorflow/core:framework_lite",
+        "//tensorflow/core/framework:resource_base",
+        "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:ExecutionEngine",
+        "@tsl//tsl/platform:thread_annotations",
+    ],
+)
+
+cc_library(
+    name = "tf_gpu_runtime_wrappers",
+    srcs = if_gpu_is_configured([
+        "tf_gpu_runtime_wrappers.cc",
+    ]),
+    hdrs =
+        if_gpu_is_configured([
+            "tf_gpu_runtime_wrappers.h",
+        ]),
+    copts = if_cuda_is_configured(["-DGOOGLE_CUDA=1"]) + if_rocm_is_configured([
+        "-DTENSORFLOW_USE_ROCM=1",
+    ]),
+    deps = [
+        "//tensorflow/core:framework",
+        "//tensorflow/core/framework:resource_base",
+        "//tensorflow/core/platform:logging",
+        "//tensorflow/core/platform:mutex",
+        "//tensorflow/core/platform:refcount",
+        "//tensorflow/core/platform:status",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/strings",
+        "@llvm-project//mlir:mlir_runner_utils",
+        "@tsl//tsl/platform:hash",
+        "@tsl//tsl/platform:thread_annotations",
+        "@xla//xla/stream_executor:stream",
+    ] + if_cuda_is_configured([
+        "@local_config_cuda//cuda:cuda_headers",
+        "@xla//xla/stream_executor/cuda:stream_executor_cuda",
+    ]) + if_rocm_is_configured([
+        "@local_config_rocm//rocm:rocm_headers",
+        "@xla//xla/stream_executor/rocm:stream_executor_rocm",
+    ]),
+)
+
+tf_proto_library(
+    name = "compile_cache_item_proto",
+    srcs = ["compile_cache_item.proto"],
+)

tensorflow/compiler/mlir/tools/kernel_gen/compile_cache_item.proto

@@ -0,0 +1,25 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+
+syntax = "proto3";
+
+package mlir.kernel_gen;
+
+// Protocolbuffer representing a compilation input and output. This will be used
+// for caching JIT compiles of kernels.
+message CompilationCacheItem {
+  string original_module = 1;
+  string result_module = 2;
+}

tensorflow/compiler/mlir/tools/kernel_gen/hlo_to_kernel.cc

@@ -0,0 +1,250 @@
+// Copyright 2020 The TensorFlow Runtime Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//===- hlo_to_kernel.cc -----------------------------------------*- C++ -*-===//
+//
+// This file implements the entry point to compile a hlo op to a kernel.
+//
+//===----------------------------------------------------------------------===//
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/CodeGen/CommandFlags.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/IR/Module.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Support/CodeGen.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/SourceMgr.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/raw_ostream.h"
+#include "llvm/Target/TargetMachine.h"
+#include "llvm/TargetParser/Host.h"
+#include "llvm/TargetParser/Triple.h"
+#include "mlir/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.h"  // from @llvm-project
+#include "mlir/ExecutionEngine/OptUtils.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"  // from @llvm-project
+#include "mlir/Target/LLVMIR/Export.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/init_mlir.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.h"
+#include "xla/service/llvm_ir/llvm_command_line_options.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+namespace kernel_gen {
+namespace {
+
+static llvm::codegen::RegisterCodeGenFlags CGF;
+
+std::unique_ptr<llvm::TargetMachine> GetTargetMachine(
+    llvm::StringRef host_triple, llvm::Module* module) {
+  llvm::Triple triple(module->getTargetTriple());
+  if (triple.getTriple().empty()) {
+    if (!host_triple.empty()) {
+      triple = llvm::Triple(host_triple);
+    } else {
+      triple = llvm::Triple(llvm::sys::getDefaultTargetTriple());
+    }
+    module->setTargetTriple(llvm::Triple(triple.getTriple()));
+  }
+
+  std::string error;
+  const llvm::Target* target =
+      llvm::TargetRegistry::lookupTarget("", triple, error);
+  if (!target) {
+    return nullptr;
+  }
+
+  llvm::TargetOptions target_options =
+      llvm::codegen::InitTargetOptionsFromCodeGenFlags(llvm::Triple());
+  return std::unique_ptr<llvm::TargetMachine>(target->createTargetMachine(
+      triple, "generic", "", target_options, llvm::Reloc::Model::PIC_));
+}
+
+// Compiles the given MLIR module via LLVM into an executable binary format.
+absl::StatusOr<std::string> EmitToBinary(llvm::StringRef host_triple,
+                                         mlir::ModuleOp module) {
+  // Translate the module.
+  llvm::LLVMContext llvm_context;
+  mlir::registerLLVMDialectTranslation(*module->getContext());
+  std::unique_ptr<llvm::Module> llvm_module =
+      mlir::translateModuleToLLVMIR(module, llvm_context);
+
+  auto target_machine = GetTargetMachine(host_triple, llvm_module.get());
+  llvm_module->setDataLayout(target_machine->createDataLayout());
+
+  // Run LLVM's mid-level optimizer to clean up the IR.
+  if (mlir::makeOptimizingTransformer(
+          /*optLevel=*/2, /*sizeLevel=*/0,
+          target_machine.get())(llvm_module.get())) {
+    return absl::InternalError("Failed to run LLVM optimizer passess");
+  }
+
+  // Set up the output stream.
+  llvm::SmallString<8> outstr;
+  llvm::raw_svector_ostream ostream(outstr);
+  ostream.SetUnbuffered();
+
+  llvm::legacy::PassManager codegen_passes;
+  codegen_passes.add(new llvm::TargetLibraryInfoWrapperPass(
+      llvm::Triple(llvm_module->getTargetTriple())));
+
+  if (target_machine->addPassesToEmitFile(codegen_passes, ostream, nullptr,
+                                          llvm::CodeGenFileType::ObjectFile,
+                                          false)) {
+    return absl::InternalError("Failed add passes to emit file");
+  }
+  codegen_passes.run(*llvm_module);
+  return ostream.str().str();
+}
+
+absl::Status Run(llvm::StringRef input_file, llvm::StringRef output_file,
+                 llvm::StringRef host_triple,
+                 llvm::ArrayRef<std::string> architectures,
+                 llvm::ArrayRef<int64_t> tile_sizes,
+                 llvm::ArrayRef<int64_t> unroll_factors, bool print_ptx,
+                 bool print_llvmir, bool enable_ftz, bool index_64bit,
+                 bool jit_compile, bool jit_i64_indexed_for_large_tensors) {
+  // Read TF code.
+  std::string hlo_code;
+  TF_RETURN_IF_ERROR(
+      ReadFileToString(Env::Default(), input_file.str(), &hlo_code));
+
+  // Compile.
+  mlir::DialectRegistry registry;
+  mlir::memref::registerAllocationOpInterfaceExternalModels(registry);
+  mlir::MLIRContext context(registry);
+
+  llvm::SourceMgr source_mgr;
+  mlir::SourceMgrDiagnosticHandler source_mgr_handler(source_mgr, &context);
+
+  TF_ASSIGN_OR_RETURN(
+      mlir::OwningOpRef<mlir::ModuleOp> module,
+      GenerateKernelForHloCode(context, hlo_code, architectures, tile_sizes,
+                               unroll_factors, print_ptx, print_llvmir,
+                               enable_ftz, index_64bit, jit_compile,
+                               jit_i64_indexed_for_large_tensors,
+                               /*apply_cl_options=*/true));
+
+  // Get binary.
+  TF_ASSIGN_OR_RETURN(std::string binary, EmitToBinary(host_triple, *module));
+
+  // Write .a file.
+  TF_RETURN_IF_ERROR(
+      WriteStringToFile(Env::Default(), output_file.str(), binary));
+  return absl::OkStatus();
+}
+
+}  // namespace
+}  // namespace kernel_gen
+}  // namespace tensorflow
+
+int main(int argc, char** argv) {
+  llvm::cl::opt<std::string> input_file("input", llvm::cl::desc("input file"),
+                                        llvm::cl::value_desc("filename"),
+                                        llvm::cl::init("foo.mlir"));
+  llvm::cl::opt<std::string> output_file(
+      "output", llvm::cl::desc("output file"), llvm::cl::value_desc("filename"),
+      llvm::cl::init("foo.bin"));
+  llvm::cl::opt<bool> index_64bit("index_64bit",
+                                  llvm::cl::desc("enable 64 bit indexing"),
+                                  llvm::cl::init(false));
+  llvm::cl::opt<bool> print_ptx(
+      "print-ptx",
+      llvm::cl::desc("print generated PTX code per target architecture."),
+      llvm::cl::init(false));
+  llvm::cl::opt<bool> print_llvmir(
+      "print-llvmir", llvm::cl::desc("print llvm ir during lowering to ptx."),
+      llvm::cl::init(false));
+  llvm::cl::opt<bool> enable_ftz(
+      "enable_ftz",
+      llvm::cl::desc(
+          "enable the denormal flush to zero mode when generating code."),
+      llvm::cl::init(false));
+  llvm::cl::opt<bool> jit_compile(
+      "jit", llvm::cl::desc("Generate only a JIT compiler invocation."),
+      llvm::cl::init(false));
+  llvm::cl::opt<std::string> host_triple(
+      "host-triple", llvm::cl::desc("Override host triple for module"));
+  llvm::cl::list<std::string> architectures(
+      "arch", llvm::cl::desc("target architectures (e.g. sm_70 or compute_75)"),
+      llvm::cl::ZeroOrMore, llvm::cl::CommaSeparated);
+  llvm::cl::list<int64_t> tile_sizes(
+      "tile_sizes", llvm::cl::desc("tile sizes to use"), llvm::cl::ZeroOrMore,
+      llvm::cl::CommaSeparated);
+  llvm::cl::list<int64_t> unroll_factors(
+      "unroll_factors",
+      llvm::cl::desc("factors to unroll by, separated by commas"),
+      llvm::cl::ZeroOrMore, llvm::cl::CommaSeparated);
+  llvm::cl::opt<bool> jit_i64_indexed_for_large_tensors(
+      "jit_i64_indexed_for_large_tensors",
+      llvm::cl::desc(
+          "Enable JIT compilation of i64-indexed kernels for large inputs."),
+      llvm::cl::init(false));
+
+  tensorflow::InitMlir y(&argc, &argv);
+
+#ifdef TF_LLVM_X86_AVAILABLE
+  LLVMInitializeX86Target();
+  LLVMInitializeX86TargetInfo();
+  LLVMInitializeX86TargetMC();
+  LLVMInitializeX86AsmPrinter();
+#endif
+
+#ifdef TF_LLVM_AARCH64_AVAILABLE
+  LLVMInitializeAArch64Target();
+  LLVMInitializeAArch64TargetInfo();
+  LLVMInitializeAArch64TargetMC();
+  LLVMInitializeAArch64AsmPrinter();
+#endif
+
+  mlir::registerPassManagerCLOptions();
+  mlir::registerMLIRContextCLOptions();
+
+  // Forward cli options to XLA, as it will reset llvm options internally
+  // during the first invocation.
+  auto& xla_llvm_global_options =
+      xla::llvm_ir::LLVMCommandLineOptionsLock::GetGlobalOptions();
+  xla_llvm_global_options.insert(xla_llvm_global_options.end(), argv + 1,
+                                 argv + argc);
+  llvm::cl::ParseCommandLineOptions(argc, argv, "TF op kernel generator\n");
+
+  auto status = tensorflow::kernel_gen::Run(
+      input_file, output_file, host_triple, architectures, tile_sizes,
+      unroll_factors, print_ptx, print_llvmir, enable_ftz, index_64bit,
+      jit_compile, jit_i64_indexed_for_large_tensors);
+  if (!status.ok()) {
+    LOG(ERROR) << status;
+    return 1;
+  }
+  return 0;
+}

tensorflow/compiler/mlir/tools/kernel_gen/ir/BUILD

@@ -0,0 +1,81 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//tensorflow/compiler/mlir/tools/kernel_gen:friends",
+        # Allow visibility from the mlir language server.
+        "//learning/brain/mlir/mlir_lsp_server:__pkg__",
+    ],
+    licenses = ["notice"],
+)
+
+td_library(
+    name = "td_files",
+    srcs = [
+        "tf_framework_ops.td",
+        "tf_status.td",
+    ],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "@llvm-project//mlir:AllocationOpInterfaceTdFiles",
+        "@llvm-project//mlir:ControlFlowInterfacesTdFiles",
+        "@llvm-project//mlir:OpBaseTdFiles",
+        "@llvm-project//mlir:SideEffectInterfacesTdFiles",
+    ],
+)
+
+gentbl_cc_library(
+    name = "tf_framework_ops_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {
+        "tf_framework_ops.h.inc": ["-gen-op-decls"],
+        "tf_framework_ops.cc.inc": ["-gen-op-defs"],
+        "tf_framework_dialect.h.inc": ["-gen-dialect-decls"],
+        "tf_framework_dialect.cc.inc": ["-gen-dialect-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tf_framework_ops.td",
+    deps = [":td_files"],
+)
+
+gentbl_cc_library(
+    name = "tf_status_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {
+        "tf_status.h.inc": ["-gen-enum-decls"],
+        "tf_status.cc.inc": ["-gen-enum-defs"],
+    },
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "tf_status.td",
+    deps = [":td_files"],
+)
+
+cc_library(
+    name = "tf_framework_ops",
+    srcs = [
+        "tf_framework_ops.cc",
+        "tf_framework_ops.cc.inc",
+        "tf_framework_ops.h.inc",
+    ],
+    hdrs = ["tf_framework_ops.h"],
+    deps = [
+        ":tf_framework_ops_inc_gen",
+        ":tf_status_inc_gen",
+        "//tensorflow/core/protobuf:error_codes_proto_impl_cc",
+        "@com_google_absl//absl/status",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:AllocationOpInterface",
+        "@llvm-project//mlir:BufferizationDialect",
+        "@llvm-project//mlir:BytecodeOpInterface",
+        "@llvm-project//mlir:ControlFlowInterfaces",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+        "@xla//xla/tsl/protobuf:error_codes_proto_impl_cc",
+    ],
+)

tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.cc

@@ -0,0 +1,174 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file defines the operations used in the tf_framework dialect.
+
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
+
+#include <optional>
+
+#include "absl/status/status.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
+#include "mlir/IR/Builders.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_status.cc.inc"
+
+// Generated dialect definitions.
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_dialect.cc.inc"
+#include "xla/tsl/protobuf/error_codes.pb.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace tf_framework {
+
+void TFFrameworkDialect::initialize() {
+  addOperations<
+#define GET_OP_LIST
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.cc.inc"
+      >();
+  addTypes<JITCallableType, OpKernelContextType>();
+}
+
+/// Parse a type registered to this dialect.
+Type TFFrameworkDialect::parseType(DialectAsmParser &parser) const {
+  StringRef keyword;
+  if (parser.parseKeyword(&keyword)) return Type();
+
+  if (keyword == "op_kernel_context") {
+    return OpKernelContextType::get(getContext());
+  }
+  if (keyword == "jit_callable") {
+    return JITCallableType::get(getContext());
+  }
+
+  parser.emitError(parser.getNameLoc(), "unknown TF Framework type: ")
+      << keyword;
+  return Type();
+}
+
+/// Print a type registered to this dialect.
+void TFFrameworkDialect::printType(Type type, DialectAsmPrinter &os) const {
+  if (mlir::isa<OpKernelContextType>(type)) {
+    os << "op_kernel_context";
+    return;
+  }
+  if (mlir::isa<JITCallableType>(type)) {
+    os << "jit_callable";
+    return;
+  }
+  llvm_unreachable("unexpected TF Framework type kind");
+}
+
+//===----------------------------------------------------------------------===//
+// TFAllocOp
+//===----------------------------------------------------------------------===//
+LogicalResult TFAllocOp::verify() {
+  TFAllocOp op = *this;
+  // Check that the total number of operands matches the number of dynamic
+  // dimensions specified in the memref type.
+  unsigned result_dyn_dims = op.getType().getNumDynamicDims();
+  unsigned dyn_sizes_count = op.getDynSizes().size();
+  if (dyn_sizes_count != result_dyn_dims)
+    return op.emitOpError()
+           << "`dyn_sizes` count " << dyn_sizes_count
+           << " does not match dynamic dimensions count in the result type"
+           << op.getType();
+  return success();
+}
+
+std::optional<Operation *> TFAllocOp::buildDealloc(OpBuilder &builder,
+                                                   Value alloc) {
+  auto funcop = alloc.getParentRegion()->getParentOfType<func::FuncOp>();
+  return TFDeallocOp::create(builder, alloc.getLoc(), funcop.getArgument(0),
+                             alloc)
+      .getOperation();
+}
+
+std::optional<Value> TFAllocOp::buildClone(OpBuilder &builder, Value alloc) {
+  // TODO(herhut): We should have our own clone op if one of these survives.
+  return mlir::bufferization::CloneOp::create(builder, alloc.getLoc(), alloc)
+      .getResult();
+}
+
+//===----------------------------------------------------------------------===//
+// JITExecuteOp
+//===----------------------------------------------------------------------===//
+
+std::optional<Operation *> JITExecuteOp::buildDealloc(OpBuilder &builder,
+                                                      Value alloc) {
+  auto funcop = alloc.getParentRegion()->getParentOfType<func::FuncOp>();
+  return TFDeallocOp::create(builder, alloc.getLoc(), funcop.getArgument(0),
+                             alloc)
+      .getOperation();
+}
+
+std::optional<Value> JITExecuteOp::buildClone(OpBuilder &builder, Value alloc) {
+  // TODO(herhut): We should have our own clone op if one of these survives.
+  return mlir::bufferization::CloneOp::create(builder, alloc.getLoc(), alloc)
+      .getResult();
+}
+
+absl::StatusCode ConvertAttrToEnumValue(ErrorCode error_code) {
+  using ::tensorflow::error::Code;
+  switch (error_code) {
+    case ErrorCode::OK:
+      return absl::StatusCode::kOk;
+    case ErrorCode::CANCELLED:
+      return absl::StatusCode::kCancelled;
+    case ErrorCode::UNKNOWN:
+      return absl::StatusCode::kUnknown;
+    case ErrorCode::INVALID_ARGUMENT:
+      return absl::StatusCode::kInvalidArgument;
+    case ErrorCode::DEADLINE_EXCEEDED:
+      return absl::StatusCode::kDeadlineExceeded;
+    case ErrorCode::NOT_FOUND:
+      return absl::StatusCode::kNotFound;
+    case ErrorCode::ALREADY_EXISTS:
+      return absl::StatusCode::kAlreadyExists;
+    case ErrorCode::PERMISSION_DENIED:
+      return absl::StatusCode::kPermissionDenied;
+    case ErrorCode::UNAUTHENTICATED:
+      return absl::StatusCode::kUnauthenticated;
+    case ErrorCode::RESOURCE_EXHAUSTED:
+      return absl::StatusCode::kResourceExhausted;
+    case ErrorCode::FAILED_PRECONDITION:
+      return absl::StatusCode::kFailedPrecondition;
+    case ErrorCode::ABORTED:
+      return absl::StatusCode::kAborted;
+    case ErrorCode::OUT_OF_RANGE:
+      return absl::StatusCode::kOutOfRange;
+    case ErrorCode::UNIMPLEMENTED:
+      return absl::StatusCode::kUnimplemented;
+    case ErrorCode::INTERNAL:
+      return absl::StatusCode::kInternal;
+    case ErrorCode::UNAVAILABLE:
+      return absl::StatusCode::kUnavailable;
+    case ErrorCode::DATA_LOSS:
+      return absl::StatusCode::kDataLoss;
+  }
+}
+
+}  // namespace tf_framework
+}  // namespace kernel_gen
+}  // namespace mlir
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.cc.inc"

tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h

@@ -0,0 +1,69 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file defines the operations used in the TFFramework dialect.
+//
+#ifndef TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_IR_TF_FRAMEWORK_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_IR_TF_FRAMEWORK_OPS_H_
+
+#include "absl/status/status.h"
+#include "mlir/Bytecode/BytecodeOpInterface.h"  // from @llvm-project
+#include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/OpImplementation.h"  // from @llvm-project
+#include "mlir/IR/TypeSupport.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Interfaces/ControlFlowInterfaces.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_status.h.inc"
+#include "tensorflow/core/protobuf/error_codes.pb.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace tf_framework {
+
+/// OpKernelContextType corresponds to C++ class OpKernelContext defined in
+/// tensorflow/core/framework/op_kernel.h
+class OpKernelContextType
+    : public Type::TypeBase<OpKernelContextType, Type, TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr StringLiteral name =
+      "kernel_gen.tf_framework.op_kernel_context";
+};
+
+class JITCallableType
+    : public Type::TypeBase<JITCallableType, Type, TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr StringLiteral name = "kernel_gen.tf_framework.jit_callable";
+};
+
+absl::StatusCode ConvertAttrToEnumValue(ErrorCode error_code);
+
+}  // namespace tf_framework
+}  // namespace kernel_gen
+}  // namespace mlir
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_dialect.h.inc"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_IR_TF_FRAMEWORK_OPS_H_

tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.td

@@ -0,0 +1,335 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This is the operation definition file for TF Framework ops.
+
+#ifndef TF_FRAMEWORK_OPS
+#define TF_FRAMEWORK_OPS
+
+include "mlir/IR/OpBase.td"
+include "mlir/Dialect/Bufferization/IR/AllocationOpInterface.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_status.td"
+
+def TFFramework_Dialect : Dialect {
+  let name = "tf_framework";
+
+  let summary = "Types and operations for tf_framework dialect";
+  let description = [{
+    This dialect contains operations and types for that correspond to
+    TensorFlow C++ Framework.
+  }];
+  let cppNamespace = "::mlir::kernel_gen::tf_framework";
+
+  let extraClassDeclaration = [{
+    static constexpr StringRef kTFEntryAttrName = "tf_entry";
+    static constexpr size_t kAllocationAlignment = 16;
+  }];
+
+  let useDefaultTypePrinterParser = 1;
+}
+
+def TFFramework_OpKernelContextType : DialectType<TFFramework_Dialect,
+    CPred<"llvm::isa<::mlir::kernel_gen::tf_framework::OpKernelContextType>($_self)">,
+          "op_kernel_construction">,
+    BuildableType<"$_builder.getType<::mlir::kernel_gen::tf_framework::OpKernelContextType>()"> {
+  let description = [{
+    OpKernelContextType corresponds to C++ class OpKernelContext defined in
+    tensorflow/core/framework/op_kernel.h
+  }];
+}
+
+def TFFramework_JITCallableType : DialectType<TFFramework_Dialect, CPred<
+    "llvm::isa<::mlir::kernel_gen::tf_framework::JITCallableType>($_self)">>,
+    BuildableType<"$_builder.getType<::mlir::kernel_gen::tf_framework::JITCallableType>()"> {
+  let description = [{
+    A `callable` represents the result of JIT compilation. Conceptually, it
+    corresponds to a function pointer.
+  }];
+}
+
+// Base class for TF Framework dialect ops.
+class TFFramework_Op<string mnemonic, list<Trait> traits = []> :
+    Op<TFFramework_Dialect, mnemonic, traits>;
+
+//===----------------------------------------------------------------------===//
+// TFAllocOp
+//===----------------------------------------------------------------------===//
+def TFFramework_TFAllocOp : TFFramework_Op<"alloc", [
+    MemoryEffects<[MemAlloc<DefaultResource>]>,
+    DeclareOpInterfaceMethods<AllocationOpInterface,
+      ["buildDealloc", "buildClone"]>
+  ]> {
+  let summary = "allocation of tensors that uses TF Framework";
+  let description = [{
+    Allocation of tensors during kernel execution in the Compute method.
+
+    This should be used to allocate any temporary or output memref. If
+    `output_index` and `input_indices` are given, attempts to forward one of
+    the input tensors to the output by calling `OpKernelContext::forward_input`.
+
+    If the attributes are missing or the forwarding fails, calls
+    `Allocator::AllocateRaw` in tensorflow/core/framework/allocator.h.
+  }];
+
+  let arguments = (ins
+    TFFramework_OpKernelContextType:$ctx,
+    Variadic<Index>:$dyn_sizes,
+    OptionalAttr<I32ArrayAttr>:$input_indices,
+    OptionalAttr<I32Attr>:$output_index
+  );
+  let results = (outs Res<AnyMemRef, "", [MemAlloc<DefaultResource>]>:$result);
+
+  let builders = [
+    OpBuilder<(ins "MemRefType":$memref_type, "Value":$ctx),
+    [{
+      $_state.addOperands(ctx);
+      $_state.types.push_back(memref_type);
+    }]>,
+    OpBuilder<(ins "MemRefType":$memref_type, "Value":$ctx,
+      "ValueRange":$dyn_sizes),
+    [{
+      build($_builder, $_state, memref_type, ctx);
+      $_state.addOperands(dyn_sizes);
+    }]>];
+
+  let extraClassDeclaration = [{
+    MemRefType getType() { return llvm::cast<MemRefType>(getResult().getType()); }
+    static constexpr StringRef kReuseOutputAttrName = "reuse_output";
+    static constexpr StringRef kReuseInputCandidatesAttrName =
+        "reuse_input_candidates";
+  }];
+  let assemblyFormat = [{
+    `(` $ctx (`,` $dyn_sizes^ )? `)` attr-dict `:` type($result)
+  }];
+
+  let hasVerifier = 1;
+}
+
+//===----------------------------------------------------------------------===//
+// TFAssertOp
+//===----------------------------------------------------------------------===//
+def TFFramework_TFAssertOp : TFFramework_Op<"assert"> {
+  let summary = "Assert operation with message attribute and error code";
+  let description = [{
+    Assert operation that propagates the error message to TF Framework.
+  }];
+
+  let arguments = (ins
+    TFFramework_OpKernelContextType:$ctx,
+    I1:$arg,
+    TFFramework_ErrorCodeAttr:$error_code,
+    StrAttr:$msg
+  );
+
+  let assemblyFormat = "$ctx `,` $arg `,` $error_code `,` $msg attr-dict";
+}
+
+//===----------------------------------------------------------------------===//
+// TFDeallocOp
+//===----------------------------------------------------------------------===//
+def TFFramework_TFDeallocOp : TFFramework_Op<"dealloc",
+    [MemoryEffects<[MemFree]>]> {
+  let summary = "deallocation of tensors that uses TF Framework";
+  let description = [{
+    Deallocation of tensors during kernel execution in the Compute method.
+
+    This should be used to deallocate any temporary memref that was allocated
+    with `tf_framework.alloc`.
+    Corresponds to `Allocator::DeallocateRaw` in
+    tensorflow/core/framework/allocator.h.
+  }];
+
+  let arguments = (ins TFFramework_OpKernelContextType:$ctx,
+                   Arg<AnyRankedOrUnrankedMemRef, "", [MemFree]>:$memref);
+  let assemblyFormat = "`(` $ctx `,` $memref `)` attr-dict `:` type($memref)";
+}
+
+//===----------------------------------------------------------------------===//
+// IsValidMemRefOp
+//===----------------------------------------------------------------------===//
+def TFFramework_IsValidMemRefOp : TFFramework_Op<"is_valid_memref", []> {
+  let summary = "Op to check if the memref is valid.";
+  let description = [{
+    The op checks if the allocation was successful so that the underlying ptr
+    in the descriptor is not equal to NULL. Also, we check if the number of
+    elements is not zero to support empty shapes inputs correctly.
+
+    is_valid_memref(memref) = ptr(memref) != NULL || num_elements(memref) == 0
+  }];
+
+  let arguments = (ins AnyMemRef:$arg);
+  let results = (outs I1:$result);
+  let assemblyFormat = [{
+    `(` $arg `)` attr-dict `:` type($arg) `->` type($result)
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// NullContextOp
+//===----------------------------------------------------------------------===//
+def TFFramework_NullContextOp : TFFramework_Op<"null_context",
+    [Pure]> {
+  let summary = "Creates a fake TF context that will be lowered to nullptr";
+  let description = [{Needed for testing}];
+  let results = (outs TFFramework_OpKernelContextType:$result);
+  let assemblyFormat = "attr-dict `:` type($result)";
+}
+
+//===----------------------------------------------------------------------===//
+// NullMemRefOp
+//===----------------------------------------------------------------------===//
+def TFFramework_NullMemRefOp : TFFramework_Op<"null_memref", []> {
+  let summary = "Op to construct unranked memref with 0-rank";
+  let description = [{
+    The op is needed to construct a throw-away result after error reporting
+    happened. It constructs a ranked memref descriptor with 0-rank and
+    {NULL, NULL, 0} underlying ranked descriptor of type memref<elem_type>.
+    In ranked case it constructs a ranked memref descriptor depending on the
+    result type, but sets both allocated and aligned pointers to NULL.
+  }];
+
+  let results = (outs AnyRankedOrUnrankedMemRef:$result);
+  let assemblyFormat = "attr-dict `:` type($result)";
+}
+
+//===----------------------------------------------------------------------===//
+// ReportErrorOp
+//===----------------------------------------------------------------------===//
+def TFFramework_ReportErrorOp : TFFramework_Op<"report_error", []> {
+  let summary = "Operation that propagates error message to TF Framework";
+  let description = [{
+    Error reporting operation that corresponds to
+    `OpKernelContext::CtxFailureWithWarning`.
+  }];
+
+  let arguments = (ins
+    TFFramework_OpKernelContextType:$ctx,
+    TFFramework_ErrorCodeAttr:$error_code,
+    StrAttr:$msg
+  );
+
+  let assemblyFormat = "$ctx `,` $error_code `,` $msg attr-dict";
+}
+
+//===----------------------------------------------------------------------===//
+// JITCompileOp
+//===----------------------------------------------------------------------===//
+
+def TFFramework_JITCompileOp : TFFramework_Op<"jit_compile", [SingleBlock]> {
+  let summary = "Invokes JIT compilation through the TF framework";
+  let description = [{
+    The op takes an optional TF context, so that it can be added at a later
+    stage in the compilation pipeline. The op's body corresponds to a function
+    body and the corresponding `tf_framework.jit_compile_yield` represents the
+    result value(s).
+  }];
+
+  let arguments = (ins
+    Optional<TFFramework_OpKernelContextType>:$ctx
+  );
+  let results = (outs TFFramework_JITCallableType:$result);
+  let regions = (region SizedRegion<1>:$body);
+
+  let assemblyFormat = "$ctx $body attr-dict";
+}
+
+//===----------------------------------------------------------------------===//
+// JITCompileYieldOp
+//===----------------------------------------------------------------------===//
+
+def TFFramework_JITCompileYieldOp : TFFramework_Op<"jit_compile_yield", [
+    Pure, ReturnLike, Terminator, HasParent<"JITCompileOp">]> {
+  let summary = "Yields the results in a `tf_framework.jit_compile` op's body";
+  let description = [{
+    See `tf_framework.jit_compile`.
+  }];
+
+  let arguments = (ins
+    AnyTypeOf<[AnyRankedOrUnrankedMemRef, AnyTensor]>:$result
+  );
+
+  let assemblyFormat = "$result attr-dict `:` type($result)";
+}
+
+//===----------------------------------------------------------------------===//
+// JITCompileFromStrOp
+//===----------------------------------------------------------------------===//
+
+def TFFramework_JITCompileFromStrOp : TFFramework_Op<"jit_compile_from_str",
+    []> {
+  let summary = "Invokes JIT compilation through the TF framework";
+  let description = [{
+    This operation is similar to `tf_framework.jit_compile`. Instead of a body,
+    the to-be-compiled function is represented as a string attribute. The string
+    shall be the serialized form of a module with a single function named
+    `main`.
+  }];
+
+  let arguments = (ins
+    Optional<TFFramework_OpKernelContextType>:$ctx,
+    StrAttr:$code,
+    I64ArrayAttr:$tileSizes,
+    I64ArrayAttr:$unrollFactors,
+    BoolAttr:$enableFtz,
+    BoolAttr:$index64Bit,
+    BoolAttr:$cpuCodegen
+  );
+  let results = (outs TFFramework_JITCallableType:$result);
+
+  let extraClassDeclaration = [{
+    static constexpr StringRef kJITEntryFunctionName = "main";
+  }];
+
+  let assemblyFormat = "($ctx^ `,`)? $code attr-dict";
+}
+
+//===----------------------------------------------------------------------===//
+// JITExecuteOp
+//===----------------------------------------------------------------------===//
+
+def TFFramework_JITExecuteOp : TFFramework_Op<"jit_execute", [
+    AttrSizedOperandSegments,
+    MemoryEffects<[MemAlloc<DefaultResource>]>,
+    DeclareOpInterfaceMethods<AllocationOpInterface,
+        ["buildDealloc", "buildClone"]>]> {
+  let summary = "Executes a JIT-compiled function through the TF framework";
+  let description = [{
+    The op takes an optional TF context, so that it can be added at a later
+    stage in the compilation pipeline. The callable must be a JIT-compiled
+    function that is the result of either `tf_framework.jit_compile` or
+    `tf_framework.jit_compile_from_str`. The remaining operands must be tensor
+    or memref arguments and will be forwarded to the callable function. The
+    result types must match those of the callable function. Otherwise, the
+    execution behavior is undefined.
+  }];
+
+  let arguments = (ins
+    Optional<TFFramework_OpKernelContextType>:$ctx,
+    TFFramework_JITCallableType:$callable,
+    Variadic<AnyTypeOf<[AnyRankedOrUnrankedMemRef, AnyTensor]>>:$inputs
+  );
+  let results = (outs
+    Res<AnyTypeOf<[AnyRankedOrUnrankedMemRef, AnyTensor]> , "",
+        [MemAlloc<DefaultResource>]>:$result);
+
+  let assemblyFormat = [{
+    (`ctx` `(` $ctx^ `)`)? $callable `(` $inputs `)` attr-dict
+        `:` type($inputs) `->` type($result)
+  }];
+}
+
+#endif // TF_FRAMEWORK_OPS

tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_status.td

@@ -0,0 +1,81 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This is the enum definition file for TF Status attribute.
+
+#ifndef TF_FRAMEWORK_STATUS
+#define TF_FRAMEWORK_STATUS
+
+include "mlir/IR/EnumAttr.td"
+include "mlir/IR/OpBase.td"
+
+def TFFramework_ERROR_CODE_OK
+    : I32EnumAttrCase<"OK", 0>;
+def TFFramework_ERROR_CODE_CANCELLED
+    : I32EnumAttrCase<"CANCELLED", 1>;
+def TFFramework_ERROR_CODE_UNKNOWN
+    : I32EnumAttrCase<"UNKNOWN", 2>;
+def TFFramework_ERROR_CODE_INVALID_ARGUMENT
+    : I32EnumAttrCase<"INVALID_ARGUMENT", 3>;
+def TFFramework_ERROR_CODE_DEADLINE_EXCEEDED
+    : I32EnumAttrCase<"DEADLINE_EXCEEDED", 4>;
+def TFFramework_ERROR_CODE_NOT_FOUND
+    : I32EnumAttrCase<"NOT_FOUND", 5>;
+def TFFramework_ERROR_CODE_ALREADY_EXISTS
+    : I32EnumAttrCase<"ALREADY_EXISTS", 6>;
+def TFFramework_ERROR_CODE_PERMISSION_DENIED
+    : I32EnumAttrCase<"PERMISSION_DENIED", 7>;
+def TFFramework_ERROR_CODE_UNAUTHENTICATED
+    : I32EnumAttrCase<"UNAUTHENTICATED", 16>;
+def TFFramework_ERROR_CODE_RESOURCE_EXHAUSTED
+    : I32EnumAttrCase<"RESOURCE_EXHAUSTED", 8>;
+def TFFramework_ERROR_CODE_FAILED_PRECONDITION
+    : I32EnumAttrCase<"FAILED_PRECONDITION", 9>;
+def TFFramework_ERROR_CODE_ABORTED
+    : I32EnumAttrCase<"ABORTED", 10>;
+def TFFramework_ERROR_CODE_OUT_OF_RANGE
+    : I32EnumAttrCase<"OUT_OF_RANGE", 11>;
+def TFFramework_ERROR_CODE_UNIMPLEMENTED
+    : I32EnumAttrCase<"UNIMPLEMENTED", 12>;
+def TFFramework_ERROR_CODE_INTERNAL
+    : I32EnumAttrCase<"INTERNAL", 13>;
+def TFFramework_ERROR_CODE_UNAVAILABLE
+    : I32EnumAttrCase<"UNAVAILABLE", 14>;
+def TFFramework_ERROR_CODE_DATA_LOSS
+    : I32EnumAttrCase<"DATA_LOSS", 15>;
+
+def TFFramework_ErrorCodeAttr : I32EnumAttr<"ErrorCode", "error code", [
+  TFFramework_ERROR_CODE_OK,
+  TFFramework_ERROR_CODE_CANCELLED,
+  TFFramework_ERROR_CODE_UNKNOWN,
+  TFFramework_ERROR_CODE_INVALID_ARGUMENT,
+  TFFramework_ERROR_CODE_DEADLINE_EXCEEDED,
+  TFFramework_ERROR_CODE_NOT_FOUND,
+  TFFramework_ERROR_CODE_ALREADY_EXISTS,
+  TFFramework_ERROR_CODE_PERMISSION_DENIED,
+  TFFramework_ERROR_CODE_UNAUTHENTICATED,
+  TFFramework_ERROR_CODE_RESOURCE_EXHAUSTED,
+  TFFramework_ERROR_CODE_FAILED_PRECONDITION,
+  TFFramework_ERROR_CODE_ABORTED,
+  TFFramework_ERROR_CODE_OUT_OF_RANGE,
+  TFFramework_ERROR_CODE_UNIMPLEMENTED,
+  TFFramework_ERROR_CODE_INTERNAL,
+  TFFramework_ERROR_CODE_UNAVAILABLE,
+  TFFramework_ERROR_CODE_DATA_LOSS
+]> {
+  let cppNamespace = "::mlir::kernel_gen::tf_framework";
+}
+
+#endif // TF_FRAMEWORK_STATUS

tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.cc

@@ -0,0 +1,506 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+//===- kernel_creator.cc ----------------------------------------*- C++ -*-===//
+//
+// This file implements the function to compile a TF kernel function to gpu
+// binary (hsaco for AMD, cubin for NVIDIA) or to a gpu binary with host side.
+//
+//===----------------------------------------------------------------------===//
+#include "tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.h"
+
+#include <string>
+
+#include "absl/status/status.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"  // from @llvm-project
+#include "mlir/Conversion/ComplexToStandard/ComplexToStandard.h"  // from @llvm-project
+#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"  // from @llvm-project
+#include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"  // from @llvm-project
+#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"  // from @llvm-project
+#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"  // from @llvm-project
+#include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h"  // from @llvm-project
+#include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"  // from @llvm-project
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Arith/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Dialect/Bufferization/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Dialect/Complex/IR/Complex.h"  // from @llvm-project
+#include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"  // from @llvm-project
+#include "mlir/Dialect/GPU/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Dialect/Linalg/Passes.h"  // from @llvm-project
+#include "mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Dialect/NVVM/Transforms/OptimizeForNVVM.h"  // from @llvm-project
+#include "mlir/Dialect/SCF/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypeInterfaces.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Interfaces/DataLayoutInterfaces.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "mlir/Target/LLVMIR/Dialect/Builtin/BuiltinToLLVMIRTranslation.h"  // from @llvm-project
+#include "mlir/Target/LLVMIR/Dialect/GPU/GPUToLLVMIRTranslation.h"  // from @llvm-project
+#include "mlir/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.h"  // from @llvm-project
+#include "mlir/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.h"  // from @llvm-project
+#include "mlir/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "mlir/Transforms/Passes.h"  // from @llvm-project
+#include "stablehlo/conversions/linalg/transforms/Passes.h"  // from @stablehlo
+#include "stablehlo/dialect/ChloOps.h"  // from @stablehlo
+#include "stablehlo/transforms/Passes.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewriters.h"
+#include "xla/mlir_hlo/deallocation/transforms/passes.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+#include "xla/mlir_hlo/mhlo/transforms/passes.h"
+#include "xla/mlir_hlo/transforms/gpu_passes.h"
+#include "xla/mlir_hlo/transforms/passes.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/statusor.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/statusor.h"
+
+namespace tensorflow {
+namespace kernel_gen {
+namespace {
+
+using mlir::Value;
+using mlir::func::FuncOp;
+using mlir::memref::RankOp;
+
+constexpr llvm::StringRef kGpuBinaryAttrName = "gpu.binary";
+
+/// Check if the size of the allocation is less than the given size. The
+/// transformation is only applied to small buffers since large buffers could
+/// exceed the stack space.
+bool IsSmallAlloc(Value alloc) {
+  constexpr unsigned kMaximumSizeInBytes = 64;
+  constexpr unsigned kMaxRankOfAllocatedMemRef = 1;
+
+  auto type = mlir::dyn_cast<mlir::ShapedType>(alloc.getType());
+  if (!type || !alloc.getDefiningOp<mlir::memref::AllocOp>()) return false;
+  if (!type.hasStaticShape()) {
+    // Check if the dynamic shape dimension of the alloc is produced by RankOp
+    // or SelectOp(_, RankOp, RankOp).
+    // If this is the case, it is likely to be small. Furthermore, the dimension
+    // is limited to the maximum rank of the allocated memref to avoid large
+    // values by multiplying several small values.
+    if (type.getRank() <= kMaxRankOfAllocatedMemRef) {
+      for (Value alloc_arg : alloc.getDefiningOp()->getOperands()) {
+        if (auto select = alloc_arg.getDefiningOp<mlir::arith::SelectOp>()) {
+          if (!select.getTrueValue().getDefiningOp<RankOp>() ||
+              !select.getFalseValue().getDefiningOp<RankOp>())
+            return false;
+        } else if (!alloc_arg.getDefiningOp<RankOp>()) {
+          return false;
+        }
+      }
+      return true;
+    }
+    return false;
+  }
+  unsigned bitwidth = mlir::DataLayout::closest(alloc.getDefiningOp())
+                          .getTypeSizeInBits(type.getElementType());
+  return type.getNumElements() * bitwidth <= kMaximumSizeInBytes * 8;
+}
+
+absl::Status LowerHloToJITInvocation(mlir::ModuleOp module,
+                                     llvm::ArrayRef<int64_t> tile_sizes,
+                                     llvm::ArrayRef<int64_t> unroll_factors,
+                                     bool enable_ftz, bool index_64bit,
+                                     bool jit_i64_indexed_for_large_tensors,
+                                     bool apply_cl_options) {
+  mlir::PassManager pm(module.getContext());
+  if (apply_cl_options) applyTensorflowAndCLOptions(pm);
+
+  pm.addNestedPass<FuncOp>(
+      mlir::kernel_gen::transforms::CreateFuncToJITInvocationPass(
+          tile_sizes, unroll_factors, enable_ftz, index_64bit,
+          /*cpu_codegen=*/false, jit_i64_indexed_for_large_tensors));
+  pm.addPass(mlir::kernel_gen::tf_framework::CreateEmbedTFFrameworkPass());
+  pm.addNestedPass<FuncOp>(
+      mlir::bufferization::createEmptyTensorToAllocTensorPass());
+  pm.addPass(mlir::createComputeOpAndFuncBufferizePass());
+
+  pm.addPass(mlir::createFinalBufferizePass(
+      /*alignment=*/64,
+      mlir::kernel_gen::transforms::populateExtraBufferizeDialects,
+      mlir::kernel_gen::transforms::populateExtraBufferizePatterns));
+
+  if (failed(pm.run(module))) {
+    return absl::InternalError("Lowering HLO to JIT invocation failed.");
+  }
+  return absl::OkStatus();
+}
+
+absl::Status LowerHlotoLoops(mlir::ModuleOp module,
+                             llvm::ArrayRef<int64_t> tile_sizes,
+                             llvm::ArrayRef<int64_t> unroll_factors,
+                             bool enable_ftz, bool index_64bit,
+                             bool jit_i64_indexed_for_large_tensors,
+                             bool apply_cl_options) {
+  mlir::PassManager pm(module.getContext());
+  if (apply_cl_options) applyTensorflowAndCLOptions(pm);
+  if (jit_i64_indexed_for_large_tensors) {
+    pm.addNestedPass<FuncOp>(
+        mlir::kernel_gen::transforms::CreateFuncToJITInvocationPass(
+            tile_sizes, unroll_factors, enable_ftz, index_64bit,
+            /*cpu_codegen=*/false,
+            /*jit_i64_indexed_for_large_tensors=*/true));
+  }
+
+  pm.addNestedPass<mlir::func::FuncOp>(
+      mlir::stablehlo::createChloLegalizeToStablehloPass());
+  pm.addPass(mlir::mhlo::createStablehloLegalizeToHloPass());
+
+  pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
+  pm.addNestedPass<FuncOp>(mlir::createCSEPass());
+  pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
+  pm.addNestedPass<FuncOp>(mlir::kernel_gen::createShapeSimplificationPass());
+  pm.addNestedPass<FuncOp>(mlir::kernel_gen::createMergeAssumingOpsPass());
+  pm.addNestedPass<FuncOp>(mlir::kernel_gen::createBroadcastPropagationPass());
+  pm.addNestedPass<FuncOp>(mlir::createCanonicalizerPass());
+  pm.addNestedPass<FuncOp>(mlir::createCSEPass());
+
+  // Transform HLO operations to Linalg and standard.
+  mlir::mhlo::HloLegalizeToStablehloPassOptions options;
+  options.allow_xla_features_ = true;  // MinimumBroadcastOp
+  pm.addPass(mlir::mhlo::createHloLegalizeToStablehloPass(options));
+
+  // Allowing implicit captures of scalar inputs causes crashes in kernelgen
+  // JIT kernels.
+  mlir::stablehlo::StablehloLegalizeToLinalgPassOptions linalg_options;
+  linalg_options.captureScalarInputs = false;
+  pm.addNestedPass<FuncOp>(
+      mlir::stablehlo::createStablehloLegalizeToLinalgPass(linalg_options));
+  pm.addPass(mlir::mhlo::createStablehloLegalizeToHloPass());
+
+  // Convert tensor.reshape to MHLO before running any bufferization passes.
+  // We'll need to teach this pipeline how to properly handle tensor.reshape
+  // in order to fully deprecate MHLO but in the meantime this ublocks lots of
+  // cleanup.
+  pm.addNestedPass<FuncOp>(
+      ::mlir::kernel_gen::createLegalizeTensorReshapePass());
+
+  // Remove the remaining references to unsigned types after all HLO compute
+  // operations were converted.
+  pm.addPass(mlir::stablehlo::createStablehloConvertToSignlessPass());
+
+  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addNestedPass<FuncOp>(mlir::createCSEPass());
+
+  // Convert operations from the Complex dialect to the Standard/Math dialects.
+  pm.addNestedPass<FuncOp>(::mlir::createConvertComplexToStandardPass());
+
+  // Fuse linalg operations.
+  pm.addPass(mlir::memref::createResolveShapedTypeResultDimsPass());
+  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addNestedPass<FuncOp>(mlir::createLinalgElementwiseOpFusionPass());
+
+  // Partial bufferization: Transforms in particular HLO and Linalg operations
+  // to their corresponding LHLO operations and converts the function signature.
+  // Leaves shape operations untouched.
+  //
+  // TODO(pifon): Rename the pass to CreateHloLinalgBufferizePass or bufferize
+  // in 2 steps: first Linalg, then Hlo. That would need refactoring of
+  // BufferizeTypeConverter.
+  pm.addNestedPass<FuncOp>(
+      mlir::bufferization::createEmptyTensorToAllocTensorPass());
+  pm.addPass(mlir::createComputeOpAndFuncBufferizePass());
+  pm.addNestedPass<FuncOp>(::mlir::createCanonicalizerPass());
+  pm.addNestedPass<FuncOp>(::mlir::createCSEPass());
+  // Remove copies which are introduced by canonicalizing
+  // BufferCastOp(TensorLoadOp).
+  pm.addNestedPass<FuncOp>(
+      mlir::kernel_gen::transforms::CreateCopyCleanupPass());
+  // Find candidates for buffer reuse. This is only successful if buffer size
+  // equality can be determined based on `linalg.generic` operations.
+  pm.addNestedPass<FuncOp>(
+      mlir::kernel_gen::transforms::CreateBufferReusePass());
+  // Approximate Tanh using standard operations.
+  pm.addNestedPass<FuncOp>(
+      ::mlir::mhlo::createLegalizeTrigonometricToApproximationPass());
+  // Transform the Linalg ops inside of the loop nest into parallel loops.
+  pm.addNestedPass<FuncOp>(::mlir::createConvertLinalgToParallelLoopsPass());
+
+  // Canonicalize the code to simplify index computations. This is needed so
+  // that loop bounds have the same value.
+  pm.addNestedPass<FuncOp>(::mlir::createCanonicalizerPass());
+  // Run CSE to ensure that loads and stores to the same subview get
+  // recognized as such.
+  pm.addNestedPass<FuncOp>(::mlir::createCSEPass());
+  // Collapse and tile parallel loops for GPU only.
+  pm.addNestedPass<FuncOp>(mlir::createCollapseParallelLoopsTo1DPass());
+  pm.addNestedPass<FuncOp>(
+      mlir::createTileLoopsPass(tile_sizes, unroll_factors));
+
+  pm.addNestedPass<FuncOp>(::mlir::createCanonicalizerPass());
+  pm.addNestedPass<FuncOp>(::mlir::createCSEPass());
+  if (failed(pm.run(module))) {
+    return absl::InternalError("Lowering HLO to loops failed.");
+  }
+  return absl::OkStatus();
+}
+
+absl::Status LowerLoopsToGPU(mlir::ModuleOp module, bool index_64bit,
+                             bool apply_cl_options) {
+  mlir::PassManager pm(module.getContext());
+  if (apply_cl_options) applyTensorflowAndCLOptions(pm);
+
+  // Greedily map the remaining loop to GPU hardware dimensions.
+  pm.addNestedPass<FuncOp>(mlir::createGpuMapParallelLoopsPass());
+
+  // Expand memref_reshape to its ranked form so that we can propagate
+  // scalars and avoid allocation.
+  pm.addNestedPass<FuncOp>(mlir::arith::createArithExpandOpsPass());
+  pm.addNestedPass<FuncOp>(mlir::memref::createExpandOpsPass());
+  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addPass(mlir::kernel_gen::transforms::CreateShapeToDescriptorsPass());
+  // Before bufferizing further, remove unused tensor_to_memref, so that we do
+  // not create allocations for tensor computations that are not actually
+  // needed.
+  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addNestedPass<FuncOp>(mlir::createCSEPass());
+  // Before inserting more allocs, map the ones we already have to the
+  // tf runtime. That ensures that all allocations for the actual computation
+  // end up on the device, whereas allocations for shape computation and host
+  // side things remain on the host.
+  // Longer term, this should be handled by proper device placement.
+  pm.addPass(mlir::kernel_gen::tf_framework::CreateEmbedTFFrameworkPass());
+  // Now lower the shape computations, bufferize all remaining ops and insert
+  // deallocs.
+  pm.addPass(mlir::createFinalBufferizePass(
+      /*alignment=*/64,
+      mlir::kernel_gen::transforms::populateExtraBufferizeDialects,
+      mlir::kernel_gen::transforms::populateExtraBufferizePatterns));
+  pm.addNestedPass<FuncOp>(::mlir::bufferization::createBufferHoistingPass());
+  pm.addNestedPass<FuncOp>(mlir::bufferization::createPromoteBuffersToStackPass(
+      [](Value alloc) { return IsSmallAlloc(alloc); }));
+  // Free all temporaries,
+  pm.addNestedPass<FuncOp>(mlir::deallocation::createBufferDeallocationPass());
+  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addNestedPass<FuncOp>(::mlir::createConvertLinalgToLoopsPass());
+
+  // Apply the mapping and go to GPU. We cannot do this earlier as the GPU
+  // dialect requires memrefs.
+  pm.addNestedPass<FuncOp>(mlir::createConvertParallelLoopToGpuPass());
+
+  // Some basic cleanup.
+  pm.addNestedPass<FuncOp>(::mlir::createCanonicalizerPass());
+  pm.addNestedPass<FuncOp>(::mlir::createCSEPass());
+  // Make loops with min bounds into a conditional plus static bounds.
+  pm.addNestedPass<FuncOp>(mlir::createForLoopSpecializationPass());
+  // Take launches to launches with kernels.
+  pm.addPass(mlir::createGpuLaunchSinkIndexComputationsPass());
+  const std::string gpuDataLayoutSpec =
+      index_64bit ? "#dlti.dl_spec<#dlti.dl_entry<index,64:i64>>"
+                  : "#dlti.dl_spec<#dlti.dl_entry<index,32:i32>>";
+  pm.addPass(
+      mlir::createGpuKernelOutliningPass({.dataLayoutStr = gpuDataLayoutSpec}));
+
+  pm.addPass(::mlir::createLowerAffinePass());
+  // Constraints are removed as late as possible and before lowering to CFG.
+  pm.addNestedPass<FuncOp>(::mlir::createConvertShapeConstraintsPass());
+  pm.addNestedPass<FuncOp>(::mlir::createCanonicalizerPass());
+  pm.addPass(::mlir::createSCFToControlFlowPass());
+  // Map asserts to the tensorflow framework.
+  pm.addPass(mlir::kernel_gen::tf_framework::CreateRewriteTFFrameworkAssert());
+  if (failed(pm.run(module))) {
+    return absl::InternalError("Lowering to GPU kernels failed.");
+  }
+  return absl::OkStatus();
+}
+
+absl::Status LowerKernelBodiesToLowLevelIr(mlir::ModuleOp module,
+                                           bool apply_cl_options,
+                                           const std::string& architecture) {
+#if !defined(TENSORFLOW_USE_ROCM) && !defined(GOOGLE_CUDA)
+  return absl::InternalError(
+      "Neither TENSORFLOW_USE_ROCM nor GOOGLE_CUDA are defined."
+      " Did you specify either --config=rocm or --config=cuda ?");
+#endif
+
+#if TENSORFLOW_USE_ROCM
+  auto gpu_modules = module.getOps<::mlir::gpu::GPUModuleOp>();
+  for (::mlir::gpu::GPUModuleOp gpu_module : gpu_modules) {
+    gpu_module.walk([&](mlir::gpu::GPUFuncOp gpu_kernel) {
+      if (gpu_kernel.isKernel()) {
+        gpu_kernel->setAttr(
+            "rocdl.max_flat_work_group_size",
+            mlir::IntegerAttr::get(
+                mlir::IntegerType::get(module.getContext(), 32), 1024));
+      }
+    });
+  }
+#endif
+
+  mlir::PassManager pm(module.getContext());
+  // We cannot verify as the signature of the kernel is rewritten.
+  // pm.enableVerifier(false);
+  if (apply_cl_options) tensorflow::applyTensorflowAndCLOptions(pm);
+  auto& kernelPm = pm.nest<::mlir::gpu::GPUModuleOp>();
+  kernelPm.addPass(::mlir::createSCFToControlFlowPass());
+#if TENSORFLOW_USE_ROCM
+  kernelPm.addPass(mlir::createGpuKernelToRocdlPass(architecture));
+#elif GOOGLE_CUDA
+  kernelPm.addPass(mlir::createGpuKernelToNvvmPass());
+  kernelPm.addPass(mlir::NVVM::createNVVMOptimizeForTargetPass());
+#endif
+  // Remove all location information to prevent a debug build.
+  pm.addPass(::mlir::createStripDebugInfoPass());
+
+  if (failed(pm.run(module))) {
+    return absl::InternalError("Lowering to low-level device IR failed.");
+  }
+
+  return absl::OkStatus();
+}
+
+absl::Status AmendKernelLLVMIRWithStaticKnowledge(mlir::ModuleOp module,
+                                                  bool apply_cl_options) {
+  mlir::PassManager pm(module.getContext());
+  if (apply_cl_options) applyTensorflowAndCLOptions(pm);
+
+  pm.addNestedPass<FuncOp>(
+      mlir::kernel_gen::transforms::CreatePropagateShapeKnowledgeToKernels());
+  pm.addNestedPass<FuncOp>(
+      mlir::kernel_gen::transforms::CreatePropagateTfAbiKnowledgeToKernels());
+
+  return failed(pm.run(module))
+             ? absl::InternalError(
+                   "Amending LLVMIR with static knowledge failed.")
+             : absl::OkStatus();
+}
+
+absl::Status GenerateDeviceCode(mlir::ModuleOp module,
+                                llvm::StringRef gpu_binary_attr_name,
+                                llvm::ArrayRef<std::string> architectures,
+                                bool print_ptx, bool print_llvmir,
+                                bool enable_ftz, bool apply_cl_options) {
+  mlir::PassManager pm(module.getContext());
+  if (apply_cl_options) applyTensorflowAndCLOptions(pm);
+  mlir::registerLLVMDialectTranslation(*module->getContext());
+
+  auto& kernel_pm = pm.nest<mlir::gpu::GPUModuleOp>();
+  // Remove debug information to ensure we do not create debug PTX.
+  kernel_pm.addPass(mlir::createStripDebugInfoPass());
+  kernel_pm.addPass(mlir::kernel_gen::transforms::CreateGpuKernelToBlobPass(
+      gpu_binary_attr_name, architectures, print_ptx, print_llvmir,
+      enable_ftz));
+
+  return failed(pm.run(module))
+             ? absl::InternalError("Generating device code failed.")
+             : absl::OkStatus();
+}
+
+absl::Status LowerHostSideToFinalForm(mlir::ModuleOp module,
+                                      bool apply_cl_options) {
+  mlir::PassManager pm(module.getContext());
+  if (apply_cl_options) applyTensorflowAndCLOptions(pm);
+
+  pm.addPass(mlir::kernel_gen::transforms::CreateTFKernelToLLVMPass(
+      kGpuBinaryAttrName));
+  pm.addPass(mlir::createReconcileUnrealizedCastsPass());
+  pm.addPass(mlir::createCanonicalizerPass());
+  pm.addPass(mlir::createCSEPass());
+
+  return failed(pm.run(module))
+             ? absl::InternalError("Final lowering of host side failed.")
+             : absl::OkStatus();
+}
+
+}  // namespace
+
+absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> SetupContextAndParseModule(
+    mlir::MLIRContext& context, llvm::StringRef tf_code) {
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::chlo::ChloDialect, mlir::mhlo::MhloDialect>();
+  registry.insert<mlir::shape::ShapeDialect, mlir::scf::SCFDialect,
+                  mlir::cf::ControlFlowDialect>();
+  registry.insert<mlir::complex::ComplexDialect, mlir::func::FuncDialect>();
+  mlir::registerBuiltinDialectTranslation(registry);
+  mlir::registerGPUDialectTranslation(registry);
+  mlir::registerLLVMDialectTranslation(registry);
+  mlir::registerNVVMDialectTranslation(registry);
+  mlir::registerROCDLDialectTranslation(registry);
+  context.appendDialectRegistry(registry);
+  mlir::OwningOpRef<mlir::ModuleOp> module =
+      mlir::parseSourceString<mlir::ModuleOp>(tf_code, &context);
+  if (!module) {
+    return absl::Status(absl::StatusCode::kInvalidArgument,
+                        "invalid kernel IR");
+  }
+  return module;
+}
+
+absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> GenerateKernelForHloCode(
+    mlir::MLIRContext& context, llvm::StringRef tf_code,
+    llvm::ArrayRef<std::string> architectures,
+    llvm::ArrayRef<int64_t> tile_sizes, llvm::ArrayRef<int64_t> unroll_factors,
+    bool print_ptx, bool print_llvmir, bool enable_ftz, bool index_64bit,
+    bool jit_compile, bool jit_i64_indexed_for_large_tensors,
+    bool apply_cl_options) {
+  if (jit_compile && jit_i64_indexed_for_large_tensors) {
+    return absl::Status(
+        absl::StatusCode::kInvalidArgument,
+        "jit compilation for large tensors "
+        "(`jit_i64_indexed_for_large_tensors`) and unconditioned jit "
+        "compilation (`jit`) must not be requested together");
+  }
+
+  TF_ASSIGN_OR_RETURN(mlir::OwningOpRef<mlir::ModuleOp> module,
+                      SetupContextAndParseModule(context, tf_code));
+
+  if (jit_compile) {
+    assert(!jit_i64_indexed_for_large_tensors &&
+           "expect to have reported an error earlier");
+    TF_RETURN_IF_ERROR(LowerHloToJITInvocation(
+        module.get(), tile_sizes, unroll_factors, enable_ftz, index_64bit,
+        /*jit_i64_indexed_for_large_tensors=*/false, apply_cl_options));
+  } else {
+    TF_RETURN_IF_ERROR(LowerHlotoLoops(
+        module.get(), tile_sizes, unroll_factors, enable_ftz, index_64bit,
+        jit_i64_indexed_for_large_tensors, apply_cl_options));
+    TF_RETURN_IF_ERROR(
+        LowerLoopsToGPU(module.get(), index_64bit, apply_cl_options));
+
+    // Note: we're just passing the first architecture out of the list. This
+    // should be sufficient for now, but in the future perhaps we'll need
+    // restructure this code to generate separate MLIR modules for each
+    // architecture.
+    const std::string& first_architecture =
+        !architectures.empty() ? architectures[0] : "";
+    TF_RETURN_IF_ERROR(LowerKernelBodiesToLowLevelIr(
+        module.get(), apply_cl_options, first_architecture));
+    TF_RETURN_IF_ERROR(
+        AmendKernelLLVMIRWithStaticKnowledge(module.get(), apply_cl_options));
+    TF_RETURN_IF_ERROR(GenerateDeviceCode(
+        module.get(), kGpuBinaryAttrName, architectures, print_ptx,
+        print_llvmir, enable_ftz, apply_cl_options));
+  }
+
+  TF_RETURN_IF_ERROR(LowerHostSideToFinalForm(module.get(), apply_cl_options));
+
+  return module;
+}
+
+}  // namespace kernel_gen
+}  // namespace tensorflow

tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.h

@@ -0,0 +1,58 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+//===- kernel_creator.h -----------------------------------------*- C++ -*-===//
+//
+// This file declares the function to compile a TF kernel function to gpu
+// binary (hsaco for AMD, cubin for NVIDIA) or to a gpu binary with host side.
+//
+//===----------------------------------------------------------------------===//
+#ifndef TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_KERNEL_CREATOR_H_
+#define TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_KERNEL_CREATOR_H_
+
+#include <cstdint>
+#include <string>
+#include <utility>
+
+#include "absl/status/statusor.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/StringRef.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "tensorflow/core/platform/statusor.h"
+
+namespace tensorflow {
+namespace kernel_gen {
+
+// Parses tf_code to create a module. An MLIRContext is taken in case any
+// unexpected dialects are needed.
+absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> SetupContextAndParseModule(
+    mlir::MLIRContext& context, llvm::StringRef tf_code);
+
+// Converts TF code to LLVM with or without GPU support.
+absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> GenerateKernelForHloCode(
+    mlir::MLIRContext& context, llvm::StringRef tf_code,
+    llvm::ArrayRef<std::string> architectures,
+    llvm::ArrayRef<int64_t> tile_sizes, llvm::ArrayRef<int64_t> unroll_factors,
+    bool print_ptx, bool print_llvmir, bool enable_ftz, bool index_64bit,
+    bool jit_compile, bool jit_i64_indexed_for_large_tensors,
+    bool apply_cl_options);
+
+}  // namespace kernel_gen
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_KERNEL_CREATOR_H_

tensorflow/compiler/mlir/tools/kernel_gen/tests/BUILD

@@ -0,0 +1,27 @@
+load("//tensorflow:tensorflow.default.bzl", "filegroup")
+load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+glob_lit_tests(
+    name = "all_tests",
+    data = [":test_utilities"],
+    driver = "//tensorflow/compiler/mlir:run_lit.sh",
+    test_file_exts = ["mlir"],
+)
+
+# Bundle together all of the test utilities that are used by tests.
+filegroup(
+    name = "test_utilities",
+    testonly = True,
+    data = [
+        "//tensorflow/compiler/mlir:tf-opt",
+        "//tensorflow/compiler/mlir/tools/kernel_gen:kernel-gen-opt",
+        "@llvm-project//llvm:FileCheck",
+        "@llvm-project//mlir:run_lit.sh",
+        "@xla//xla/mlir_hlo:mlir-hlo-opt",
+    ],
+)

tensorflow/compiler/mlir/tools/kernel_gen/tests/broadcast_propagation.mlir

@@ -0,0 +1,238 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s --split-input-file --mhlo-broadcast-propagation | \
+// RUN: FileCheck %s
+
+// CHECK-LABEL: @single_bcast
+// CHECK-SAME:  %[[ARG0:.*]]: tensor<16x?xf32>, %[[ARG1:.*]]: tensor<16x?xf32>, %[[SHAPE:.*]]: tensor<3xindex>
+func.func @single_bcast(%arg0 : tensor<16x?xf32>, %arg1 : tensor<16x?xf32>,
+    %shape : tensor<3xindex>) -> tensor<?x16x?xf32> {
+  // CHECK-DAG: %[[BCASTED_ARG0:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[SHAPE]])
+  // CHECK-DAG: %[[BCASTED_ARG1:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[SHAPE]])
+  // CHECK-DAG: %[[ADD:.*]] = mhlo.add %[[BCASTED_ARG0]], %[[BCASTED_ARG1]] : [[BCASTED_TY:tensor<\?x16x\?xf32>]]
+  // CHECK-DAG: %[[MUL:.*]] = mhlo.multiply %[[ADD]], %[[ADD]] : [[BCASTED_TY]]
+  // CHECK-DAG: %[[SUB:.*]] = mhlo.subtract %[[ADD]], %[[ADD]] : [[BCASTED_TY]]
+  // CHECK-DAG: %[[DIV:.*]] = mhlo.divide %[[MUL]], %[[SUB]] : [[BCASTED_TY]]
+  // CHECK:     return %[[DIV]] : [[BCASTED_TY]]
+  %0 = mhlo.add %arg0, %arg1 : tensor<16x?xf32>
+  %1 = mhlo.multiply %0, %0 : tensor<16x?xf32>
+  %2 = mhlo.subtract %0, %0 : tensor<16x?xf32>
+  %3 = mhlo.divide %1, %2 : tensor<16x?xf32>
+  %4 = "mhlo.dynamic_broadcast_in_dim"(%3, %shape) {
+      broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} :
+      (tensor<16x?xf32>, tensor<3xindex>) -> tensor<?x16x?xf32>
+  func.return %4 : tensor<?x16x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @single_bcast_ensure_order
+// CHECK-SAME:  %[[ARG0:.*]]: tensor<16x?xf32>, %[[ARG1:.*]]: tensor<16x?xf32>, %[[SHAPE:.*]]: tensor<3xindex>
+func.func @single_bcast_ensure_order(%arg0 : tensor<16x?xf32>, %arg1 : tensor<16x?xf32>,
+    %shape : tensor<3xindex>) -> tensor<?x16x?xf32> {
+  // CHECK-DAG: %[[BCASTED_ARG0:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[SHAPE]])
+  // CHECK-DAG: %[[BCASTED_ARG1:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[SHAPE]])
+  // CHECK-DAG: %[[ADD:.*]] = mhlo.add %[[BCASTED_ARG0]], %[[BCASTED_ARG1]] : [[BCASTED_TY:tensor<\?x16x\?xf32>]]
+  // CHECK-DAG: %[[MUL:.*]] = mhlo.multiply %[[ADD]], %[[ADD]] : [[BCASTED_TY]]
+  // CHECK-DAG: %[[SUB:.*]] = mhlo.subtract %[[MUL]], %[[MUL]] : [[BCASTED_TY]]
+  // CHECK-DAG: %[[DIV:.*]] = mhlo.divide %[[ADD]], %[[SUB]] : [[BCASTED_TY]]
+  // CHECK:     return %[[DIV]] : [[BCASTED_TY]]
+  %0 = mhlo.add %arg0, %arg1 : tensor<16x?xf32>
+  %1 = mhlo.multiply %0, %0 : tensor<16x?xf32>
+  %2 = mhlo.subtract %1, %1 : tensor<16x?xf32>
+  %3 = mhlo.divide %0, %2 : tensor<16x?xf32>
+  %4 = "mhlo.dynamic_broadcast_in_dim"(%3, %shape) {
+      broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} :
+      (tensor<16x?xf32>, tensor<3xindex>) -> tensor<?x16x?xf32>
+  func.return %4 : tensor<?x16x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @double_bcasts
+// CHECK-SAME:  %[[ARG0:.*]]: tensor<16x?xf32>, %[[ARG1:.*]]: tensor<16x?xf32>, %[[SHAPE0:.*]]: tensor<3xindex>, %[[SHAPE1:.*]]: tensor<3xindex>
+func.func @double_bcasts(%arg0 : tensor<16x?xf32>, %arg1 : tensor<16x?xf32>,
+    %shape0 : tensor<3xindex>, %shape1 : tensor<3xindex>)
+    -> (tensor<?x16x?xf32>, tensor<?x16x?xf32>) {
+  // CHECK-DAG: %[[BCASTED_ARG00:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[SHAPE0]]) [[BCAST_DIMS0:<{broadcast_dimensions = dense<\[1, 2\]> : tensor<2xi64>}>]]
+  // CHECK-DAG: %[[BCASTED_ARG01:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[SHAPE1]]) [[BCAST_DIMS1:<{broadcast_dimensions = dense<\[0, 2\]> : tensor<2xi64>}>]]
+  // CHECK-DAG: %[[BCASTED_ARG10:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[SHAPE0]]) [[BCAST_DIMS0]]
+  // CHECK-DAG: %[[BCASTED_ARG11:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[SHAPE1]]) [[BCAST_DIMS1]]
+  // CHECK-DAG: %[[ADD0:.*]] = mhlo.add %[[BCASTED_ARG00]], %[[BCASTED_ARG10]] : [[BCAST_TY:tensor<\?x16x\?xf32>]]
+  // CHECK-DAG: %[[ADD1:.*]] = mhlo.add %[[BCASTED_ARG01]], %[[BCASTED_ARG11]] : [[BCAST_TY]]
+  // CHECK-DAG: %[[MUL0:.*]] = mhlo.multiply %[[ADD0]], %[[ADD0]] : [[BCAST_TY]]
+  // CHECK-DAG: %[[MUL1:.*]] = mhlo.multiply %[[ADD1]], %[[ADD1]] : [[BCAST_TY]]
+  // CHECK-DAG: %[[SUB0:.*]] = mhlo.subtract %[[ADD0]], %[[ADD0]] : [[BCAST_TY]]
+  // CHECK-DAG: %[[SUB1:.*]] = mhlo.subtract %[[ADD1]], %[[ADD1]] : [[BCAST_TY]]
+  // CHECK-DAG: %[[DIV0:.*]] = mhlo.divide %[[MUL0]], %[[SUB0]] : [[BCAST_TY]]
+  // CHECK-DAG: %[[DIV1:.*]] = mhlo.divide %[[MUL1]], %[[SUB1]] : [[BCAST_TY]]
+  // CHECK:     return %[[DIV0]], %[[DIV1]] : [[BCAST_TY]], [[BCAST_TY]]
+  %0 = mhlo.add %arg0, %arg1 : tensor<16x?xf32>
+  %1 = mhlo.multiply %0, %0 : tensor<16x?xf32>
+  %2 = mhlo.subtract %0, %0 : tensor<16x?xf32>
+  %3 = mhlo.divide %1, %2 : tensor<16x?xf32>
+  %4 = "mhlo.dynamic_broadcast_in_dim"(%3, %shape0) {
+      broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>} :
+      (tensor<16x?xf32>, tensor<3xindex>) -> tensor<?x16x?xf32>
+  %5 = "mhlo.dynamic_broadcast_in_dim"(%3, %shape1) {
+      broadcast_dimensions = dense<[0, 2]> : tensor<2xi64>} :
+      (tensor<16x?xf32>, tensor<3xindex>) -> tensor<?x16x?xf32>
+  func.return %4, %5 : tensor<?x16x?xf32>, tensor<?x16x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @late_output_dimensions
+// CHECK: %[[ARG0:.*]]: tensor<?x32xf32>, %[[ARG1:.*]]: tensor<?x32xf32>, %[[ARG2:.*]]: tensor<?x?x?xf32>
+func.func @late_output_dimensions(%arg0 : tensor<?x32xf32>, %arg1 : tensor<?x32xf32>,
+    %arg2 : tensor<?x?x?xf32>) -> tensor<?x?x32xf32> {
+  // CHECK-DAG: %[[SHAPE:.*]] = shape.shape_of %[[ARG2]]
+  // CHECK-DAG: %[[BCASTED_ARG0:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[SHAPE]]) <{broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>}>
+  // CHECK-DAG: %[[BCASTED_ARG1:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[SHAPE]]) <{broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>}>
+  // CHECK-DAG: %[[SUB:.*]] = mhlo.subtract %[[BCASTED_ARG0]], %[[BCASTED_ARG1]] : [[BCAST_TY:tensor<\?x\?x32xf32>]]
+  // CHECK-DAG: %[[ADD:.*]] = mhlo.add %[[SUB]], %[[SUB]] : [[BCAST_TY]]
+  // CHECK:     return %[[ADD]] : [[BCAST_TY]]
+  %0 = mhlo.subtract %arg0, %arg1 : tensor<?x32xf32>
+  %1 = shape.shape_of %arg2 : tensor<?x?x?xf32> -> tensor<3xindex>
+  %2 = "mhlo.dynamic_broadcast_in_dim"(%0, %1)
+      {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} :
+      (tensor<?x32xf32>, tensor<3xindex>) -> tensor<?x?x32xf32>
+  %3 = mhlo.add %2, %2 : tensor<?x?x32xf32>
+  func.return %3 : tensor<?x?x32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @very_late_output_dimensions
+// CHECK: (%[[ARG0:.*]]: tensor<?x32xf32>, %[[ARG1:.*]]: tensor<?x32xf32>, %[[ARG2:.*]]: tensor<?x?x?xf32>)
+func.func @very_late_output_dimensions(%arg0 : tensor<?x32xf32>,
+    %arg1 : tensor<?x32xf32>, %arg2 : tensor<?x?x?xf32>) -> tensor<?x?x32xf32> {
+  // CHECK-DAG: %[[S:.*]] = shape.shape_of %[[ARG2:.*]]
+  // CHECK-DAG: %[[BCASTED_ARG0:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG0]], %[[S]])
+  // CHECK-DAG: %[[BCASTED_ARG1:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG1]], %[[S]])
+  // CHECK-DAG: %[[ACC0:.*]] = mhlo.add %[[BCASTED_ARG0]], %[[BCASTED_ARG1]]
+  // CHECK-DAG: %[[ACC1:.*]] = mhlo.multiply %[[ACC0]], %[[BCASTED_ARG1]]
+  // CHECK-DAG: %[[ACC2:.*]] = mhlo.subtract %[[ACC1]], %[[BCASTED_ARG1]]
+  // CHECK-DAG: %[[ACC3:.*]] = mhlo.divide %[[ACC2]], %[[BCASTED_ARG1]]
+  // CHECK:     return %[[ACC3]]
+  %acc0 = mhlo.add %arg0, %arg1 : tensor<?x32xf32>
+  %acc1 = mhlo.multiply %acc0, %arg1 : tensor<?x32xf32>
+  %acc2 = mhlo.subtract %acc1, %arg1 : tensor<?x32xf32>
+  %acc3 = mhlo.divide %acc2, %arg1 : tensor<?x32xf32>
+  %1 = shape.shape_of %arg2 : tensor<?x?x?xf32> -> tensor<3xindex>
+  %3 = "mhlo.dynamic_broadcast_in_dim"(%acc3, %1) <{broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>}> : (tensor<?x32xf32>, tensor<3xindex>) -> tensor<?x?x32xf32>
+  func.return %3 : tensor<?x?x32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @propagate_within_block
+// CHECK: (%[[ARG0:.*]]: tensor<?x32xf32>, %[[ARG1:.*]]: tensor<?x32xf32>, %[[ARG2:.*]]: tensor<?x32xf32>, %[[W:.*]]: !shape.witness, %[[SHAPE:.*]]: tensor<3xindex>)
+func.func @propagate_within_block(%arg0 : tensor<?x32xf32>, %arg1 : tensor<?x32xf32>,
+    %arg2 : tensor<?x32xf32>, %w : !shape.witness, %shape : tensor<3xindex>)
+    -> tensor<?x?x32xf32> {
+  // CHECK-DAG: %[[SUB:.*]] = mhlo.subtract %[[ARG0]], %[[ARG1]] : tensor<?x32xf32>
+  // CHECK:     %[[RESULT:.*]] = shape.assuming %[[W]]
+  // CHECK-DAG:   %[[BCASTED_ARG2:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG2]], %[[SHAPE]])
+  // CHECK-DAG:   %[[BCASTED_SUB:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[SUB]], %[[SHAPE]])
+  // CHECK-DAG:   %[[ADD:.*]] = mhlo.add %[[BCASTED_SUB]], %[[BCASTED_ARG2]] : tensor<?x?x32xf32>
+  // CHECK:       shape.assuming_yield %[[ADD]]
+  // CHECK:     return %[[RESULT]]
+  %sub = mhlo.subtract %arg0, %arg1 : tensor<?x32xf32>
+  %result = shape.assuming %w -> tensor<?x?x32xf32> {
+    %add = mhlo.add %sub, %arg2 : tensor<?x32xf32>
+    %bcasted = "mhlo.dynamic_broadcast_in_dim"(%add, %shape)
+        {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} :
+        (tensor<?x32xf32>, tensor<3xindex>) -> tensor<?x?x32xf32>
+    shape.assuming_yield %bcasted : tensor<?x?x32xf32>
+  }
+  func.return %result : tensor<?x?x32xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @propagate_within_block_2
+// CHECK-SAME:  %[[ARG:.*]]: tensor<?x?x?xf32>, %[[S:.*]]: tensor<3xindex>, %[[W:.*]]: !shape.witness
+func.func @propagate_within_block_2(%arg : tensor<?x?x?xf32>,
+    %shape: tensor<3xindex>, %w: !shape.witness) -> tensor<?x?x?xf32> {
+  // CHECK-DAG: %[[TMP:.*]] = "mhlo.dynamic_broadcast_in_dim"(%arg0, %[[S]])
+  // CHECK:     %[[RES:.*]] = shape.assuming %[[W]]
+  // CHECK-DAG:   %[[INNER_RES:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[TMP]], %[[S]])
+  // CHECK:       shape.assuming_yield %[[INNER_RES]]
+  // CHECK:     return %[[RES]]
+  %0 = "mhlo.dynamic_broadcast_in_dim"(%arg, %shape)
+      {broadcast_dimensions = dense<[0, 1, 2]> : tensor<3xi64>}
+      : (tensor<?x?x?xf32>, tensor<3xindex>) -> tensor<?x?x?xf32>
+  %1 = shape.assuming %w -> tensor<?x?x?xf32> {
+    %2 = "mhlo.dynamic_broadcast_in_dim"(%0, %shape)
+        {broadcast_dimensions = dense<[0, 1, 2]> : tensor<3xi64>}
+        : (tensor<?x?x?xf32>, tensor<3xindex>) -> tensor<?x?x?xf32>
+    shape.assuming_yield %2 : tensor<?x?x?xf32>
+  }
+  func.return %1 : tensor<?x?x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @propagate_across_bcasts_cst_src
+// CHECK-SAME:  %[[ARG:.*]]: tensor<1xindex>
+func.func @propagate_across_bcasts_cst_src(%s : tensor<1xindex>) -> tensor<?xi1> {
+  // CHECK-DAG: %[[C1:.*]] = mhlo.constant dense<true> : tensor<i1>
+  // CHECK-DAG: %[[RES:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[C1]], %[[ARG]]) <{broadcast_dimensions = dense<> : tensor<0xi64>}> : (tensor<i1>, tensor<1xindex>) -> tensor<?xi1>
+  // CHECK:     return %[[RES]]
+  %0 = mhlo.constant dense<true> : tensor<i1>
+  %1 = "mhlo.dynamic_broadcast_in_dim"(%0, %s)
+      {broadcast_dimensions = dense<> : tensor<0xi64>}
+      : (tensor<i1>, tensor<1xindex>) -> tensor<?xi1>
+  %2 = "mhlo.dynamic_broadcast_in_dim"(%1, %s)
+      {broadcast_dimensions = dense<0> : tensor<1xi64>}
+      : (tensor<?xi1>, tensor<1xindex>) -> tensor<?xi1>
+  func.return %2 : tensor<?xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @compose_bcast_dims
+// CHECK-SAME:  %[[ARG:.*]]: tensor<?x?xi1>, %[[S0:.*]]: tensor<3xindex>, %[[S1:.*]]: tensor<4xindex>
+func.func @compose_bcast_dims(%arg : tensor<?x?xi1>, %s0 : tensor<3xindex>, %s1 : tensor<4xindex>) -> tensor<1x?x1x?xi1> {
+  // CHECK-DAG: %[[RES:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG]], %[[S1]]) <{broadcast_dimensions = dense<[1, 3]> : tensor<2xi64>}> : (tensor<?x?xi1>, tensor<4xindex>) -> tensor<1x?x1x?xi1>
+  // CHECK:     return %[[RES]]
+  %1 = "mhlo.dynamic_broadcast_in_dim"(%arg, %s0)
+      {broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>}
+      : (tensor<?x?xi1>, tensor<3xindex>) -> tensor<1x?x?xi1>
+  %2 = "mhlo.dynamic_broadcast_in_dim"(%1, %s1)
+      {broadcast_dimensions = dense<[0, 1, 3]> : tensor<3xi64>}
+      : (tensor<1x?x?xi1>, tensor<4xindex>) -> tensor<1x?x1x?xi1>
+  func.return %2 : tensor<1x?x1x?xi1>
+}
+
+// -----
+
+// CHECK-LABEL: @propagate_across_bcasts
+// CHECK-SAME:  %[[ARG:.*]]: tensor<?x?x?xf32>, %[[S:.*]]: tensor<3xindex>
+func.func @propagate_across_bcasts(%arg : tensor<?x?x?xf32>, %shape : tensor<3xindex>) -> tensor<?x?x?xf32> {
+  // CHECK-DAG: %[[RES:.*]] = "mhlo.dynamic_broadcast_in_dim"(%[[ARG]], %[[S]]) <{broadcast_dimensions = dense<[0, 1, 2]> : tensor<3xi64>}> : (tensor<?x?x?xf32>, tensor<3xindex>) -> tensor<?x?x?xf32>
+  // CHECK:     return %[[RES]]
+  %0 = "mhlo.dynamic_broadcast_in_dim"(%arg, %shape)
+      {broadcast_dimensions = dense<[0, 1, 2]> : tensor<3xi64>}
+      : (tensor<?x?x?xf32>, tensor<3xindex>) -> tensor<?x?x?xf32>
+  %1 = "mhlo.dynamic_broadcast_in_dim"(%0, %shape)
+      {broadcast_dimensions = dense<[0, 1, 2]> : tensor<3xi64>}
+      : (tensor<?x?x?xf32>, tensor<3xindex>) -> tensor<?x?x?xf32>
+  %2 = "mhlo.dynamic_broadcast_in_dim"(%1, %shape)
+      {broadcast_dimensions = dense<[0, 1, 2]> : tensor<3xi64>}
+      : (tensor<?x?x?xf32>, tensor<3xindex>) -> tensor<?x?x?xf32>
+  func.return %2 : tensor<?x?x?xf32>
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/buffer_deallocation.mlir

@@ -0,0 +1,45 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s --buffer-deallocation | FileCheck %s
+
+// CHECK-LABEL: @jit_execute_allocation
+// CHECK-SAME:  %[[CTX:.*]]: !tf_framework.op_kernel_context, %[[ARG:.*]]: memref<*xf32>, %[[PRED:.*]]: i1, %[[CALLABLE:.*]]: !tf_framework.jit_callable, %[[SIZE:.*]]: index, %[[SHAPE:.*]]: memref<?xindex>
+func.func @jit_execute_allocation(%ctx: !tf_framework.op_kernel_context,
+    %arg: memref<*xf32>, %pred: i1, %callable: !tf_framework.jit_callable,
+    %size: index, %shape: memref<?xindex>) -> memref<*xf32> {
+  // CHECK:   %[[RES:.*]] = scf.if %[[PRED]]
+  // CHECK:     %[[JIT_EXECUTE:.*]] = tf_framework.jit_execute ctx(%[[CTX]]) %[[CALLABLE]](%[[ARG]])
+  // CHECK:     %[[INNER_RES:.*]] = bufferization.clone %[[JIT_EXECUTE]]
+  // CHECK:     tf_framework.dealloc(%[[CTX]], %[[JIT_EXECUTE]])
+  // CHECK:     scf.yield %[[INNER_RES]]
+  // CHECK:   else
+  // CHECK:     %[[ALLOC:.*]] = tf_framework.alloc(%[[CTX]], %[[SIZE]])
+  // CHECK:     %[[RESHAPE:.*]] = memref.reshape %[[ALLOC]](%[[SHAPE]])
+  // CHECK:     %[[INNER_RES:.*]] = bufferization.clone %[[RESHAPE]]
+  // CHECK:     tf_framework.dealloc(%[[CTX]], %[[ALLOC]])
+  // CHECK:     scf.yield %[[INNER_RES]]
+  // CHECK:   return %[[RES]]
+  %res = scf.if %pred -> (memref<*xf32>) {
+    %inner_res = tf_framework.jit_execute ctx(%ctx) %callable(%arg)
+        : memref<*xf32> -> memref<*xf32>
+    scf.yield %inner_res : memref<*xf32>
+  } else {
+    %alloc = tf_framework.alloc(%ctx, %size) : memref<?xf32>
+    %inner_res = memref.reshape %alloc(%shape)
+        : (memref<?xf32>, memref<?xindex>) -> memref<*xf32>
+    scf.yield %inner_res : memref<*xf32>
+  }
+  return %res : memref<*xf32>
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/buffer_reuse.mlir

@@ -0,0 +1,555 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s --buffer-reuse | FileCheck %s
+
+// CHECK-LABEL: @unique_reuse_output
+func.func @unique_reuse_output() -> (index, memref<2x3xi64>) attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_output = 1 : i32
+  %result_0 = arith.constant 1 : index
+  %result_1 = memref.alloc() : memref<2x3xi64>
+  func.return %result_0, %result_1 : index, memref<2x3xi64>
+}
+
+// CHECK-LABEL: @ambiguous_reuse_output
+func.func @ambiguous_reuse_output(%pred : i1)
+    -> (memref<2x3xi64>, memref<2x3xi64>) attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK: reuse_output = -1 : i32
+  %mem = memref.alloc() : memref<2x3xi64>
+  %other_mem = memref.alloc() : memref<2x3xi64>
+  cf.cond_br %pred, ^bb0, ^bb1
+^bb0:
+  func.return %mem, %other_mem : memref<2x3xi64>, memref<2x3xi64>
+^bb1:
+  func.return %other_mem, %mem : memref<2x3xi64>, memref<2x3xi64>
+}
+
+// CHECK-LABEL: @direct_reuse
+func.func @direct_reuse(%not_a_memref : index,
+                   %smaller : memref<5xi64>,
+                   %greater : memref<7xi64>,
+                   %different_element_type : memref<2x3xf32>,
+                   %reusable_0 : memref<2x3xi64>,
+                   %reusable_1 : memref<6xi64>) -> memref<2x3xi64>
+                   attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [1 : i32, 2 : i32, 3 : i32, 4 : i32, 5 : i32]
+  %result = memref.alloc() : memref<2x3xi64>
+  func.return %result : memref<2x3xi64>
+}
+
+// CHECK-LABEL: @local_reuse_with_memref_maps
+func.func @local_reuse_with_memref_maps(
+    %arg : memref<?xi64, strided<[3], offset: 2>>, %n : index)
+    -> memref<?xi64, strided<[3], offset: 2>> attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [0 : i32]
+  %result = memref.alloc(%n) : memref<?xi64, strided<[3], offset: 2>>
+  linalg.generic {
+    indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
+    iterator_types = ["parallel"]
+  } ins(%arg : memref<?xi64, strided<[3], offset: 2>>)
+    outs(%result : memref<?xi64, strided<[3], offset: 2>>) {
+  ^bb0(%a : i64, %b : i64):
+    linalg.yield %a : i64
+  }
+  func.return %result : memref<?xi64, strided<[3], offset: 2>>
+}
+
+// CHECK-LABEL: @local_reuse_with_broadcasting_memref_maps
+func.func @local_reuse_with_broadcasting_memref_maps(
+    %arg0 : memref<i64>, %arg1 : memref<?xi64>, %n : index)
+    -> memref<?xi64> attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [0 : i32, 1 : i32]
+  %result = memref.alloc(%n) : memref<?xi64>
+  linalg.generic {
+    indexing_maps = [affine_map<(i) -> ()>, affine_map<(i) -> (i)>, affine_map<(i) -> (i)>],
+    iterator_types = ["parallel"]
+  } ins(%arg0, %arg1 : memref<i64>, memref<?xi64>)
+    outs(%result : memref<?xi64>) {
+  ^bb0(%a : i64, %b : i64, %c : i64):
+    %add = arith.addi %a, %b : i64
+    linalg.yield %add : i64
+  }
+  func.return %result : memref<?xi64>
+}
+
+// CHECK-LABEL: @local_reuse_with_broadcasting_memref_maps2
+func.func @local_reuse_with_broadcasting_memref_maps2(
+    %arg0 : memref<?xi64>, %arg1 : memref<?xi64>)
+    -> memref<i64> attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [0 : i32, 1 : i32]
+  %result = memref.alloc() : memref<i64>
+  linalg.generic {
+    indexing_maps = [affine_map<(i) -> (i)>, affine_map<(i) -> (i)>, affine_map<(i) -> ()>],
+    iterator_types = ["parallel"]
+  } ins(%arg0, %arg1 : memref<?xi64>, memref<?xi64>)
+    outs(%result : memref<i64>) {
+  ^bb0(%a : i64, %b : i64, %c : i64):
+    %add = arith.addi %a, %b : i64
+    linalg.yield %add : i64
+  }
+  func.return %result : memref<i64>
+}
+
+// CHECK-LABEL: @local_reuse_with_broadcasting_memref_maps3
+func.func @local_reuse_with_broadcasting_memref_maps3(
+    %arg0 : memref<i64>, %arg1 : memref<?xi64>)
+    -> memref<i64> attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [0 : i32, 1 : i32]
+  %result = memref.alloc() : memref<i64>
+  linalg.generic {
+    indexing_maps = [affine_map<(i) -> ()>, affine_map<(i) -> (i)>, affine_map<(i) -> ()>],
+    iterator_types = ["parallel"]
+  } ins(%arg0, %arg1 : memref<i64>, memref<?xi64>)
+    outs(%result : memref<i64>) {
+  ^bb0(%a : i64, %b : i64, %c : i64):
+    %add = arith.addi %a, %b : i64
+    linalg.yield %add : i64
+  }
+  func.return %result : memref<i64>
+}
+
+// CHECK-LABEL: @nolocal_reuse_with_broadcasting_memref_maps
+func.func @nolocal_reuse_with_broadcasting_memref_maps(
+    %arg0 : memref<?xi64>, %arg1 : memref<?x?xi64>, %n : index)
+    -> memref<?xi64> attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [1 : i32]
+  %result = memref.alloc(%n) : memref<?xi64>
+  linalg.generic {
+    indexing_maps = [affine_map<(i,j) -> (i)>, affine_map<(i,j) -> (i,j)>, affine_map<(i,j) -> (j)>],
+    iterator_types = ["parallel", "parallel"]
+  } ins(%arg0, %arg1 : memref<?xi64>, memref<?x?xi64>)
+    outs(%result : memref<?xi64>) {
+  ^bb0(%a : i64, %b : i64, %c : i64):
+    %add = arith.addi %a, %b : i64
+    linalg.yield %add : i64
+  }
+  func.return %result : memref<?xi64>
+}
+
+// CHECK-LABEL: @nolocal_reuse_with_broadcasting_memref_maps2
+func.func @nolocal_reuse_with_broadcasting_memref_maps2(
+    %arg0 : memref<?x?x?xi64>, %arg1 : memref<?x?x?xi64>, %n : index)
+    -> memref<?x?xi64> attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = []
+  %result = memref.alloc(%n, %n) : memref<?x?xi64>
+  linalg.generic {
+    indexing_maps = [affine_map<(i,j,k) -> (i,j,k)>, affine_map<(i,j,k) -> (i,j,k)>, affine_map<(i,j,k) -> (k,i)>],
+    iterator_types = ["parallel", "parallel", "parallel"]
+  } ins(%arg0, %arg1 : memref<?x?x?xi64>, memref<?x?x?xi64>)
+    outs(%result : memref<?x?xi64>) {
+  ^bb0(%a : i64, %b : i64, %c : i64):
+    %add = arith.addi %a, %b : i64
+    linalg.yield %add : i64
+  }
+  func.return %result : memref<?x?xi64>
+}
+
+// CHECK-LABEL: @memref.reinterpret_cast_alias
+func.func @memref.reinterpret_cast_alias(%arg : memref<f32>, %n : index)
+    -> memref<?xf32> attributes {tf_entry} {
+  %c0 = arith.constant 0 : index
+  %reinterpreted = memref.reinterpret_cast %arg to
+      offset: [0],
+      sizes: [%n],
+      strides: [%c0]: memref<f32> to memref<?xf32, strided<[?], offset: ?>>
+
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [0 : i32]
+  %result = memref.alloc(%n) : memref<?xf32>
+
+  // reinterpreted (arg) and result are of same size.
+  linalg.generic {
+    indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+    iterator_types = ["parallel"]
+  } ins(%reinterpreted : memref<?xf32, strided<[?], offset: ?>>) outs(%result : memref<?xf32>) {
+  ^bb0(%a : f32, %b : f32):
+    linalg.yield %a : f32
+  }
+
+  func.return %result : memref<?xf32>
+}
+
+// CHECK-LABEL: @memref.cast_alias
+func.func @memref.cast_alias(%arg : memref<*xf32>, %n : index)
+    -> memref<?xf32> attributes {tf_entry} {
+  %casted = memref.cast %arg : memref<*xf32> to memref<?xf32>
+
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [0 : i32]
+  %result = memref.alloc(%n) : memref<?xf32>
+
+  // reinterpreted (arg) and result are of same size.
+  linalg.generic {
+    indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+    iterator_types = ["parallel"]
+  } ins(%casted : memref<?xf32>) outs(%result : memref<?xf32>) {
+  ^bb0(%a : f32, %b : f32):
+    linalg.yield %a : f32
+  }
+
+  func.return %result : memref<?xf32>
+}
+
+// CHECK-LABEL: @indirect_size_equality
+func.func @indirect_size_equality(%arg0 : memref<?xi64>,
+                             %arg1 : memref<?xi64>,
+                             %n : index) -> memref<?xi64>
+                             attributes {tf_entry} {
+  // arg0 and arg1 are equal in size.
+  linalg.generic {
+    indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+    iterator_types = ["parallel"]
+  } ins(%arg0 : memref<?xi64>) outs(%arg1 : memref<?xi64>) {
+  ^bb0(%a : i64, %b : i64):
+    linalg.yield %a : i64
+  }
+
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [0 : i32, 1 : i32]
+  %result = memref.alloc(%n) : memref<?xi64>
+
+  // arg0 and result are equal in size.
+  linalg.generic {
+    indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+    iterator_types = ["parallel"]
+  } ins(%arg0 : memref<?xi64>) outs(%result : memref<?xi64>) {
+  ^bb0(%a : i64, %b : i64):
+    linalg.yield %a : i64
+  }
+
+  func.return %result : memref<?xi64>
+}
+
+// CHECK-LABEL: @livetimes_incompatible
+func.func @livetimes_incompatible(%arg0 : memref<3xi64>)
+    -> memref<3xi64> attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = []
+  %result = memref.alloc() : memref<3xi64>
+
+  // Use newly allocated buffer.
+  %c0 = arith.constant 0 : index
+  %0 = memref.load %result[%c0] : memref<3xi64>
+
+  // Use argument buffer again.
+  %1 = memref.load %arg0[%c0] : memref<3xi64>
+
+  func.return %result : memref<3xi64>
+}
+
+// CHECK-LABEL: @never_used
+func.func @never_used(%arg0 : memref<3xi64>) -> memref<3xi64> attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [0 : i32]
+  %result = memref.alloc() : memref<3xi64>
+  %c0 = arith.constant 0 : index
+  %0 = memref.load %arg0[%c0] : memref<3xi64>
+  func.return %result : memref<3xi64>
+}
+
+// CHECK-LABEL: @branching_reuse
+func.func @branching_reuse(%pred : i1, %arg : memref<6xi64>) -> memref<6xi64>
+    attributes {tf_entry} {
+  cf.cond_br %pred, ^bb0, ^bb1
+^bb0:
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [1 : i32]
+  %mem0 = memref.alloc() : memref<6xi64>
+
+  // Keep buffer argument live in this branch. Reuse is still possible because
+  // the newly allocated buffer was not used yet.
+  %c0 = arith.constant 0 : index
+  memref.load %arg[%c0] : memref<6xi64>
+
+  cf.br ^bb2(%mem0 : memref<6xi64>)
+^bb1:
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [1 : i32]
+  %mem1 = memref.alloc() : memref<6xi64>
+  cf.br ^bb2(%mem1 : memref<6xi64>)
+^bb2(%result : memref<6xi64>):
+  func.return %result : memref<6xi64>
+}
+
+// CHECK-LABEL: @branching_no_reuse
+func.func @branching_no_reuse(%pred : i1, %arg : memref<6xi64>) -> memref<6xi64>
+    attributes {tf_entry} {
+  cf.cond_br %pred, ^bb0, ^bb1
+^bb0:
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = []
+  %mem0 = memref.alloc() : memref<6xi64>
+
+  // Use newly allocated memory immediately.
+  %c0 = arith.constant 0 : index
+  memref.load %mem0[%c0] : memref<6xi64>
+
+  // Keep buffer argument live in this branch and prevent reuse.
+  memref.load %arg[%c0] : memref<6xi64>
+
+  cf.br ^bb2(%mem0 : memref<6xi64>)
+^bb1:
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [1 : i32]
+  %mem1 = memref.alloc() : memref<6xi64>
+  cf.br ^bb2(%mem1 : memref<6xi64>)
+^bb2(%result : memref<6xi64>):
+  func.return %result : memref<6xi64>
+}
+
+// CHECK-LABEL: @branching_reuse_if
+func.func @branching_reuse_if(%pred : i1, %arg : memref<6xi64>)
+    -> memref<6xi64> attributes {tf_entry} {
+  %result = scf.if %pred -> (memref<6xi64>) {
+    // CHECK: memref.alloc
+    // CHECK-SAME: reuse_input_candidates = [1 : i32]
+    %mem0 = memref.alloc() : memref<6xi64>
+
+    // Keep buffer argument live in this branch. Reuse is still possible because
+    // the newly allocated buffer was not used yet.
+    %c0 = arith.constant 0 : index
+    memref.load %arg[%c0] : memref<6xi64>
+
+    scf.yield %mem0 : memref<6xi64>
+  } else {
+    // CHECK: memref.alloc
+    // CHECK-SAME: reuse_input_candidates = [1 : i32]
+    %mem1 = memref.alloc() : memref<6xi64>
+    scf.yield %mem1 : memref<6xi64>
+  }
+  func.return %result : memref<6xi64>
+}
+
+// CHECK-LABEL: @branching_no_reuse_if
+func.func @branching_no_reuse_if(%pred : i1, %arg : memref<6xi64>) -> memref<6xi64>
+    attributes {tf_entry} {
+  %result = scf.if %pred -> (memref<6xi64>) {
+    // CHECK: memref.alloc
+    // CHECK-SAME: reuse_input_candidates = []
+    %mem0 = memref.alloc() : memref<6xi64>
+
+    // Use newly allocated memory immediately.
+    %c0 = arith.constant 0 : index
+    memref.load %mem0[%c0] : memref<6xi64>
+
+    // Keep buffer argument live in this branch and prevent reuse.
+    memref.load %arg[%c0] : memref<6xi64>
+
+    scf.yield %mem0 : memref<6xi64>
+  } else {
+    // CHECK: memref.alloc
+    // CHECK-SAME: reuse_input_candidates = [1 : i32]
+    %mem1 = memref.alloc() : memref<6xi64>
+    scf.yield %mem1 : memref<6xi64>
+  }
+  func.return %result : memref<6xi64>
+}
+
+// CHECK-LABEL: @alloc_before_branching
+// New buffer is first used in the blocks succeeding its allocation block. In
+// both/all cases the newly allocated buffer is used after the buffer argument
+// is no longer live. Because these first uses are not block-local the analysis
+// does not detect this case (yet). It is correct but incomplete.
+func.func @alloc_before_branching(%pred : i1, %arg : memref<6xi64>) -> memref<6xi64>
+    attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = []
+  %mem = memref.alloc() : memref<6xi64>
+  %c0 = arith.constant 0 : index
+  cf.cond_br %pred, ^bb0, ^bb1
+^bb0:
+  // Last use of `arg` before first use of `mem` (can reuse).
+  memref.load %arg[%c0] : memref<6xi64>
+  memref.load %mem[%c0] : memref<6xi64>
+  func.return %mem : memref<6xi64>
+^bb1:
+  // Last use of `arg` before first use of `mem` (can reuse).
+  memref.load %arg[%c0] : memref<6xi64>
+  memref.load %mem[%c0] : memref<6xi64>
+  func.return %mem : memref<6xi64>
+}
+
+// CHECK-LABEL: @alloc_before_branching_2
+func.func @alloc_before_branching_2(%pred : i1, %arg : memref<6xi64>)
+    -> memref<6xi64> attributes {tf_entry} {
+  // CHECK: memref.alloc()
+  // CHECK-SAME: reuse_input_candidates = []
+  %mem = memref.alloc() : memref<6xi64>
+  %c0 = arith.constant 0 : index
+  cf.cond_br %pred, ^bb0, ^bb1
+^bb0:
+  // Last use of `arg` after first use of `mem` (cannot reuse).
+  memref.load %mem[%c0] : memref<6xi64>
+  memref.load %arg[%c0] : memref<6xi64>
+  func.return %mem : memref<6xi64>
+^bb1:
+  // Last use of `arg` before first use of `mem` (can reuse).
+  memref.load %arg[%c0] : memref<6xi64>
+  memref.load %mem[%c0] : memref<6xi64>
+  func.return %mem : memref<6xi64>
+}
+
+// CHECK-LABEL: @alloc_before_branching_if
+// New buffer is first used in the blocks succeeding its allocation block. In
+// both/all cases the newly allocated buffer is used after the buffer argument
+// is no longer live. Because these first uses are not block-local the analysis
+// does not detect this case (yet). It is correct but incomplete.
+func.func @alloc_before_branching_if(%pred : i1, %arg : memref<6xi64>) -> memref<6xi64>
+    attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = []
+  %mem = memref.alloc() : memref<6xi64>
+  %result = scf.if %pred -> (memref<6xi64>) {
+    // Last use of `arg` before first use of `mem` (can reuse).
+    %c0 = arith.constant 0 : index
+    memref.load %arg[%c0] : memref<6xi64>
+    memref.load %mem[%c0] : memref<6xi64>
+    scf.yield %mem : memref<6xi64>
+  } else {
+    // Last use of `arg` before first use of `mem` (can reuse).
+    %c0 = arith.constant 0 : index
+    memref.load %arg[%c0] : memref<6xi64>
+    memref.load %mem[%c0] : memref<6xi64>
+    scf.yield %mem : memref<6xi64>
+  }
+  func.return %result : memref<6xi64>
+}
+
+// CHECK-LABEL: @alloc_before_branching_2_if
+func.func @alloc_before_branching_2_if(%pred : i1, %arg : memref<6xi64>)
+    -> memref<6xi64> attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = []
+  %mem = memref.alloc() : memref<6xi64>
+  %result = scf.if %pred -> (memref<6xi64>) {
+    // Last use of `arg` after first use of `mem` (cannot reuse).
+    %c0 = arith.constant 0 : index
+    memref.load %mem[%c0] : memref<6xi64>
+    memref.load %arg[%c0] : memref<6xi64>
+    scf.yield %mem : memref<6xi64>
+  } else {
+    // Last use of `arg` before first use of `mem` (can reuse).
+    %c0 = arith.constant 0 : index
+    memref.load %arg[%c0] : memref<6xi64>
+    memref.load %mem[%c0] : memref<6xi64>
+    scf.yield %mem : memref<6xi64>
+  }
+  func.return %result : memref<6xi64>
+}
+
+// CHECK-LABEL: @abs_unranked_i64
+func.func @abs_unranked_i64(%arg : memref<*xi64>,
+                       %arg_shape : memref<?xindex>,
+                       %flat_shape : memref<1xindex>,
+                       %arg_size : index) -> memref<*xi64>
+                       attributes {tf_entry} {
+  %flat_arg = memref.reshape %arg(%flat_shape)
+      : (memref<*xi64>, memref<1xindex>) -> memref<?xi64>
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [0 : i32, 2 : i32], reuse_output = 0 : i32
+  %flat_result = memref.alloc(%arg_size) : memref<?xi64>
+  linalg.generic {
+    indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+    iterator_types = ["parallel"]
+  } ins(%flat_arg : memref<?xi64>) outs(%flat_result : memref<?xi64>) {
+  ^bb0(%a : i64, %b : i64):
+    %c0 = arith.constant 0 : i64
+    %a_pos = arith.cmpi sge, %a, %c0 : i64
+    %a_neg = arith.subi %c0, %a : i64
+    %a_abs = arith.select %a_pos, %a, %a_neg : i64
+    linalg.yield %a_abs : i64
+  }
+  %result = memref.reshape %flat_result(%arg_shape)
+      : (memref<?xi64>, memref<?xindex>) -> memref<*xi64>
+  func.return %result : memref<*xi64>
+}
+
+// CHECK-LABEL: @old_buffer_alias_outside_block
+func.func @old_buffer_alias_outside_block(%arg: memref<3xf32>)
+    attributes {llvm.emit_c_interface, tf_entry} {
+  %c0 = arith.constant 0 : index
+  %c1 = arith.constant 1 : index
+  %true = arith.constant true
+
+  // Alias outside of the block with the new buffer allocation.
+  %alias = memref.cast %arg : memref<3xf32> to memref<3xf32>
+
+  scf.if %true {
+
+    // Allocation and use of new buffer.
+    // CHECK: memref.alloc
+    // CHECK-SAME: reuse_input_candidates = [0 : i32]
+    %mem = memref.alloc() : memref<3xf32>
+    %use = memref.load %mem[%c0] : memref<3xf32>
+
+  } else {
+  }
+  func.return
+}
+
+// CHECK-LABEL: @index_element_type
+func.func @index_element_type(%arg : memref<2x3xindex>) -> memref<2x3xindex>
+    attributes {tf_entry} {
+  // CHECK: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [0 : i32]
+  %result = memref.alloc() : memref<2x3xindex>
+  func.return %result : memref<2x3xindex>
+}
+
+// Example as it occurs in the `tf.Abs` kernel for `f32`.
+// CHECK-LABEL: @abs_f32
+func.func @abs_f32(%arg0: memref<*xf32>) -> memref<*xf32>
+    attributes {llvm.emit_c_interface, tf_entry} {
+  %c0 = arith.constant 0 : index
+  %0 = shape.shape_of %arg0 : memref<*xf32> -> tensor<?xindex>
+  %1 = shape.num_elements %0 : tensor<?xindex> -> index
+  // CHECK-LABEL: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = []
+  %2 = memref.alloc() : memref<1xindex>
+  memref.store %1, %2[%c0] : memref<1xindex>
+  %3 = memref.reshape %arg0(%2)
+      : (memref<*xf32>, memref<1xindex>) -> memref<?xf32>
+  %4 = memref.dim %3, %c0 : memref<?xf32>
+  %5 = arith.index_cast %4 : index to i64
+  // CHECK-LABEL: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = []
+  %6 = memref.alloc() : memref<1xi64>
+  memref.store %5, %6[%c0] : memref<1xi64>
+  %7 = memref.load %6[%c0] : memref<1xi64>
+  %8 = arith.index_cast %7 : i64 to index
+  // CHECK-LABEL: memref.alloc
+  // CHECK-SAME: reuse_input_candidates = [0 : i32]
+  %9 = memref.alloc(%8) : memref<?xf32>
+  linalg.generic {
+    indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
+    iterator_types = ["parallel"]
+  } ins(%3 : memref<?xf32>) outs(%9 : memref<?xf32>) {
+  ^bb0(%arg1: f32, %arg2: f32):
+    %12 = math.absf %arg1 : f32
+    linalg.yield %12 : f32
+  }
+  %10 = bufferization.to_buffer %0 : tensor<?xindex> to memref<?xindex>
+  %11 = memref.reshape %9(%10)
+      : (memref<?xf32>, memref<?xindex>) -> memref<*xf32>
+  func.return %11 : memref<*xf32>
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/bufferize.mlir

@@ -0,0 +1,34 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s --split-input-file  --mlir-print-op-generic |\
+// RUN: mlir-hlo-opt --allow-unregistered-dialect \
+// RUN:   --computeop-and-func-bufferize |\
+// RUN: kernel-gen-opt \
+// RUN:    --kernelgen-final-bufferize | FileCheck %s
+// RUN: kernel-gen-opt %s --split-input-file  --mlir-print-op-generic |\
+// RUN: mlir-hlo-opt --allow-unregistered-dialect \
+// RUN:   --computeop-and-func-bufferize |\
+// RUN: kernel-gen-opt \
+// RUN:    --kernelgen-final-bufferize --promote-buffers-to-stack | FileCheck %s
+
+// CHECK-LABEL: @jit_execute
+// CHECK-SAME: (%[[F:.*]]: !tf_framework.jit_callable, %[[ARG:.*]]: memref<*xf32>) -> memref<*xf32>
+func.func @jit_execute(%f : !tf_framework.jit_callable, %arg : tensor<*xf32>)
+    -> tensor<*xf32> {
+  // CHECK: %[[RES:.*]] = tf_framework.jit_execute %[[F]](%[[ARG]]) : memref<*xf32> -> memref<*xf32>
+  // CHECK: return %[[RES]] : memref<*xf32>
+  %0 = tf_framework.jit_execute %f(%arg) : tensor<*xf32> -> tensor<*xf32>
+  func.return %0 : tensor<*xf32>
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/copy_cleanup.mlir

@@ -0,0 +1,202 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s --copy-cleanup --split-input-file | FileCheck %s
+
+#map0 = affine_map<(d0)[s0] -> (d0 * s0)>
+#map1 = affine_map<(d0) -> (d0)>
+builtin.module {
+  func.func @Copy(%lhs: memref<?xi16>, %rhs: memref<?xi16>) {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %size = memref.dim %rhs, %c0 : memref<?xi16>
+    %lhsCasted = memref.reinterpret_cast %lhs to offset: [0], sizes: [%size], strides: [%c0] : memref<?xi16> to memref<?xi16, #map0>
+    %lhsAlloc = memref.alloc(%size) : memref<?xi16>
+    memref.copy %lhsCasted, %lhsAlloc : memref<?xi16, #map0> to memref<?xi16>
+    %rhsCasted = memref.reinterpret_cast %rhs to offset: [0], sizes: [%size], strides: [%c1] : memref<?xi16> to memref<?xi16, #map0>
+    %rhsAlloc = memref.alloc(%size) : memref<?xi16>
+    memref.copy %rhsCasted, %rhsAlloc : memref<?xi16, #map0> to memref<?xi16>
+    %outputAlloc = memref.alloc(%size) : memref<?xi16>
+    linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel"]} ins(%lhsAlloc, %rhsAlloc : memref<?xi16>, memref<?xi16>) outs(%outputAlloc : memref<?xi16>) {
+    ^bb0(%arg1: i16, %arg2: i16, %arg3 : i16):
+      %and = arith.andi %arg1, %arg2 : i16
+      linalg.yield %and : i16
+    }
+    func.return
+  }
+}
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL:   func @Copy(
+// CHECK-SAME:              %[[LHS:.*]]: memref<?xi16>, %[[RHS:.*]]: memref<?xi16>) {
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[SIZE:.*]] = memref.dim %[[RHS]], %[[C0]] : memref<?xi16>
+// CHECK: %[[LHS_CASTED:.*]] = memref.reinterpret_cast %[[LHS]] to offset: [0], sizes: [%[[SIZE]]], strides: [%[[C0]]] : memref<?xi16> to memref<?xi16, #[[$MAP0]]>
+// CHECK: %[[RHS_CASTED:.*]] = memref.reinterpret_cast %[[RHS]] to offset: [0], sizes: [%[[SIZE]]], strides: [%[[C1]]] : memref<?xi16> to memref<?xi16, #[[$MAP0]]>
+// CHECK: %[[OUTPUT:.*]] = memref.alloc(%[[SIZE]]) : memref<?xi16>
+// CHECK: linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel"]} ins(%[[LHS_CASTED]], %[[RHS_CASTED]] : memref<?xi16, #[[$MAP0]]>, memref<?xi16, #[[$MAP0]]>) outs(%{{.*}} : memref<?xi16>) {
+// CHECK: ^bb0(%[[ARG1:.*]]: i16, %[[ARG2:.*]]: i16, %[[ARG3:.*]]: i16):
+// CHECK:   %[[AND:.*]] = arith.andi %[[ARG1]], %[[ARG2]] : i16
+// CHECK:   linalg.yield %[[AND]] : i16
+
+// -----
+
+// The target of the copy is also used to write.
+
+#map0 = affine_map<(d0)[s0] -> (d0 * s0)>
+#map1 = affine_map<(d0) -> (d0)>
+builtin.module {
+  func.func @CopyWithWrite(%lhs: memref<?xi16>, %rhs: memref<?xi16>) {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %size = memref.dim %rhs, %c0 : memref<?xi16>
+    %lhsCasted = memref.reinterpret_cast %lhs to offset: [0], sizes: [%size], strides: [%c0] : memref<?xi16> to memref<?xi16, #map0>
+    %lhsAlloc = memref.alloc(%size) : memref<?xi16>
+    memref.copy %lhsCasted, %lhsAlloc : memref<?xi16, #map0> to memref<?xi16>
+    %rhsCasted = memref.reinterpret_cast %rhs to offset: [0], sizes: [%size], strides: [%c1] : memref<?xi16> to memref<?xi16, #map0>
+    %rhsAlloc = memref.alloc(%size) : memref<?xi16>
+    memref.copy %rhsCasted, %rhsAlloc : memref<?xi16, #map0> to memref<?xi16>
+    linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel"]} ins(%lhsAlloc, %rhsAlloc : memref<?xi16>, memref<?xi16>) outs(%rhsAlloc : memref<?xi16>) {
+    ^bb0(%arg1: i16, %arg2: i16, %arg3 : i16):
+      %and = arith.andi %arg1, %arg2 : i16
+      linalg.yield %and : i16
+    }
+    func.return
+  }
+}
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL:   func @CopyWithWrite(
+// CHECK-SAME:              %[[LHS:.*]]: memref<?xi16>, %[[RHS:.*]]: memref<?xi16>) {
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[SIZE:.*]] = memref.dim %[[RHS]], %[[C0]] : memref<?xi16>
+// CHECK: %[[LHS_CASTED:.*]] = memref.reinterpret_cast %[[LHS]] to offset: [0], sizes: [%[[SIZE]]], strides: [%[[C0]]] : memref<?xi16> to memref<?xi16, #[[$MAP0]]>
+// CHECK: %[[RHS_CASTED:.*]] = memref.reinterpret_cast %[[RHS]] to offset: [0], sizes: [%[[SIZE]]], strides: [%[[C1]]] : memref<?xi16> to memref<?xi16, #[[$MAP0]]>
+// CHECK: %[[RHS_ALLOC:.*]] = memref.alloc(%[[SIZE]]) : memref<?xi16>
+// CHECK: memref.copy %[[RHS_CASTED]], %[[RHS_ALLOC]]
+// CHECK: linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel"]} ins(%[[LHS_CASTED]], %[[RHS_ALLOC]] : memref<?xi16, #[[$MAP0]]>, memref<?xi16>) outs(%[[RHS_ALLOC]] : memref<?xi16>) {
+// CHECK: ^bb0(%[[ARG1:.*]]: i16, %[[ARG2:.*]]: i16, %[[ARG3:.*]]: i16):
+// CHECK:   %[[AND:.*]] = arith.andi %[[ARG1]], %[[ARG2]] : i16
+// CHECK:   linalg.yield %[[AND]] : i16
+
+// -----
+
+// The source of the copy is mutated.
+
+#map0 = affine_map<(d0)[s0] -> (d0 * s0)>
+#map1 = affine_map<(d0) -> (d0)>
+builtin.module {
+  func.func @CopyWithMutation(%lhs: memref<?xi16>, %rhs: memref<?xi16>) {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c42 = arith.constant 42 : i16
+    %size = memref.dim %rhs, %c0 : memref<?xi16>
+    %lhsCasted = memref.reinterpret_cast %lhs to offset: [0], sizes: [%size], strides: [%c0] : memref<?xi16> to memref<?xi16, #map0>
+    %lhsAlloc = memref.alloc(%size) : memref<?xi16>
+    memref.copy %lhsCasted, %lhsAlloc : memref<?xi16, #map0> to memref<?xi16>
+    %rhsCasted = memref.reinterpret_cast %rhs to offset: [0], sizes: [%size], strides: [%c1] : memref<?xi16> to memref<?xi16, #map0>
+    %rhsAlloc = memref.alloc(%size) : memref<?xi16>
+    memref.copy %rhsCasted, %rhsAlloc : memref<?xi16, #map0> to memref<?xi16>
+    %outputAlloc = memref.alloc(%size) : memref<?xi16>
+    memref.store %c42, %lhs[%c0] : memref<?xi16>
+    linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel"]} ins(%lhsAlloc, %rhsAlloc : memref<?xi16>, memref<?xi16>) outs(%outputAlloc : memref<?xi16>) {
+    ^bb0(%arg1: i16, %arg2: i16, %arg3 : i16):
+      %and = arith.andi %arg1, %arg2 : i16
+      linalg.yield %and : i16
+    }
+    func.return
+  }
+}
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0)[s0] -> (d0 * s0)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL:   func @CopyWithMutation(
+// CHECK-SAME:              %[[LHS:.*]]: memref<?xi16>, %[[RHS:.*]]: memref<?xi16>) {
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[SIZE:.*]] = memref.dim %[[RHS]], %[[C0]] : memref<?xi16>
+// CHECK: %[[LHS_CASTED:.*]] = memref.reinterpret_cast %[[LHS]] to offset: [0], sizes: [%[[SIZE]]], strides: [%[[C0]]] : memref<?xi16> to memref<?xi16, #[[$MAP0]]>
+// CHECK: %[[LHS_ALLOC:.*]] = memref.alloc(%[[SIZE]]) : memref<?xi16>
+// CHECK: memref.copy %[[LHS_CASTED]], %[[LHS_ALLOC]]
+// CHECK: %[[RHS_CASTED:.*]] = memref.reinterpret_cast %[[RHS]] to offset: [0], sizes: [%[[SIZE]]], strides: [%[[C1]]] : memref<?xi16> to memref<?xi16, #[[$MAP0]]>
+// CHECK: %[[RHS_ALLOC:.*]] = memref.alloc(%[[SIZE]]) : memref<?xi16>
+// CHECK: memref.copy %[[RHS_CASTED]], %[[RHS_ALLOC]]
+// CHECK: %[[OUTPUT_ALLOC:.*]] = memref.alloc(%[[SIZE]]) : memref<?xi16>
+// CHECK: linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel"]} ins(%[[LHS_ALLOC]], %[[RHS_ALLOC]] : memref<?xi16>, memref<?xi16>) outs(%[[OUTPUT_ALLOC]] : memref<?xi16>) {
+// CHECK: ^bb0(%[[ARG1:.*]]: i16, %[[ARG2:.*]]: i16, %[[ARG3:.*]]: i16):
+// CHECK:   %[[AND:.*]] = arith.andi %[[ARG1]], %[[ARG2]] : i16
+// CHECK:   linalg.yield %[[AND]] : i16
+
+// -----
+
+#map0 = affine_map<(d0) -> (d0)>
+builtin.module  {
+  func.func @testCopyAfterLinalg(%arg2: memref<4xi32>, %arg3: memref<4xi32>, %arg4: memref<4xi32>) {
+    %0 = memref.alloc() : memref<4xi32>
+    linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel"]} ins(%arg2, %arg3 : memref<4xi32>, memref<4xi32>) outs(%0 : memref<4xi32>) {
+    ^bb0(%arg5: i32, %arg6: i32, %arg7: i32):
+      %1 = arith.addi %arg5, %arg6 : i32
+      linalg.yield %1 : i32
+    }
+    memref.copy %0, %arg4 : memref<4xi32> to memref<4xi32>
+    func.return
+  }
+}
+
+// CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL:func @testCopyAfterLinalg(
+// CHECK-SAME: %[[ARG0:.*]]: memref<4xi32>, %[[ARG1:.*]]: memref<4xi32>, %[[ARG2:.*]]: memref<4xi32>) {
+// CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%[[ARG0]], %[[ARG1]] : memref<4xi32>, memref<4xi32>) outs(%[[ARG2]] : memref<4xi32>) {
+// CHECK: ^bb0(%[[ARG3:.*]]: i32, %[[ARG4:.*]]: i32, %[[ARG5:.*]]: i32):
+// CHECK:   %[[R1:.*]] = arith.addi %[[ARG3]], %[[ARG4]] : i32
+// CHECK:   linalg.yield %[[R1]] : i32
+// CHECK: }
+// CHECK: return
+// CHECK: }
+
+// -----
+
+#map0 = affine_map<(d0) -> (d0)>
+builtin.module  {
+  func.func @testCopyAfterLinalgMutated(%arg2: memref<4xi32>, %arg3: memref<4xi32>, %arg4: memref<4xi32>) {
+    %c0 = arith.constant 0 : index
+    %c42 = arith.constant 42 : i32
+    %0 = memref.alloc() : memref<4xi32>
+    linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel"]} ins(%arg2, %arg3 : memref<4xi32>, memref<4xi32>) outs(%0 : memref<4xi32>) {
+    ^bb0(%arg5: i32, %arg6: i32, %arg7: i32):
+      %1 = arith.addi %arg5, %arg6 : i32
+      linalg.yield %1 : i32
+    }
+    memref.copy %0, %arg4 : memref<4xi32> to memref<4xi32>
+    memref.store %c42, %0[%c0] : memref<4xi32>
+    func.return
+  }
+}
+
+// CHECK: #[[$MAP0:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-LABEL:func @testCopyAfterLinalgMutated(
+// CHECK-SAME: %[[ARG0:.*]]: memref<4xi32>, %[[ARG1:.*]]: memref<4xi32>, %[[ARG2:.*]]: memref<4xi32>) {
+// CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<4xi32>
+// CHECK: linalg.generic {indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP0]]], iterator_types = ["parallel"]} ins(%[[ARG0]], %[[ARG1]] : memref<4xi32>, memref<4xi32>) outs(%[[ALLOC]] : memref<4xi32>) {
+// CHECK: ^bb0(%[[ARG3:.*]]: i32, %[[ARG4:.*]]: i32, %[[ARG5:.*]]: i32):
+// CHECK:   %[[R1:.*]] = arith.addi %[[ARG3]], %[[ARG4]] : i32
+// CHECK:   linalg.yield %[[R1]] : i32
+// CHECK: }
+// CHECK: memref.copy %[[ALLOC]], %[[ARG2]]
+// CHECK: return
+// CHECK: }

tensorflow/compiler/mlir/tools/kernel_gen/tests/embed_tf_framework.mlir

@@ -0,0 +1,100 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s -split-input-file -verify-diagnostics -embed-tf-framework |\
+// RUN: FileCheck %s
+
+// CHECK-LABEL: func @tf_entry(
+// CHECK-SAME:    [[CTX:%.*]]: !tf_framework.op_kernel_context,
+// CHECK-SAME:    [[SIZE_0:%.*]]: index,
+// CHECK-SAME:    [[SIZE_2:%.*]]: index) -> index attributes {tf_entry} {
+func.func @tf_entry(%size_0 : index , %size_2 : index) -> index
+    attributes {tf_entry} {
+  %buf = memref.alloc(%size_0, %size_2)[] : memref<?x10x?xf32>
+  memref.dealloc %buf : memref<?x10x?xf32>
+  func.return %size_0 : index
+}
+// CHECK-NEXT: [[BUF:%.*]] = tf_framework.alloc
+// CHECK-SAME:   ([[CTX]], [[SIZE_0]], [[SIZE_2]]) : memref<?x10x?xf32>
+
+// CHECK-NEXT: [[IS_VALID:%.*]] = tf_framework.is_valid_memref([[BUF]])
+// CHECK-SAME: memref<?x10x?xf32> -> i1
+
+// CHECK-NEXT: tf_framework.assert [[CTX]], [[IS_VALID]],
+// CHECK-SAME:   RESOURCE_EXHAUSTED, "failed to allocate memory"
+
+// CHECK-NEXT: tf_framework.dealloc([[CTX]], [[BUF]]) : memref<?x10x?xf32>
+// CHECK-NEXT: return [[SIZE_0]] : index
+
+// -----
+
+// CHECK-LABEL: func @non_tf_entry(
+// CHECK-SAME:    [[SIZE_0:%.*]]: index, [[SIZE_2:%.*]]: index) -> index
+func.func @non_tf_entry(%size_0 : index , %size_2 : index) -> index {
+  func.return %size_0 : index
+}
+
+// -----
+
+func.func @tf_entry_no_ctx(%size : index) attributes {tf_entry} {
+  // expected-error @+1 {{failed to legalize operation 'memref.alloc' that was explicitly marked illegal}}
+  %buf = memref.alloc()[%size] : memref<64xf32, affine_map<(d0)[s0] -> (d0 + s0)>>
+  memref.dealloc %buf : memref<64xf32, affine_map<(d0)[s0] -> (d0 + s0)>>
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL: func @assert(
+// CHECK-SAME: [[CTX:%.*]]: !tf_framework.op_kernel_context
+func.func @assert(%arg0: !tf_framework.op_kernel_context) attributes {tf_entry} {
+  %true = arith.constant true
+  cf.assert %true, "the one and only"
+  func.return
+}
+// CHECK:   [[TRUE:%.*]] = arith.constant true
+// CHECK-NEXT: tf_framework.assert [[CTX]], [[TRUE]], INVALID_ARGUMENT
+
+// -----
+
+// CHECK-LABEL: func @jit_execute
+// CHECK-SAME:  %[[CTX:.*]]: !tf_framework.op_kernel_context,
+// CHECK-SAME:  %[[F:.*]]: !tf_framework.jit_callable,
+// CHECK-SAME:  %[[ARG0:.*]]: tensor<2x?xf32>,
+// CHECK-SAME:  %[[ARG1:.*]]: tensor<2x?xf32>
+func.func @jit_execute(%ctx : !tf_framework.op_kernel_context,
+    %f : !tf_framework.jit_callable, %arg0 : tensor<2x?xf32>,
+    %arg1 : tensor<2x?xf32>) -> tensor<2x?xf32> attributes {tf_entry} {
+  // CHECK: %[[RES:.*]] = tf_framework.jit_execute
+  // CHECK-SAME: ctx(%[[CTX]]) %[[F]](%[[ARG0]], %[[ARG1]])
+  // CHECK: return %[[RES]]
+  %0 = tf_framework.jit_execute %f(%arg0, %arg1)
+      : tensor<2x?xf32>, tensor<2x?xf32> -> tensor<2x?xf32>
+  func.return %0 : tensor<2x?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @jit_compile_from_str
+// CHECK-SAME:  (%[[CTX:.*]]: !tf_framework.op_kernel_context)
+func.func @jit_compile_from_str(%ctx : !tf_framework.op_kernel_context)
+    -> !tf_framework.jit_callable attributes {tf_entry} {
+  // CHECK: %[[RES:.*]] = tf_framework.jit_compile_from_str %[[CTX]], "placeholder"
+  // CHECK: return %[[RES]]
+  %0 = tf_framework.jit_compile_from_str "placeholder" {
+      architectures = ["sm_123", "sm_456"], tileSizes = [1, 2, 3],
+      unrollFactors = [4], enableFtz = false,
+      index64Bit = false, cpuCodegen = false }
+  func.return %0 : !tf_framework.jit_callable
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/func_to_jit_invocations.mlir

@@ -0,0 +1,188 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s --split-input-file \
+// RUN:   --func-to-jit-invocation="tile-sizes=1,2,3 unroll-factors=3,2,1 \
+// RUN:       enable-ftz=false cpu-codegen=false" | \
+// RUN: FileCheck %s
+
+// RUN: kernel-gen-opt %s --split-input-file \
+// RUN:   --func-to-jit-invocation="tile-sizes=1,2,3 unroll-factors=3,2,1 \
+// RUN:       enable-ftz=false cpu-codegen=false \
+// RUN:       jit_i64_indexed_for_large_tensors=true" | \
+// RUN: FileCheck %s --check-prefix=CHECK-JFLT
+
+func.func @unary_tanh(%arg : tensor<?xf32>) -> tensor<?xf32> {
+  %0 = mhlo.tanh %arg : tensor<?xf32>
+  func.return %0 : tensor<?xf32>
+}
+
+// CHECK-LABEL: @unary_tanh
+// CHECK-SAME:  %[[ARG:.*]]: tensor<?xf32>
+// CHECK:       %[[CALLABLE:.*]] = tf_framework.jit_compile_from_str
+// CHECK-SAME:      "
+// CHECK-SAME:      module {
+// CHECK-SAME:        func @main(%[[ARG_JIT:.*]]: tensor<?xf32>) -> tensor<?xf32>
+// CHECK-SAME:          attributes {tf_entry}
+// CHECK-SAME:        {
+// CHECK-SAME:          %[[RES_JIT:.*]] = mhlo.tanh %[[ARG_JIT]]
+// CHECK-SAME:          return %[[RES_JIT]]
+// CHECK-SAME:        }
+// CHECK-SAME:      }
+// CHECK-SAME:      "
+// CHECK-SAME:      {
+// CHECK-SAME:        cpuCodegen = false
+// CHECK-SAME:        enableFtz = false
+// CHECK-SAME:        tileSizes = [1, 2, 3]
+// CHECK-SAME:        unrollFactors = [3, 2, 1]
+// CHECK-SAME:      }
+// CHECK:       %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG]])
+// CHECK:       return %[[RES]]
+
+// CHECK-JFLT-LABEL: @unary_tanh
+// CHECK-JFLT-SAME:  %[[ARG0:.*]]: tensor<?xf32>
+// CHECK-JFLT:       %[[SHAPE:.*]] = shape.shape_of %[[ARG0]]
+// CHECK-JFLT:       %[[NUM:.*]] = shape.num_elements %[[SHAPE]]
+// CHECK-JFLT:       %[[LIMIT:.*]] = arith.constant 2147483647
+// CHECK-JFLT:       %[[CMPI:.*]] = arith.cmpi sgt, %[[NUM]], %[[LIMIT]]
+// CHECK-JFLT:       %[[IF:.*]] = scf.if %[[CMPI]]
+// CHECK-JFLT:         %[[JIT:.*]] = tf_framework.jit_compile_from_str
+// CHECK-JFLT-SAME:        "module
+// CHECK-JFLT-SAME:        cpuCodegen = false
+// CHECK-JFLT-SAME:        enableFtz = false
+// CHECK-JFLT-SAME:        index64Bit = true
+// CHECK-JFLT-SAME:        tileSizes = [1, 2, 3]
+// CHECK-JFLT-SAME:        unrollFactors = [3, 2, 1]
+// CHECK-JFLT:         %[[JIT_0:.*]] = tf_framework.jit_execute %[[JIT]](%[[ARG0]])
+// CHECK-JFLT:         scf.yield %[[JIT_0]]
+// CHECK-JFLT:       else
+// CHECK-JFLT:         %[[VAL:.*]] = mhlo.tanh %[[ARG0]]
+// CHECK-JFLT:         scf.yield %[[VAL]]
+// CHECK-JFLT:       return %[[IF]]
+
+// -----
+
+func.func @binary_sub(%arg0 : tensor<*xf32>, %arg1 : tensor<*xf32>) -> tensor<*xf32> {
+  %0 = chlo.broadcast_subtract %arg0, %arg1 : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
+  func.return %0 : tensor<*xf32>
+}
+
+// CHECK-LABEL: @binary_sub
+// CHECK-SAME:  %[[ARG0:.*]]: tensor<*xf32>, %[[ARG1:.*]]: tensor<*xf32>
+// CHECK:       %[[CALLABLE:.*]] = tf_framework.jit_compile_from_str
+// CHECK-SAME:      "
+// CHECK-SAME:      module {
+// CHECK-SAME:        func @main(%[[ARG0_JIT:.*]]: tensor<*xf32>, %[[ARG1_JIT:.*]]: tensor<*xf32>) -> tensor<*xf32>
+// CHECK-SAME:          attributes {tf_entry}
+// CHECK-SAME:        {
+// CHECK-SAME:          %[[RES_JIT:.*]] = chlo.broadcast_subtract %[[ARG0_JIT]], %[[ARG1_JIT]]
+// CHECK-SAME:          return %[[RES_JIT]]
+// CHECK-SAME:        }
+// CHECK-SAME:      }
+// CHECK-SAME:      "
+// CHECK-SAME:      {
+// CHECK-SAME:        cpuCodegen = false
+// CHECK-SAME:        enableFtz = false
+// CHECK-SAME:        tileSizes = [1, 2, 3]
+// CHECK-SAME:        unrollFactors = [3, 2, 1]
+// CHECK-SAME:      }
+// CHECK:       %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG0]], %[[ARG1]])
+// CHECK:       return %[[RES]]
+
+// CHECK-JFLT-LABEL: @binary_sub
+// CHECK-JFLT:  %[[ARG0:.*]]: tensor<*xf32>, %[[ARG1:.*]]: tensor<*xf32>
+// CHECK-JFLT:  %[[SHAPE1:.*]] = shape.shape_of %[[ARG0]] : tensor<*xf32> -> tensor<?xindex>
+// CHECK-JFLT:  %[[ELEMENTCOUNT1:.*]] = shape.num_elements %[[SHAPE1]] : tensor<?xindex> -> index
+// CHECK-JFLT:  %[[LIMIT:.*]] = arith.constant 2147483647
+// CHECK-JFLT:  %[[COMP1:.*]] = arith.cmpi sgt, %[[ELEMENTCOUNT1]], %[[LIMIT]] : index
+// CHECK-JFLT:  %[[SHAPE2:.*]] = shape.shape_of %[[ARG1]] : tensor<*xf32> -> tensor<?xindex>
+// CHECK-JFLT:  %[[ELEMENTCOUNT2:.*]] = shape.num_elements %[[SHAPE2]] : tensor<?xindex> -> index
+// CHECK-JFLT:  %[[COMP2:.*]]  = arith.cmpi sgt, %[[ELEMENTCOUNT2]], %[[LIMIT]] : index
+// CHECK-JFLT:  %[[COMPRES:.*]] = arith.ori %[[COMP1]], %[[COMP2]] : i1
+// CHECK-JFLT:  %[[IFRES:.*]] = scf.if %[[COMPRES]] -> (tensor<*xf32>) {
+// CHECK-JFLT:       %[[CALLABLE:.*]] = tf_framework.jit_compile_from_str
+// CHECK-JFLT-SAME:      "
+// CHECK-JFLT-SAME:      module {
+// CHECK-JFLT-SAME:        func @main(%[[ARG0_JIT:.*]]: tensor<*xf32>, %[[ARG1_JIT:.*]]: tensor<*xf32>) -> tensor<*xf32>
+// CHECK-JFLT-SAME:          attributes {tf_entry}
+// CHECK-JFLT-SAME:        {
+// CHECK-JFLT-SAME:          %[[RES_JIT:.*]] = chlo.broadcast_subtract %[[ARG0_JIT]], %[[ARG1_JIT]]
+// CHECK-JFLT-SAME:          return %[[RES_JIT]]
+// CHECK-JFLT-SAME:        }
+// CHECK-JFLT-SAME:      }
+// CHECK-JFLT-SAME:      "
+// CHECK-JFLT-SAME:      {
+// CHECK-JFLT-SAME:        cpuCodegen = false
+// CHECK-JFLT-SAME:        enableFtz = false
+// CHECK-JFLT-SAME:        tileSizes = [1, 2, 3]
+// CHECK-JFLT-SAME:        unrollFactors = [3, 2, 1]
+// CHECK-JFLT-SAME:      }
+// CHECK-JFLT:       %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG0]], %[[ARG1]])
+// CHECK-JFLT:       scf.yield %[[RES]] : tensor<*xf32>
+// CHECK-JFLT:     } else {
+// CHECK-JFLT:       %[[RES2:.*]] = chlo.broadcast_subtract %[[ARG0]], %[[ARG1]] : (tensor<*xf32>, tensor<*xf32>) -> tensor<*xf32>
+// CHECK-JFLT:       scf.yield %[[RES2]] : tensor<*xf32>
+// CHECK-JFLT:     }
+// CHECK-JFLT:     return %[[IFRES]]
+
+// -----
+
+func.func @reciprocal(%arg0: tensor<*xf32>)
+    -> tensor<*xf32> attributes {tf_entry, llvm.emit_c_interface} {
+  %0 = mhlo.constant dense<1.0> : tensor<f32>
+  %1 = chlo.broadcast_divide %0, %arg0 : (tensor<f32>, tensor<*xf32>) -> tensor<*xf32>
+  func.return %1 : tensor<*xf32>
+}
+
+// CHECK-LABEL: @reciprocal
+// CHECK-SAME:  %[[ARG:.*]]: tensor<*xf32>
+// CHECK:       %[[CALLABLE:.*]] = tf_framework.jit_compile_from_str
+// CHECK-SAME:    module {
+// CHECK-SAME:     func.func @main(%[[ARG0_JIT:.*]]: tensor<*xf32>) -> tensor<*xf32>
+// CHECK-SAME:       %[[CST:.*]] = mhlo.constant dense<1.000000e+00> : tensor<f32>
+// CHECK-SAME:       %[[RES_JIT:.*]] = chlo.broadcast_divide %[[CST]], %[[ARG0_JIT]] : (tensor<f32>, tensor<*xf32>) -> tensor<*xf32>
+// CHECK=SAME:       return %[[RES_JIT]] : tensor<*xf32>
+// CHECK-SAME:     }
+// CHECK-SAME:   }
+// CHECK-SAME:   {
+// CHECK-SAME:     cpuCodegen = false,
+// CHECK-SAME:     enableFtz = false,
+// CHECK-SAME:     index64Bit = false,
+// CHECK-SAME:     tileSizes = [1, 2, 3],
+// CHECK-SAME:     unrollFactors = [3, 2, 1]
+// CHECK-SAME:   }
+// CHECK:      %[[RES:.*]] = tf_framework.jit_execute %[[CALLABLE]](%[[ARG]]) : tensor<*xf32> -> tensor<*xf32>
+// CHECK:      return %[[RES]] : tensor<*xf32>
+
+// CHECK-JFLT-LABEL: @reciprocal
+// CHECK-JFLT-SAME:  %[[ARG0:.*]]: tensor<*xf32>
+// CHECK-JFLT:       %[[SHAPE:.*]] = shape.shape_of %[[ARG0]]
+// CHECK-JFLT:       %[[NUM:.*]] = shape.num_elements %[[SHAPE]]
+// CHECK-JFLT:       %[[LIMIT:.*]] = arith.constant 2147483647
+// CHECK-JFLT:       %[[CMPI:.*]] = arith.cmpi sgt, %[[NUM]], %[[LIMIT]]
+// CHECK-JFLT:       %[[IF:.*]] = scf.if %[[CMPI]]
+// CHECK-JFLT:         %[[JIT:.*]] = tf_framework.jit_compile_from_str
+// CHECK-JFLT-SAME:        "module
+// CHECK-JFLT-SAME:        cpuCodegen = false
+// CHECK-JFLT-SAME:        enableFtz = false
+// CHECK-JFLT-SAME:        index64Bit = true
+// CHECK-JFLT-SAME:        tileSizes = [1, 2, 3]
+// CHECK-JFLT-SAME:        unrollFactors = [3, 2, 1]
+// CHECK-JFLT:         %[[JIT_0:.*]] = tf_framework.jit_execute %[[JIT]](%[[ARG0]])
+// CHECK-JFLT:         scf.yield %[[JIT_0]]
+// CHECK-JFLT:       else
+// CHECK-JFLT:         %[[CST:.*]] = mhlo.constant dense<1.000000e+00> : tensor<f32>
+// CHECK-JFLT:         %[[VAL:.*]] = chlo.broadcast_divide %[[CST]], %[[ARG0]] : (tensor<f32>, tensor<*xf32>) -> tensor<*xf32>
+// CHECK-JFLT:         scf.yield %[[VAL]]
+// CHECK-JFLT:       return %[[IF]]

tensorflow/compiler/mlir/tools/kernel_gen/tests/hlo_to_kernel/BUILD

@@ -0,0 +1,31 @@
+load("//tensorflow/compiler/mlir:glob_lit_test.bzl", "glob_lit_tests")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+glob_lit_tests(
+    name = "all_tests",
+    data = [":test_utilities"],
+    default_tags = [
+        # We need access to the CUDA SDK.
+        "gpu",
+        "no_rocm",
+        "nomsan",
+        "notsan",
+    ],
+    driver = "//tensorflow/compiler/mlir:run_lit.sh",
+    hermetic_cuda_data_dir = "%S/../../../../../../../../cuda_nvcc",
+    test_file_exts = ["mlir"],
+)
+
+# Bundle together all of the test utilities that are used by tests.
+filegroup(
+    name = "test_utilities",
+    testonly = True,
+    data = [
+        "//tensorflow/compiler/mlir/tools/kernel_gen:hlo_to_kernel",
+        "@llvm-project//mlir:run_lit.sh",
+    ],
+)

tensorflow/compiler/mlir/tools/kernel_gen/tests/hlo_to_kernel/add_v2.mlir

@@ -0,0 +1,146 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: hlo_to_kernel --input=%s --output=%t --unroll_factors=4 --tile_sizes=256 --arch=sm_70
+
+func.func @AddV2(%arg0: tensor<*xf32>, %arg1: tensor<*xf32>) -> tensor<*xf32> attributes {llvm.emit_c_interface, tf_entry} {
+  %0 = shape.const_shape [1, 1, 1, 1, 1] : tensor<5xindex>
+  %c5 = arith.constant 5 : index
+  %1 = shape.const_shape [1, 1, 1, 1] : tensor<4xindex>
+  %c4 = arith.constant 4 : index
+  %2 = shape.const_shape [1, 1, 1] : tensor<3xindex>
+  %c3 = arith.constant 3 : index
+  %3 = shape.const_shape [1, 1] : tensor<2xindex>
+  %c2 = arith.constant 2 : index
+  %4 = shape.const_shape [1] : tensor<1xindex>
+  %c1 = arith.constant 1 : index
+  %5 = shape.shape_of %arg0 : tensor<*xf32> -> tensor<?xindex>
+  %6 = shape.shape_of %arg1 : tensor<*xf32> -> tensor<?xindex>
+  %7 = shape.num_elements %5 : tensor<?xindex> -> index
+  %8 = arith.cmpi eq, %7, %c1 : index
+  %c_empty = arith.constant dense<> : tensor<0xindex>
+  %9 = scf.if %8 -> (tensor<*xf32>) {
+    %14 = shape.num_elements %6 : tensor<?xindex> -> index
+    %from_elements = tensor.from_elements %14 : tensor<1xindex>
+    %15 = tensor.reshape %arg0(%c_empty) : (tensor<*xf32>, tensor<0xindex>) -> tensor<f32>
+    %16 = tensor.reshape %arg1(%from_elements) : (tensor<*xf32>, tensor<1xindex>) -> tensor<?xf32>
+    %17 = chlo.broadcast_add %15, %16 : (tensor<f32>, tensor<?xf32>) -> tensor<?xf32>
+    %cast = tensor.cast %17 : tensor<?xf32> to tensor<*xf32>
+    scf.yield %cast : tensor<*xf32>
+  } else {
+    %14 = shape.num_elements %6 : tensor<?xindex> -> index
+    %15 = arith.cmpi eq, %14, %c1 : index
+    %16 = scf.if %15 -> (tensor<*xf32>) {
+      %17 = shape.num_elements %5 : tensor<?xindex> -> index
+      %from_elements = tensor.from_elements %17 : tensor<1xindex>
+      %18 = tensor.reshape %arg0(%from_elements) : (tensor<*xf32>, tensor<1xindex>) -> tensor<?xf32>
+      %19 = tensor.reshape %arg1(%c_empty) : (tensor<*xf32>, tensor<0xindex>) -> tensor<f32>
+      %20 = chlo.broadcast_add %18, %19 : (tensor<?xf32>, tensor<f32>) -> tensor<?xf32>
+      %cast = tensor.cast %20 : tensor<?xf32> to tensor<*xf32>
+      scf.yield %cast : tensor<*xf32>
+    } else {
+      %17 = shape.shape_eq %5, %6 : tensor<?xindex>, tensor<?xindex>
+      %18 = scf.if %17 -> (tensor<*xf32>) {
+        %19 = shape.any %5, %6 : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
+        %20 = shape.num_elements %19 : tensor<?xindex> -> index
+        %from_elements = tensor.from_elements %20 : tensor<1xindex>
+        %21 = tensor.reshape %arg0(%from_elements) : (tensor<*xf32>, tensor<1xindex>) -> tensor<?xf32>
+        %22 = tensor.reshape %arg1(%from_elements) : (tensor<*xf32>, tensor<1xindex>) -> tensor<?xf32>
+        %23 = chlo.broadcast_add %21, %22 : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
+        %cast = tensor.cast %23 : tensor<?xf32> to tensor<*xf32>
+        scf.yield %cast : tensor<*xf32>
+      } else {
+        %19:2 = mhlo.minimum_broadcast_shapes %5, %6 : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>, tensor<?xindex>
+        %20 = shape.rank %19#0 : tensor<?xindex> -> index
+        %21 = shape.rank %19#1 : tensor<?xindex> -> index
+        %22 = arith.cmpi sgt, %20, %21 : index
+        %23 = arith.select %22, %20, %21 : index
+        %24 = arith.cmpi ule, %23, %c1 : index
+        %25 = scf.if %24 -> (tensor<*xf32>) {
+          %26 = shape.broadcast %19#0, %4 : tensor<?xindex>, tensor<1xindex> -> tensor<?xindex>
+          %cast = tensor.cast %26 : tensor<?xindex> to tensor<1xindex>
+          %27 = tensor.reshape %arg0(%cast) : (tensor<*xf32>, tensor<1xindex>) -> tensor<?xf32>
+          %28 = shape.broadcast %19#1, %4 : tensor<?xindex>, tensor<1xindex> -> tensor<?xindex>
+          %cast_0 = tensor.cast %28 : tensor<?xindex> to tensor<1xindex>
+          %29 = tensor.reshape %arg1(%cast_0) : (tensor<*xf32>, tensor<1xindex>) -> tensor<?xf32>
+          %30 = chlo.broadcast_add %27, %29 : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
+          %cast_1 = tensor.cast %30 : tensor<?xf32> to tensor<*xf32>
+          scf.yield %cast_1 : tensor<*xf32>
+        } else {
+          %26 = arith.cmpi ule, %23, %c2 : index
+          %27 = scf.if %26 -> (tensor<*xf32>) {
+            %28 = shape.broadcast %19#0, %3 : tensor<?xindex>, tensor<2xindex> -> tensor<?xindex>
+            %cast = tensor.cast %28 : tensor<?xindex> to tensor<2xindex>
+            %29 = tensor.reshape %arg0(%cast) : (tensor<*xf32>, tensor<2xindex>) -> tensor<?x?xf32>
+            %30 = shape.broadcast %19#1, %3 : tensor<?xindex>, tensor<2xindex> -> tensor<?xindex>
+            %cast_0 = tensor.cast %30 : tensor<?xindex> to tensor<2xindex>
+            %31 = tensor.reshape %arg1(%cast_0) : (tensor<*xf32>, tensor<2xindex>) -> tensor<?x?xf32>
+            %32 = chlo.broadcast_add %29, %31 : (tensor<?x?xf32>, tensor<?x?xf32>) -> tensor<?x?xf32>
+            %cast_1 = tensor.cast %32 : tensor<?x?xf32> to tensor<*xf32>
+            scf.yield %cast_1 : tensor<*xf32>
+          } else {
+            %28 = arith.cmpi ule, %23, %c3 : index
+            %29 = scf.if %28 -> (tensor<*xf32>) {
+              %30 = shape.broadcast %19#0, %2 : tensor<?xindex>, tensor<3xindex> -> tensor<?xindex>
+              %cast = tensor.cast %30 : tensor<?xindex> to tensor<3xindex>
+              %31 = tensor.reshape %arg0(%cast) : (tensor<*xf32>, tensor<3xindex>) -> tensor<?x?x?xf32>
+              %32 = shape.broadcast %19#1, %2 : tensor<?xindex>, tensor<3xindex> -> tensor<?xindex>
+              %cast_0 = tensor.cast %32 : tensor<?xindex> to tensor<3xindex>
+              %33 = tensor.reshape %arg1(%cast_0) : (tensor<*xf32>, tensor<3xindex>) -> tensor<?x?x?xf32>
+              %34 = chlo.broadcast_add %31, %33 : (tensor<?x?x?xf32>, tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
+              %cast_1 = tensor.cast %34 : tensor<?x?x?xf32> to tensor<*xf32>
+              scf.yield %cast_1 : tensor<*xf32>
+            } else {
+              %30 = arith.cmpi ule, %23, %c4 : index
+              %31 = scf.if %30 -> (tensor<*xf32>) {
+                %32 = shape.broadcast %19#0, %1 : tensor<?xindex>, tensor<4xindex> -> tensor<?xindex>
+                %cast = tensor.cast %32 : tensor<?xindex> to tensor<4xindex>
+                %33 = tensor.reshape %arg0(%cast) : (tensor<*xf32>, tensor<4xindex>) -> tensor<?x?x?x?xf32>
+                %34 = shape.broadcast %19#1, %1 : tensor<?xindex>, tensor<4xindex> -> tensor<?xindex>
+                %cast_0 = tensor.cast %34 : tensor<?xindex> to tensor<4xindex>
+                %35 = tensor.reshape %arg1(%cast_0) : (tensor<*xf32>, tensor<4xindex>) -> tensor<?x?x?x?xf32>
+                %36 = chlo.broadcast_add %33, %35 : (tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>
+                %cast_1 = tensor.cast %36 : tensor<?x?x?x?xf32> to tensor<*xf32>
+                scf.yield %cast_1 : tensor<*xf32>
+              } else {
+                %32 = arith.cmpi ule, %23, %c5 : index
+                cf.assert %32, "Input for dynamic binary or n-ary op lowering was of a rank greater than 5"
+                %33 = shape.broadcast %19#0, %0 : tensor<?xindex>, tensor<5xindex> -> tensor<?xindex>
+                %cast = tensor.cast %33 : tensor<?xindex> to tensor<5xindex>
+                %34 = tensor.reshape %arg0(%cast) : (tensor<*xf32>, tensor<5xindex>) -> tensor<?x?x?x?x?xf32>
+                %35 = shape.broadcast %19#1, %0 : tensor<?xindex>, tensor<5xindex> -> tensor<?xindex>
+                %cast_0 = tensor.cast %35 : tensor<?xindex> to tensor<5xindex>
+                %36 = tensor.reshape %arg1(%cast_0) : (tensor<*xf32>, tensor<5xindex>) -> tensor<?x?x?x?x?xf32>
+                %37 = chlo.broadcast_add %34, %36 : (tensor<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>) -> tensor<?x?x?x?x?xf32>
+                %cast_1 = tensor.cast %37 : tensor<?x?x?x?x?xf32> to tensor<*xf32>
+                scf.yield %cast_1 : tensor<*xf32>
+              }
+              scf.yield %31 : tensor<*xf32>
+            }
+            scf.yield %29 : tensor<*xf32>
+          }
+          scf.yield %27 : tensor<*xf32>
+        }
+        scf.yield %25 : tensor<*xf32>
+      }
+      scf.yield %18 : tensor<*xf32>
+    }
+    scf.yield %16 : tensor<*xf32>
+  }
+  %10 = shape.shape_of %arg0 : tensor<*xf32> -> tensor<?xindex>
+  %11 = shape.shape_of %arg1 : tensor<*xf32> -> tensor<?xindex>
+  %12 = shape.broadcast %10, %11 : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
+  %13 = tensor.reshape %9(%12) : (tensor<*xf32>, tensor<?xindex>) -> tensor<*xf32>
+  return %13 : tensor<*xf32>
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/hlo_to_kernel/add_v2_unsigned.mlir

@@ -0,0 +1,146 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: hlo_to_kernel --input=%s --output=%t --unroll_factors=4 --tile_sizes=256 --arch=sm_70
+
+func.func @AddV2(%arg0: tensor<*xui32>, %arg1: tensor<*xui32>) -> tensor<*xui32> attributes {llvm.emit_c_interface, tf_entry} {
+  %0 = shape.const_shape [1, 1, 1, 1, 1] : tensor<5xindex>
+  %c5 = arith.constant 5 : index
+  %1 = shape.const_shape [1, 1, 1, 1] : tensor<4xindex>
+  %c4 = arith.constant 4 : index
+  %2 = shape.const_shape [1, 1, 1] : tensor<3xindex>
+  %c3 = arith.constant 3 : index
+  %3 = shape.const_shape [1, 1] : tensor<2xindex>
+  %c2 = arith.constant 2 : index
+  %4 = shape.const_shape [1] : tensor<1xindex>
+  %c1 = arith.constant 1 : index
+  %5 = shape.shape_of %arg0 : tensor<*xui32> -> tensor<?xindex>
+  %6 = shape.shape_of %arg1 : tensor<*xui32> -> tensor<?xindex>
+  %7 = shape.num_elements %5 : tensor<?xindex> -> index
+  %8 = arith.cmpi eq, %7, %c1 : index
+  %c_empty = arith.constant dense<> : tensor<0xindex>
+  %9 = scf.if %8 -> (tensor<*xui32>) {
+    %14 = shape.num_elements %6 : tensor<?xindex> -> index
+    %from_elements = tensor.from_elements %14 : tensor<1xindex>
+    %15 = tensor.reshape %arg0(%c_empty) : (tensor<*xui32>, tensor<0xindex>) -> tensor<ui32>
+    %16 = tensor.reshape %arg1(%from_elements) : (tensor<*xui32>, tensor<1xindex>) -> tensor<?xui32>
+    %17 = chlo.broadcast_add %15, %16 : (tensor<ui32>, tensor<?xui32>) -> tensor<?xui32>
+    %cast = tensor.cast %17 : tensor<?xui32> to tensor<*xui32>
+    scf.yield %cast : tensor<*xui32>
+  } else {
+    %14 = shape.num_elements %6 : tensor<?xindex> -> index
+    %15 = arith.cmpi eq, %14, %c1 : index
+    %16 = scf.if %15 -> (tensor<*xui32>) {
+      %17 = shape.num_elements %5 : tensor<?xindex> -> index
+      %from_elements = tensor.from_elements %17 : tensor<1xindex>
+      %18 = tensor.reshape %arg0(%from_elements) : (tensor<*xui32>, tensor<1xindex>) -> tensor<?xui32>
+      %19 = tensor.reshape %arg1(%c_empty) : (tensor<*xui32>, tensor<0xindex>) -> tensor<ui32>
+      %20 = chlo.broadcast_add %18, %19 : (tensor<?xui32>, tensor<ui32>) -> tensor<?xui32>
+      %cast = tensor.cast %20 : tensor<?xui32> to tensor<*xui32>
+      scf.yield %cast : tensor<*xui32>
+    } else {
+      %17 = shape.shape_eq %5, %6 : tensor<?xindex>, tensor<?xindex>
+      %18 = scf.if %17 -> (tensor<*xui32>) {
+        %19 = shape.any %5, %6 : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
+        %20 = shape.num_elements %19 : tensor<?xindex> -> index
+        %from_elements = tensor.from_elements %20 : tensor<1xindex>
+        %21 = tensor.reshape %arg0(%from_elements) : (tensor<*xui32>, tensor<1xindex>) -> tensor<?xui32>
+        %22 = tensor.reshape %arg1(%from_elements) : (tensor<*xui32>, tensor<1xindex>) -> tensor<?xui32>
+        %23 = chlo.broadcast_add %21, %22 : (tensor<?xui32>, tensor<?xui32>) -> tensor<?xui32>
+        %cast = tensor.cast %23 : tensor<?xui32> to tensor<*xui32>
+        scf.yield %cast : tensor<*xui32>
+      } else {
+        %19:2 = mhlo.minimum_broadcast_shapes %5, %6 : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>, tensor<?xindex>
+        %20 = shape.rank %19#0 : tensor<?xindex> -> index
+        %21 = shape.rank %19#1 : tensor<?xindex> -> index
+        %22 = arith.cmpi sgt, %20, %21 : index
+        %23 = arith.select %22, %20, %21 : index
+        %24 = arith.cmpi ule, %23, %c1 : index
+        %25 = scf.if %24 -> (tensor<*xui32>) {
+          %26 = shape.broadcast %19#0, %4 : tensor<?xindex>, tensor<1xindex> -> tensor<?xindex>
+          %cast = tensor.cast %26 : tensor<?xindex> to tensor<1xindex>
+          %27 = tensor.reshape %arg0(%cast) : (tensor<*xui32>, tensor<1xindex>) -> tensor<?xui32>
+          %28 = shape.broadcast %19#1, %4 : tensor<?xindex>, tensor<1xindex> -> tensor<?xindex>
+          %cast_0 = tensor.cast %28 : tensor<?xindex> to tensor<1xindex>
+          %29 = tensor.reshape %arg1(%cast_0) : (tensor<*xui32>, tensor<1xindex>) -> tensor<?xui32>
+          %30 = chlo.broadcast_add %27, %29 : (tensor<?xui32>, tensor<?xui32>) -> tensor<?xui32>
+          %cast_1 = tensor.cast %30 : tensor<?xui32> to tensor<*xui32>
+          scf.yield %cast_1 : tensor<*xui32>
+        } else {
+          %26 = arith.cmpi ule, %23, %c2 : index
+          %27 = scf.if %26 -> (tensor<*xui32>) {
+            %28 = shape.broadcast %19#0, %3 : tensor<?xindex>, tensor<2xindex> -> tensor<?xindex>
+            %cast = tensor.cast %28 : tensor<?xindex> to tensor<2xindex>
+            %29 = tensor.reshape %arg0(%cast) : (tensor<*xui32>, tensor<2xindex>) -> tensor<?x?xui32>
+            %30 = shape.broadcast %19#1, %3 : tensor<?xindex>, tensor<2xindex> -> tensor<?xindex>
+            %cast_0 = tensor.cast %30 : tensor<?xindex> to tensor<2xindex>
+            %31 = tensor.reshape %arg1(%cast_0) : (tensor<*xui32>, tensor<2xindex>) -> tensor<?x?xui32>
+            %32 = chlo.broadcast_add %29, %31 : (tensor<?x?xui32>, tensor<?x?xui32>) -> tensor<?x?xui32>
+            %cast_1 = tensor.cast %32 : tensor<?x?xui32> to tensor<*xui32>
+            scf.yield %cast_1 : tensor<*xui32>
+          } else {
+            %28 = arith.cmpi ule, %23, %c3 : index
+            %29 = scf.if %28 -> (tensor<*xui32>) {
+              %30 = shape.broadcast %19#0, %2 : tensor<?xindex>, tensor<3xindex> -> tensor<?xindex>
+              %cast = tensor.cast %30 : tensor<?xindex> to tensor<3xindex>
+              %31 = tensor.reshape %arg0(%cast) : (tensor<*xui32>, tensor<3xindex>) -> tensor<?x?x?xui32>
+              %32 = shape.broadcast %19#1, %2 : tensor<?xindex>, tensor<3xindex> -> tensor<?xindex>
+              %cast_0 = tensor.cast %32 : tensor<?xindex> to tensor<3xindex>
+              %33 = tensor.reshape %arg1(%cast_0) : (tensor<*xui32>, tensor<3xindex>) -> tensor<?x?x?xui32>
+              %34 = chlo.broadcast_add %31, %33 : (tensor<?x?x?xui32>, tensor<?x?x?xui32>) -> tensor<?x?x?xui32>
+              %cast_1 = tensor.cast %34 : tensor<?x?x?xui32> to tensor<*xui32>
+              scf.yield %cast_1 : tensor<*xui32>
+            } else {
+              %30 = arith.cmpi ule, %23, %c4 : index
+              %31 = scf.if %30 -> (tensor<*xui32>) {
+                %32 = shape.broadcast %19#0, %1 : tensor<?xindex>, tensor<4xindex> -> tensor<?xindex>
+                %cast = tensor.cast %32 : tensor<?xindex> to tensor<4xindex>
+                %33 = tensor.reshape %arg0(%cast) : (tensor<*xui32>, tensor<4xindex>) -> tensor<?x?x?x?xui32>
+                %34 = shape.broadcast %19#1, %1 : tensor<?xindex>, tensor<4xindex> -> tensor<?xindex>
+                %cast_0 = tensor.cast %34 : tensor<?xindex> to tensor<4xindex>
+                %35 = tensor.reshape %arg1(%cast_0) : (tensor<*xui32>, tensor<4xindex>) -> tensor<?x?x?x?xui32>
+                %36 = chlo.broadcast_add %33, %35 : (tensor<?x?x?x?xui32>, tensor<?x?x?x?xui32>) -> tensor<?x?x?x?xui32>
+                %cast_1 = tensor.cast %36 : tensor<?x?x?x?xui32> to tensor<*xui32>
+                scf.yield %cast_1 : tensor<*xui32>
+              } else {
+                %32 = arith.cmpi ule, %23, %c5 : index
+                cf.assert %32, "Input for dynamic binary or n-ary op lowering was of a rank greater than 5"
+                %33 = shape.broadcast %19#0, %0 : tensor<?xindex>, tensor<5xindex> -> tensor<?xindex>
+                %cast = tensor.cast %33 : tensor<?xindex> to tensor<5xindex>
+                %34 = tensor.reshape %arg0(%cast) : (tensor<*xui32>, tensor<5xindex>) -> tensor<?x?x?x?x?xui32>
+                %35 = shape.broadcast %19#1, %0 : tensor<?xindex>, tensor<5xindex> -> tensor<?xindex>
+                %cast_0 = tensor.cast %35 : tensor<?xindex> to tensor<5xindex>
+                %36 = tensor.reshape %arg1(%cast_0) : (tensor<*xui32>, tensor<5xindex>) -> tensor<?x?x?x?x?xui32>
+                %37 = chlo.broadcast_add %34, %36 : (tensor<?x?x?x?x?xui32>, tensor<?x?x?x?x?xui32>) -> tensor<?x?x?x?x?xui32>
+                %cast_1 = tensor.cast %37 : tensor<?x?x?x?x?xui32> to tensor<*xui32>
+                scf.yield %cast_1 : tensor<*xui32>
+              }
+              scf.yield %31 : tensor<*xui32>
+            }
+            scf.yield %29 : tensor<*xui32>
+          }
+          scf.yield %27 : tensor<*xui32>
+        }
+        scf.yield %25 : tensor<*xui32>
+      }
+      scf.yield %18 : tensor<*xui32>
+    }
+    scf.yield %16 : tensor<*xui32>
+  }
+  %10 = shape.shape_of %arg0 : tensor<*xui32> -> tensor<?xindex>
+  %11 = shape.shape_of %arg1 : tensor<*xui32> -> tensor<?xindex>
+  %12 = shape.broadcast %10, %11 : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
+  %13 = tensor.reshape %9(%12) : (tensor<*xui32>, tensor<?xindex>) -> tensor<*xui32>
+  return %13 : tensor<*xui32>
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/hlo_to_kernel/minimum.mlir

@@ -0,0 +1,146 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: hlo_to_kernel --input=%s --output=%t --unroll_factors=4 --tile_sizes=256 --arch=sm_70
+
+func.func @Minimum_GPU_DT_UINT32_DT_UINT32(%arg0: tensor<*xui32>, %arg1: tensor<*xui32>) -> tensor<*xui32> attributes {llvm.emit_c_interface, tf_entry} {
+  %0 = shape.const_shape [1, 1, 1, 1, 1] : tensor<5xindex>
+  %c5 = arith.constant 5 : index
+  %1 = shape.const_shape [1, 1, 1, 1] : tensor<4xindex>
+  %c4 = arith.constant 4 : index
+  %2 = shape.const_shape [1, 1, 1] : tensor<3xindex>
+  %c3 = arith.constant 3 : index
+  %3 = shape.const_shape [1, 1] : tensor<2xindex>
+  %c2 = arith.constant 2 : index
+  %4 = shape.const_shape [1] : tensor<1xindex>
+  %c1 = arith.constant 1 : index
+  %5 = shape.shape_of %arg0 : tensor<*xui32> -> tensor<?xindex>
+  %6 = shape.shape_of %arg1 : tensor<*xui32> -> tensor<?xindex>
+  %7 = shape.num_elements %5 : tensor<?xindex> -> index
+  %8 = arith.cmpi eq, %7, %c1 : index
+  %c_empty = arith.constant dense<> : tensor<0xindex>
+  %9 = scf.if %8 -> (tensor<*xui32>) {
+    %14 = shape.num_elements %6 : tensor<?xindex> -> index
+    %from_elements = tensor.from_elements %14 : tensor<1xindex>
+    %15 = tensor.reshape %arg0(%c_empty) : (tensor<*xui32>, tensor<0xindex>) -> tensor<ui32>
+    %16 = tensor.reshape %arg1(%from_elements) : (tensor<*xui32>, tensor<1xindex>) -> tensor<?xui32>
+    %17 = chlo.broadcast_minimum %15, %16 : (tensor<ui32>, tensor<?xui32>) -> tensor<?xui32>
+    %cast = tensor.cast %17 : tensor<?xui32> to tensor<*xui32>
+    scf.yield %cast : tensor<*xui32>
+  } else {
+    %14 = shape.num_elements %6 : tensor<?xindex> -> index
+    %15 = arith.cmpi eq, %14, %c1 : index
+    %16 = scf.if %15 -> (tensor<*xui32>) {
+      %17 = shape.num_elements %5 : tensor<?xindex> -> index
+      %from_elements = tensor.from_elements %17 : tensor<1xindex>
+      %18 = tensor.reshape %arg0(%from_elements) : (tensor<*xui32>, tensor<1xindex>) -> tensor<?xui32>
+      %19 = tensor.reshape %arg1(%c_empty) : (tensor<*xui32>, tensor<0xindex>) -> tensor<ui32>
+      %20 = chlo.broadcast_minimum %18, %19 : (tensor<?xui32>, tensor<ui32>) -> tensor<?xui32>
+      %cast = tensor.cast %20 : tensor<?xui32> to tensor<*xui32>
+      scf.yield %cast : tensor<*xui32>
+    } else {
+      %17 = shape.shape_eq %5, %6 : tensor<?xindex>, tensor<?xindex>
+      %18 = scf.if %17 -> (tensor<*xui32>) {
+        %19 = shape.any %5, %6 : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
+        %20 = shape.num_elements %19 : tensor<?xindex> -> index
+        %from_elements = tensor.from_elements %20 : tensor<1xindex>
+        %21 = tensor.reshape %arg0(%from_elements) : (tensor<*xui32>, tensor<1xindex>) -> tensor<?xui32>
+        %22 = tensor.reshape %arg1(%from_elements) : (tensor<*xui32>, tensor<1xindex>) -> tensor<?xui32>
+        %23 = chlo.broadcast_minimum %21, %22 : (tensor<?xui32>, tensor<?xui32>) -> tensor<?xui32>
+        %cast = tensor.cast %23 : tensor<?xui32> to tensor<*xui32>
+        scf.yield %cast : tensor<*xui32>
+      } else {
+        %19:2 = mhlo.minimum_broadcast_shapes %5, %6 : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>, tensor<?xindex>
+        %20 = shape.rank %19#0 : tensor<?xindex> -> index
+        %21 = shape.rank %19#1 : tensor<?xindex> -> index
+        %22 = arith.cmpi sgt, %20, %21 : index
+        %23 = arith.select %22, %20, %21 : index
+        %24 = arith.cmpi ule, %23, %c1 : index
+        %25 = scf.if %24 -> (tensor<*xui32>) {
+          %26 = shape.broadcast %19#0, %4 : tensor<?xindex>, tensor<1xindex> -> tensor<?xindex>
+          %cast = tensor.cast %26 : tensor<?xindex> to tensor<1xindex>
+          %27 = tensor.reshape %arg0(%cast) : (tensor<*xui32>, tensor<1xindex>) -> tensor<?xui32>
+          %28 = shape.broadcast %19#1, %4 : tensor<?xindex>, tensor<1xindex> -> tensor<?xindex>
+          %cast_0 = tensor.cast %28 : tensor<?xindex> to tensor<1xindex>
+          %29 = tensor.reshape %arg1(%cast_0) : (tensor<*xui32>, tensor<1xindex>) -> tensor<?xui32>
+          %30 = chlo.broadcast_minimum %27, %29 : (tensor<?xui32>, tensor<?xui32>) -> tensor<?xui32>
+          %cast_1 = tensor.cast %30 : tensor<?xui32> to tensor<*xui32>
+          scf.yield %cast_1 : tensor<*xui32>
+        } else {
+          %26 = arith.cmpi ule, %23, %c2 : index
+          %27 = scf.if %26 -> (tensor<*xui32>) {
+            %28 = shape.broadcast %19#0, %3 : tensor<?xindex>, tensor<2xindex> -> tensor<?xindex>
+            %cast = tensor.cast %28 : tensor<?xindex> to tensor<2xindex>
+            %29 = tensor.reshape %arg0(%cast) : (tensor<*xui32>, tensor<2xindex>) -> tensor<?x?xui32>
+            %30 = shape.broadcast %19#1, %3 : tensor<?xindex>, tensor<2xindex> -> tensor<?xindex>
+            %cast_0 = tensor.cast %30 : tensor<?xindex> to tensor<2xindex>
+            %31 = tensor.reshape %arg1(%cast_0) : (tensor<*xui32>, tensor<2xindex>) -> tensor<?x?xui32>
+            %32 = chlo.broadcast_minimum %29, %31 : (tensor<?x?xui32>, tensor<?x?xui32>) -> tensor<?x?xui32>
+            %cast_1 = tensor.cast %32 : tensor<?x?xui32> to tensor<*xui32>
+            scf.yield %cast_1 : tensor<*xui32>
+          } else {
+            %28 = arith.cmpi ule, %23, %c3 : index
+            %29 = scf.if %28 -> (tensor<*xui32>) {
+              %30 = shape.broadcast %19#0, %2 : tensor<?xindex>, tensor<3xindex> -> tensor<?xindex>
+              %cast = tensor.cast %30 : tensor<?xindex> to tensor<3xindex>
+              %31 = tensor.reshape %arg0(%cast) : (tensor<*xui32>, tensor<3xindex>) -> tensor<?x?x?xui32>
+              %32 = shape.broadcast %19#1, %2 : tensor<?xindex>, tensor<3xindex> -> tensor<?xindex>
+              %cast_0 = tensor.cast %32 : tensor<?xindex> to tensor<3xindex>
+              %33 = tensor.reshape %arg1(%cast_0) : (tensor<*xui32>, tensor<3xindex>) -> tensor<?x?x?xui32>
+              %34 = chlo.broadcast_minimum %31, %33 : (tensor<?x?x?xui32>, tensor<?x?x?xui32>) -> tensor<?x?x?xui32>
+              %cast_1 = tensor.cast %34 : tensor<?x?x?xui32> to tensor<*xui32>
+              scf.yield %cast_1 : tensor<*xui32>
+            } else {
+              %30 = arith.cmpi ule, %23, %c4 : index
+              %31 = scf.if %30 -> (tensor<*xui32>) {
+                %32 = shape.broadcast %19#0, %1 : tensor<?xindex>, tensor<4xindex> -> tensor<?xindex>
+                %cast = tensor.cast %32 : tensor<?xindex> to tensor<4xindex>
+                %33 = tensor.reshape %arg0(%cast) : (tensor<*xui32>, tensor<4xindex>) -> tensor<?x?x?x?xui32>
+                %34 = shape.broadcast %19#1, %1 : tensor<?xindex>, tensor<4xindex> -> tensor<?xindex>
+                %cast_0 = tensor.cast %34 : tensor<?xindex> to tensor<4xindex>
+                %35 = tensor.reshape %arg1(%cast_0) : (tensor<*xui32>, tensor<4xindex>) -> tensor<?x?x?x?xui32>
+                %36 = chlo.broadcast_minimum %33, %35 : (tensor<?x?x?x?xui32>, tensor<?x?x?x?xui32>) -> tensor<?x?x?x?xui32>
+                %cast_1 = tensor.cast %36 : tensor<?x?x?x?xui32> to tensor<*xui32>
+                scf.yield %cast_1 : tensor<*xui32>
+              } else {
+                %32 = arith.cmpi ule, %23, %c5 : index
+                cf.assert %32, "Input for dynamic binary or n-ary op lowering was of a rank greater than 5"
+                %33 = shape.broadcast %19#0, %0 : tensor<?xindex>, tensor<5xindex> -> tensor<?xindex>
+                %cast = tensor.cast %33 : tensor<?xindex> to tensor<5xindex>
+                %34 = tensor.reshape %arg0(%cast) : (tensor<*xui32>, tensor<5xindex>) -> tensor<?x?x?x?x?xui32>
+                %35 = shape.broadcast %19#1, %0 : tensor<?xindex>, tensor<5xindex> -> tensor<?xindex>
+                %cast_0 = tensor.cast %35 : tensor<?xindex> to tensor<5xindex>
+                %36 = tensor.reshape %arg1(%cast_0) : (tensor<*xui32>, tensor<5xindex>) -> tensor<?x?x?x?x?xui32>
+                %37 = chlo.broadcast_minimum %34, %36 : (tensor<?x?x?x?x?xui32>, tensor<?x?x?x?x?xui32>) -> tensor<?x?x?x?x?xui32>
+                %cast_1 = tensor.cast %37 : tensor<?x?x?x?x?xui32> to tensor<*xui32>
+                scf.yield %cast_1 : tensor<*xui32>
+              }
+              scf.yield %31 : tensor<*xui32>
+            }
+            scf.yield %29 : tensor<*xui32>
+          }
+          scf.yield %27 : tensor<*xui32>
+        }
+        scf.yield %25 : tensor<*xui32>
+      }
+      scf.yield %18 : tensor<*xui32>
+    }
+    scf.yield %16 : tensor<*xui32>
+  }
+  %10 = shape.shape_of %arg0 : tensor<*xui32> -> tensor<?xindex>
+  %11 = shape.shape_of %arg1 : tensor<*xui32> -> tensor<?xindex>
+  %12 = shape.broadcast %10, %11 : tensor<?xindex>, tensor<?xindex> -> tensor<?xindex>
+  %13 = tensor.reshape %9(%12) : (tensor<*xui32>, tensor<?xindex>) -> tensor<*xui32>
+  return %13 : tensor<*xui32>
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/hlo_to_kernel/tanh.mlir

@@ -0,0 +1,25 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: hlo_to_kernel --input=%s --output=%t --unroll_factors=4 --tile_sizes=256 --arch=sm_70
+
+func.func @tanh(%arg0: tensor<*xf32>) -> tensor<*xf32> attributes {tf_entry} {
+  %0 = shape.shape_of %arg0 : tensor<*xf32> -> tensor<?xindex>
+  %1 = shape.num_elements %0 : tensor<?xindex> -> index
+  %from_elements = tensor.from_elements %1 : tensor<1xindex>
+  %2 = tensor.reshape %arg0(%from_elements) : (tensor<*xf32>, tensor<1xindex>) -> tensor<?xf32>
+  %3 = mhlo.tanh %2 : tensor<?xf32>
+  %4 = tensor.reshape %3(%0) : (tensor<?xf32>, tensor<?xindex>) -> tensor<*xf32>
+  return %4 : tensor<*xf32>
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/invalid.mlir

@@ -0,0 +1,21 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s -split-input-file -verify-diagnostics
+
+func.func @alloc_raw(%ctx: !tf_framework.op_kernel_context, %size : index) {
+  // expected-error @+1 {{`dyn_sizes` count 1 does not match dynamic dimensions}}
+  %buf = tf_framework.alloc(%ctx, %size) : memref<?x10x?xi8>
+  func.return
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/isinf.mlir

@@ -0,0 +1,33 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: tf-opt %s --test-tf-lower-tf --xla-legalize-tf | \
+// RUN: mlir-hlo-opt \
+// RUN: --hlo-legalize-to-stablehlo=allow-xla-features \
+// RUN: --stablehlo-legalize-to-linalg \
+// RUN: --empty-tensor-to-alloc-tensor \
+// RUN: --computeop-and-func-bufferize --canonicalize | \
+// RUN: kernel-gen-opt -allow-unregistered-dialect \
+// RUN: --shape-to-descriptors \
+// RUN: --canonicalize --kernelgen-final-bufferize | \
+// RUN: FileCheck %s
+
+// Test whether all shape computations required for isinf can be lowered to
+// the standard dialect, scf and descriptors.
+// CHECK-LABEL: @isinf
+func.func @isinf(%arg0: tensor<?xf32>) -> tensor<?xi1> {
+  // CHECK-NOT: shape
+  %0 = "tf.IsInf"(%arg0) : (tensor<?xf32>) -> tensor<?xi1>
+  func.return %0 : tensor<?xi1>
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/legalize_tensor_reshape.mlir

@@ -0,0 +1,28 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt --split-input-file --legalize-tensor-reshape %s | FileCheck %s
+
+// CHECK-LABEL: func @tanh
+func.func @tanh(%arg0: tensor<*xf32>) -> tensor<*xf32> attributes {tf_entry} {
+  %0 = shape.shape_of %arg0 : tensor<*xf32> -> tensor<?xindex>
+  %1 = shape.num_elements %0 : tensor<?xindex> -> index
+  %from_elements = tensor.from_elements %1 : tensor<1xindex>
+  // CHECK: mhlo.dynamic_reshape
+  // CHECK-NOT: tensor.reshape
+  %2 = tensor.reshape %arg0(%from_elements) : (tensor<*xf32>, tensor<1xindex>) -> tensor<?xf32>
+  %3 = mhlo.tanh %2 : tensor<?xf32>
+  %4 = tensor.reshape %3(%0) : (tensor<?xf32>, tensor<?xindex>) -> tensor<*xf32>
+  return %4 : tensor<*xf32>
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/merge_assuming_ops.mlir

@@ -0,0 +1,362 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt --split-input-file --allow-unregistered-dialect \
+// RUN:   --mhlo-merge-assuming-ops --canonicalize --cse %s | \
+// RUN: FileCheck %s
+
+// Shape computations shall be reified.
+// CHECK-LABEL: @shape_of_unary
+// CHECK-SAME: (%[[ARG:.*]]: tensor<?x32xi16>)
+func.func @shape_of_unary(%arg : tensor<?x32xi16>) {
+  // CHECK: %[[SHAPE:.*]] = shape.shape_of %[[ARG]] : tensor<?x32xi16> -> tensor<?xindex>
+  // CHECK: "use"(%[[SHAPE]])
+  %0 = "mhlo.convert"(%arg) : (tensor<?x32xi16>) -> tensor<?x32xf16>
+  %1 = shape.shape_of %0 : tensor<?x32xf16> -> tensor<?xindex>
+  "use"(%1) : (tensor<?xindex>) -> ()
+  func.return
+}
+
+// -----
+
+// Shape computations shall be reified.
+// CHECK-LABEL: @shape_of_nary
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?x32xf16>, %[[ARG1:.*]]: tensor<?x32xf16>)
+func.func @shape_of_nary(%arg0 : tensor<?x32xf16>, %arg1 : tensor<?x32xf16>) {
+  // CHECK: %[[SHAPE:.*]] = shape.shape_of %[[ARG0]] : tensor<?x32xf16> -> tensor<?xindex>
+  // CHECK: "use"(%[[SHAPE]])
+  %0 = mhlo.subtract %arg0, %arg1 : tensor<?x32xf16>
+  %1 = mhlo.subtract %0, %arg1 : tensor<?x32xf16>
+  %2 = shape.shape_of %1 : tensor<?x32xf16> -> tensor<?xindex>
+  "use"(%2) : (tensor<?xindex>) -> ()
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL: @inline_bcasted_shape_operands
+// CHECK-SAME: (%[[A:.*]]: tensor<?xindex>, %[[B:.*]]: tensor<?xindex>, %[[C:.*]]: tensor<?xindex>)
+func.func @inline_bcasted_shape_operands(%a : tensor<?xindex>, %b : tensor<?xindex>,
+    %c : tensor<?xindex>) -> !shape.witness {
+  // CHECK-NOT: shape.broadcast
+  // CHECK:     %[[WITNESS:.*]] = shape.cstr_broadcastable %[[A]], %[[B]], %[[C]]
+  // CHECK:     return %[[WITNESS]] : !shape.witness
+  %0 = shape.broadcast %a, %b : tensor<?xindex>, tensor<?xindex>
+      -> tensor<?xindex>
+  %1 = shape.cstr_broadcastable %0, %c : tensor<?xindex>, tensor<?xindex>
+  func.return %1 : !shape.witness
+}
+
+// -----
+
+// CHECK-LABEL: @move_shape_of_into_assuming
+// CHECK-SAME: (%[[ARG0:.*]]: !shape.witness, %[[ARG1:.*]]: tensor<?x32xf32>)
+func.func @move_shape_of_into_assuming(%arg0 : !shape.witness,
+    %arg1 : tensor<?x32xf32>) -> tensor<2xindex> {
+  // CHECK:     %[[ASSUMING_RESULTS:.*]]:3 = shape.assuming %[[ARG0]] -> (tensor<?x32xf32>, tensor<?x32xf32>, tensor<2xindex>) {
+  // CHECK:       %[[DUMMY_TENSOR:.*]] = "dummy.tensor"() : () -> tensor<?x32xf32>
+  // CHECK:       %[[SHAPE:.*]] = shape.shape_of %[[DUMMY_TENSOR]]
+  // CHECK:       shape.assuming_yield %[[ARG1]], %[[DUMMY_TENSOR]], %[[SHAPE]]
+  // CHECK:     }
+  // CHECK-NOT: shape_of
+  // CHECK:     return %[[ASSUMING_RESULTS]]#2
+  %0:2 = shape.assuming %arg0 -> (tensor<?x32xf32>, tensor<?x32xf32>) {
+    %1 = "dummy.tensor"() : () -> tensor<?x32xf32>
+    shape.assuming_yield %arg1, %1 : tensor<?x32xf32>, tensor<?x32xf32>
+  }
+  %2 = shape.shape_of %0#1 : tensor<?x32xf32> -> tensor<2xindex>
+  "use"(%0#0, %0#1) : (tensor<?x32xf32>, tensor<?x32xf32>) -> ()
+  func.return %2 : tensor<2xindex>
+}
+
+// -----
+
+// CHECK-LABEL: @move_cstr_broadcastable_into_assuming
+// CHECK-SAME: (%[[ARG0:.*]]: !shape.witness, %[[ARG1:.*]]: tensor<2xindex>)
+func.func @move_cstr_broadcastable_into_assuming(%arg0 : !shape.witness,
+    %arg1 : tensor<2xindex>) -> !shape.witness {
+  // CHECK:     %[[ASSUMING_RESULTS:.*]]:3 = shape.assuming %[[ARG0]] -> (tensor<2xindex>, tensor<3xindex>, !shape.witness) {
+  // CHECK:       %[[DUMMY_TENSOR:.*]] = "dummy.tensor"() : () -> tensor<3xindex>
+  // CHECK:       %[[WITNESS:.*]] = shape.cstr_broadcastable %[[ARG1]], %[[DUMMY_TENSOR]]
+  // CHECK:       shape.assuming_yield %[[ARG1]], %[[DUMMY_TENSOR]], %[[WITNESS]]
+  // CHECK:     }
+  // CHECK-NOT: cstr_broadcastable
+  // CHECK:     return %[[ASSUMING_RESULTS]]#2
+  %0:2 = shape.assuming %arg0 -> (tensor<2xindex>, tensor<3xindex>) {
+    %1 = "dummy.tensor"() : () -> tensor<3xindex>
+    shape.assuming_yield %arg1, %1 : tensor<2xindex>, tensor<3xindex>
+  }
+  %1 = shape.cstr_broadcastable %arg1, %0#1 : tensor<2xindex>, tensor<3xindex>
+  "use"(%0#0, %0#1) : (tensor<2xindex>, tensor<3xindex>) -> ()
+  func.return %1 : !shape.witness
+}
+
+// -----
+
+// CHECK-LABEL: @not_move_shape_of_into_assuming
+// CHECK-SAME: (%[[W:.*]]: !shape.witness, %[[ARG0:.*]]: tensor<?x32xf32>, %[[ARG1:.*]]: tensor<?x32xf32>)
+func.func @not_move_shape_of_into_assuming(%arg0 : !shape.witness,
+    %arg1 : tensor<?x32xf32>, %arg2 : tensor<?x32xf32>) -> tensor<2xindex> {
+  // CHECK: %[[S:.*]] = shape.shape_of %[[ARG1]]
+  // CHECK: %[[ASS_RES:.*]] = shape.assuming %[[W]]
+  // CHECK:   shape.assuming_yield %[[ARG0]]
+  // CHECK: }
+  // CHECK: "some.other.op"(%[[ASS_RES]])
+  // CHECK: return %[[S]]
+  %0:2 = shape.assuming %arg0 -> (tensor<?x32xf32>, tensor<?x32xf32>) {
+    shape.assuming_yield %arg1, %arg2 : tensor<?x32xf32>, tensor<?x32xf32>
+  }
+  "some.other.op"(%0#0) : (tensor<?x32xf32>) -> ()
+  %2 = shape.shape_of %0#1 : tensor<?x32xf32> -> tensor<2xindex>
+  func.return %2 : tensor<2xindex>
+}
+
+// -----
+
+// CHECK-LABEL: @move_cstr_broadcastable_out_of_assuming
+// CHECK-SAME: (%[[ARG0:.*]]: !shape.witness, %[[ARG1:.*]]: tensor<2xindex>, %[[ARG2:.*]]: tensor<3xindex>)
+func.func @move_cstr_broadcastable_out_of_assuming(%arg0 : !shape.witness,
+    %arg1 : tensor<2xindex>, %arg2 : tensor<3xindex>) -> !shape.witness {
+  // CHECK:     %[[WITNESS:.*]] = shape.cstr_broadcastable %[[ARG1]], %[[ARG2]]
+  // CHECK-NOT: assuming
+  // CHECK-NOT: cstr_broadcastable
+  // CHECK:     return %[[WITNESS]]
+  %0 = shape.assuming %arg0 -> (!shape.witness) {
+    %1 = shape.cstr_broadcastable %arg1, %arg2 : tensor<2xindex>, tensor<3xindex>
+    shape.assuming_yield %1 : !shape.witness
+  }
+  func.return %0 : !shape.witness
+}
+
+// -----
+
+// CHECK-LABEL: @move_elementwise_into_assuming
+// CHECK-SAME:  (%[[ARG0:.*]]: !shape.witness, %[[ARG1:.*]]: tensor<?xf32>)
+func.func @move_elementwise_into_assuming(%arg0 : !shape.witness,
+    %arg1 : tensor<?xf32>) -> tensor<?xf32> {
+  // CHECK:     %[[RES:.*]] = shape.assuming %[[ARG0]]
+  // CHECK:       %[[SOME:.*]] = "some.op"
+  // CHECK:       %[[TANH:.*]] = mhlo.tanh %[[ARG1]]
+  // CHECK:       %[[BCAST_ADD:.*]] = chlo.broadcast_add %[[TANH]], %[[SOME]]
+  // CHECK:       shape.assuming_yield %[[BCAST_ADD]]
+  // CHECK-NOT: tanh
+  // CHECK-NOT: broadcast_add
+  // CHECK:     return %[[RES]]
+  %0:2 = shape.assuming %arg0 -> (tensor<?xf32>, tensor<?xf32>) {
+    %1 = "some.op"() : () -> tensor<?xf32>
+    shape.assuming_yield %arg1, %1 : tensor<?xf32>, tensor<?xf32>
+  }
+  %1 = mhlo.tanh %arg1 : tensor<?xf32>
+  %2 = chlo.broadcast_add %1, %0#1
+      : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>
+  func.return %2 : tensor<?xf32>
+}
+
+// -----
+
+// CHECK-LABEL: @move_shape_of_out_of_assuming
+// CHECK-SAME: (%[[ARG0:.*]]: !shape.witness, %[[ARG1:.*]]: tensor<2x?xf32>)
+func.func @move_shape_of_out_of_assuming(%arg0 : !shape.witness,
+    %arg1 : tensor<2x?xf32>) -> tensor<2xindex> {
+  // CHECK:     %[[SHAPE:.*]] = shape.shape_of %[[ARG1]]
+  // CHECK-NOT: assuming
+  // CHECK-NOT: cstr_broadcastable
+  // CHECK:     return %[[SHAPE]]
+    %0 = shape.assuming %arg0 -> (tensor<2xindex>) {
+    %1 = shape.shape_of %arg1 : tensor<2x?xf32> -> tensor<2xindex>
+    shape.assuming_yield %1 : tensor<2xindex>
+  }
+  func.return %0 : tensor<2xindex>
+}
+
+// -----
+
+// CHECK-LABEL: @move_shape_of_out_of_assuming
+// CHECK-SAME: (%[[ARG0:.*]]: !shape.witness, %[[ARG1:.*]]: tensor<2x?xf32>)
+func.func @move_shape_of_out_of_assuming(%arg0 : !shape.witness,
+    %arg1 : tensor<2x?xf32>) -> tensor<2xindex> {
+  // CHECK:     %[[SHAPE:.*]] = shape.shape_of %[[ARG1]]
+  // CHECK:     %{{.*}} = shape.assuming %[[ARG0]] -> (tensor<2x?xf32>) {
+  // CHECK:       %[[SOME_VAL:.*]] = "some.op"() : () -> tensor<2x?xf32>
+  // CHECK:       shape.assuming_yield %[[SOME_VAL]] : tensor<2x?xf32>
+  // CHECK:     }
+  // CHECK:     return %[[SHAPE]]
+  %0:2 = shape.assuming %arg0 -> (tensor<2x?xf32>, tensor<2xindex>) {
+    %1 = "some.op"() : () -> (tensor<2x?xf32>)
+    %2 = shape.shape_of %arg1 : tensor<2x?xf32> -> tensor<2xindex>
+    shape.assuming_yield %1, %2 : tensor<2x?xf32>, tensor<2xindex>
+  }
+  "use"(%0#0, %0#1) : (tensor<2x?xf32>, tensor<2xindex>) -> ()
+  func.return %0#1 : tensor<2xindex>
+}
+
+// -----
+
+// CHECK-LABEL: @not_move_shape_of_out_of_assuming
+// CHECK-SAME: (%[[ARG0:.*]]: !shape.witness, %[[ARG1:.*]]: tensor<2x?xf32>)
+func.func @not_move_shape_of_out_of_assuming(%arg0 : !shape.witness,
+    %arg1 : tensor<2x?xf32>) -> tensor<2xindex> {
+  // CHECK-NOT:  shape_of
+  // CHECK:      shape.assuming
+  // CHECK-SAME: {
+  // CHECK:        "some.tensor"
+  // CHECK:        shape_of
+  // CHECK:      }
+  %0 = shape.assuming %arg0 -> (tensor<2xindex>) {
+    %1 = "some.tensor"() : () -> tensor<2x?xf32>
+    %2 = shape.shape_of %1 : tensor<2x?xf32> -> tensor<2xindex>
+    shape.assuming_yield %2 : tensor<2xindex>
+  }
+  func.return %0 : tensor<2xindex>
+}
+
+// -----
+
+// CHECK: @merge_assuming_ops
+// CHECK: (%[[ARG0:.*]]: tensor<?x32xf16>, %[[ARG1:.*]]: tensor<?x32xf16>, %[[ARG2:.*]]: tensor<?x?x32xf16>)
+func.func @merge_assuming_ops(%arg0: tensor<?x32xf16>, %arg1 : tensor<?x32xf16>,
+    %arg2: tensor<?x?x32xf16>) -> tensor<?x?x32xf16> {
+  // CHECK:      %[[SHAPE0:.*]] = shape.shape_of %[[ARG0]]
+  // CHECK:      %[[SHAPE1:.*]] = shape.shape_of %[[ARG1]]
+  // CHECK:      %[[SHAPE2:.*]] = shape.shape_of %[[ARG2]]
+  // CHECK:      %[[WITNESS:.*]] = shape.cstr_broadcastable %[[SHAPE0]], %[[SHAPE1]], %[[SHAPE2]]
+  // CHECK:      %[[MERGED:.*]]:2 = shape.assuming %[[WITNESS]]
+  // CHECK-SAME: {
+  // CHECK:        "some.op"
+  // CHECK:        %[[RESULT0:.*]] = "some.producer"
+  // CHECK:        "another.op"
+  // CHECK:        %[[RESULT1:.*]] = "another.producer"
+  // CHECK:        shape.assuming_yield %[[RESULT0]], %[[RESULT1]]
+  // CHECK:      }
+  // CHECK:      return %[[MERGED]]#1
+  %0 = shape.shape_of %arg0 : tensor<?x32xf16> -> tensor<2xindex>
+  %1 = shape.shape_of %arg1 : tensor<?x32xf16> -> tensor<2xindex>
+  %2 = shape.shape_of %arg2 : tensor<?x?x32xf16> -> tensor<3xindex>
+  %3 = shape.cstr_broadcastable %0, %1 : tensor<2xindex>, tensor<2xindex>
+  %4 = shape.cstr_broadcastable %0, %1, %2 : tensor<2xindex>, tensor<2xindex>,
+      tensor<3xindex>
+  %5 = shape.assuming %3 -> (tensor<?x32xf16>) {
+    "some.op"() : () -> ()
+    %6 = "some.producer"() : () -> tensor<?x32xf16>
+    shape.assuming_yield %6 : tensor<?x32xf16>
+  }
+  %7 = shape.assuming %4 -> (tensor<?x?x32xf16>) {
+    "another.op"() : () -> ()
+    %8 = "another.producer"() : () -> tensor<?x?x32xf16>
+    shape.assuming_yield %8 : tensor<?x?x32xf16>
+  }
+  "use"(%5, %7) : (tensor<?x32xf16>, tensor<?x?x32xf16>) -> ()
+  func.return %7 : tensor<?x?x32xf16>
+}
+
+// -----
+
+// Do not merge assuming ops if witness will not dominate use.
+// CHECK: @do_not_merge_assuming_ops
+func.func @do_not_merge_assuming_ops() {
+  // CHECK: shape.assuming
+  // CHECK: shape.assuming
+  %0 = "some.witness"() : () -> !shape.witness
+  %1 = shape.assuming %0 -> (!shape.witness) {
+    %2 = "some.witness"() : () -> !shape.witness
+    shape.assuming_yield %2 : !shape.witness
+  }
+  shape.assuming %1 {
+    "some.op"() : () -> ()
+    shape.assuming_yield
+  }
+  func.return
+}
+
+// -----
+
+// CHECK:      @eliminate_extent_tensor_cast
+// CHECK-SAME: (%[[ARG:.*]]: tensor<2x?x4xf32>)
+func.func @eliminate_extent_tensor_cast(%arg : tensor<2x?x4xf32>) {
+  // CHECK-NOT:  shape_of
+  // CHECK:      %[[RESULT:.*]] = shape.shape_of %[[ARG]] : tensor<2x?x4xf32> -> tensor<3xindex>
+  // CHECK-NEXT: "use"(%[[RESULT]]) : (tensor<3xindex>) -> ()
+  %0 = shape.shape_of %arg : tensor<2x?x4xf32> -> tensor<?xindex>
+  %1 = tensor.cast %0 : tensor<?xindex> to tensor<3xindex>
+  "use"(%1) : (tensor<3xindex>) -> ()
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL: @redundant_cstr_broadcastable
+// CHECK-SAME: (%[[ARG0:.*]]: tensor<?xindex>, %[[ARG1:.*]]: tensor<?xindex>)
+func.func @redundant_cstr_broadcastable(%arg0: tensor<?xindex>,
+    %arg1 : tensor<?xindex>) {
+  // CHECK-DAG:  %[[WITNESS:.*]] = shape.cstr_broadcastable %[[ARG0]], %[[ARG1]]
+  // CHECK:      shape.assuming %[[WITNESS]]
+  %0 = shape.cstr_broadcastable %arg0, %arg1 : tensor<?xindex>, tensor<?xindex>
+  %1 = shape.cstr_broadcastable %arg0, %arg1 : tensor<?xindex>, tensor<?xindex>
+  %2 = shape.assuming_all %0, %1
+  shape.assuming %2 -> () {
+    "some.op"() : () -> ()
+    shape.assuming_yield
+  }
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL: @move_assuming_all_over_assuming_region
+// CHECK-SAME:  %[[ARG0:.*]]: tensor<?xindex>, %[[ARG1:.*]]: tensor<?xindex>, %[[ARG2:.*]]: tensor<?xindex>, %[[ARG3:.*]]: tensor<?xindex>, %[[ARG4:.*]]: tensor<?xindex>
+func.func @move_assuming_all_over_assuming_region(%arg0: tensor<?xindex>,
+    %arg1 : tensor<?xindex>, %arg2 : tensor<?xindex>, %arg3 : tensor<?xindex>,
+    %arg4 : tensor<?xindex>) {
+  // CHECK-DAG: %[[CSTR0:.*]] = shape.cstr_broadcastable %[[ARG0]], %[[ARG1]]
+  // CHECK-DAG: %[[CSTR1:.*]] = shape.cstr_broadcastable %[[ARG1]], %[[ARG2]]
+  // CHECK-DAG: %[[CSTR2:.*]] = shape.cstr_broadcastable %[[ARG2]], %[[ARG3]]
+  // CHECK-DAG: %[[CSTR3:.*]] = shape.cstr_broadcastable %[[ARG3]], %[[ARG4]]
+  // CHECK-DAG: %[[CSTR_ALL:.*]] = shape.assuming_all %[[CSTR0]], %[[CSTR1]], %[[CSTR2]], %[[CSTR3]]
+  // CHECK:     shape.assuming %[[CSTR_ALL]] {
+  // CHECK:       "some.op"()
+  // CHECK:       "some.op"()
+  // CHECK:     }
+  // CHECK:     return
+  %0 = shape.cstr_broadcastable %arg0, %arg1 : tensor<?xindex>, tensor<?xindex>
+  %1 = shape.cstr_broadcastable %arg1, %arg2 : tensor<?xindex>, tensor<?xindex>
+  %2 = shape.assuming_all %0, %1
+  shape.assuming %2 -> () {
+    "some.op"() : () -> ()
+  }
+  %3 = shape.cstr_broadcastable %arg2, %arg3 : tensor<?xindex>, tensor<?xindex>
+  %4 = shape.cstr_broadcastable %arg3, %arg4 : tensor<?xindex>, tensor<?xindex>
+  %5 = shape.assuming_all %3, %4
+  shape.assuming %5 -> () {
+    "some.op"() : () -> ()
+  }
+  func.return
+}
+
+// -----
+
+// CHECK-LABEL: @move_down_into_assuming
+// CHECK-SAME:  (%[[ARG:.*]]: tensor<?x32xi16>, %[[W:.*]]: !shape.witness)
+func.func @move_down_into_assuming(%arg0: tensor<?x32xi16>, %w: !shape.witness) -> tensor<?x32xf16> {
+  // CHECK: %[[RES:.*]] = shape.assuming %[[W]]
+  // CHECK:   %[[INNER_RES:.*]] = mhlo.convert %[[ARG]]
+  // CHECK:   shape.assuming_yield %[[INNER_RES]]
+  // CHECK: }
+  // CHECK: return %[[RES]]
+  %0 = mhlo.convert %arg0 : (tensor<?x32xi16>) -> tensor<?x32xf16>
+  "some.possibly_side_effecting_op"() : () -> ()
+  %4 = shape.assuming %w -> (tensor<?x32xf16>) {
+    shape.assuming_yield %0 : tensor<?x32xf16>
+  }
+  func.return %4 : tensor<?x32xf16>
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/ops.mlir

@@ -0,0 +1,129 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s | FileCheck %s
+// Verify the printed output can be parsed.
+// RUN: kernel-gen-opt %s | kernel-gen-opt | FileCheck %s
+// Verify the generic form can be parsed.
+// RUN: kernel-gen-opt -mlir-print-op-generic %s | kernel-gen-opt | FileCheck %s
+
+// CHECK-LABEL: func @alloc
+func.func @alloc(%ctx: !tf_framework.op_kernel_context,
+                   %size_0 : index , %size_2 : index) {
+  %buf_0 = tf_framework.alloc(%ctx) : memref<10xi8>
+  %buf_1 = tf_framework.alloc(%ctx, %size_0, %size_2) : memref<?x10x?xi8>
+  func.return
+}
+
+// CHECK-LABEL: func @forwarding_alloc
+func.func @forwarding_alloc(%ctx: !tf_framework.op_kernel_context,
+                       %size_0 : index , %size_2 : index) {
+  %buf = tf_framework.alloc(%ctx, %size_0, %size_2) {
+    input_indices = [0 : i32, 1 : i32],
+    output_index = 0 : i32
+  } : memref<?x10x?xi8>
+  func.return
+}
+
+// CHECK-LABEL: func @dealloc
+func.func @dealloc(%ctx: !tf_framework.op_kernel_context,
+              %memref : memref<?x10xf32>) {
+  tf_framework.dealloc(%ctx, %memref) : memref<?x10xf32>
+  func.return
+}
+
+// CHECK-LABEL: func @assert
+func.func @assert(%ctx: !tf_framework.op_kernel_context, %cond: i1) {
+  tf_framework.assert %ctx, %cond, "ALREADY_EXISTS", "Or maybe not"
+  func.return
+}
+
+// CHECK-LABEL: func @report_error
+func.func @report_error(%ctx: !tf_framework.op_kernel_context) {
+  tf_framework.report_error %ctx, "INVALID_ARGUMENT", "Everything is awesome"
+  func.return
+}
+
+// CHECK-LABEL: func @null_memref
+func.func @null_memref() {
+  tf_framework.null_memref : memref<*xf32>
+  func.return
+}
+
+// CHECK-LABEL: func @null_context
+func.func @null_context() {
+  tf_framework.null_context : !tf_framework.op_kernel_context
+  func.return
+}
+
+// CHECK-LABEL: func @is_valid_memref
+func.func @is_valid_memref(%buf: memref<?xf32>) -> i1 {
+  %pred = tf_framework.is_valid_memref(%buf) : memref<?xf32> -> i1
+  func.return %pred : i1
+}
+
+// CHECK-LABEL: func @jit_compile_wo_ctx
+func.func @jit_compile_wo_ctx() -> !tf_framework.jit_callable {
+  %callable = tf_framework.jit_compile {
+  ^bb0(%arg : tensor<2x?xf32>):
+    tf_framework.jit_compile_yield %arg : tensor<2x?xf32>
+  }
+  func.return %callable : !tf_framework.jit_callable
+}
+
+// CHECK-LABEL: func @jit_compile
+func.func @jit_compile(%ctx : !tf_framework.op_kernel_context)
+    -> !tf_framework.jit_callable {
+  %callable = tf_framework.jit_compile %ctx {
+  ^bb0(%arg : tensor<2x?xf32>):
+    tf_framework.jit_compile_yield %arg : tensor<2x?xf32>
+  }
+  func.return %callable : !tf_framework.jit_callable
+}
+
+// CHECK-LABEL: func @jit_compile_from_str_wo_ctx
+func.func @jit_compile_from_str_wo_ctx() -> !tf_framework.jit_callable {
+  %callable = tf_framework.jit_compile_from_str "placeholder" {
+      architectures = ["sm_123", "sm_456"], tileSizes = [1, 2, 3],
+      unrollFactors = [4], enableFtz = false,
+      index64Bit = false, cpuCodegen = false }
+  func.return %callable : !tf_framework.jit_callable
+}
+
+// CHECK-LABEL: func @jit_compile_from_str
+func.func @jit_compile_from_str(%ctx : !tf_framework.op_kernel_context)
+    -> !tf_framework.jit_callable {
+  %callable = tf_framework.jit_compile_from_str %ctx , "placeholder" {
+      architectures = ["sm_123", "sm_456"], tileSizes = [1, 2, 3],
+      unrollFactors = [4], enableFtz = false,
+      index64Bit = false, cpuCodegen = false }
+  func.return %callable : !tf_framework.jit_callable
+}
+
+// CHECK-LABEL: func @jit_execute_wo_ctx
+func.func @jit_execute_wo_ctx(%callable : !tf_framework.jit_callable,
+    %arg : tensor<2x?xf32>) -> tensor<2x?xf32> {
+  %0 = tf_framework.jit_execute %callable(%arg)
+      : tensor<2x?xf32> -> tensor<2x?xf32>
+  func.return %0 : tensor<2x?xf32>
+}
+
+// CHECK-LABEL: func @jit_execute
+func.func @jit_execute(%ctx : !tf_framework.op_kernel_context,
+    %callable : !tf_framework.jit_callable, %arg : tensor<2x?xf32>)
+    -> tensor<2x?xf32> {
+  %0 = tf_framework.jit_execute ctx(%ctx) %callable(%arg)
+      : tensor<2x?xf32> -> tensor<2x?xf32>
+  func.return %0 : tensor<2x?xf32>
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/parallel_loops_to_sequential.mlir

@@ -0,0 +1,31 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s --parallel-loops-to-sequential | FileCheck %s
+
+// CHECK-LABEL: @parallel_loop
+func.func @parallel_loop(%lb_0 : index, %lb_1 : index,
+                     %ub_0 : index, %ub_1 : index,
+                     %s_0 : index, %s_1 : index,
+                     %buf: memref<?x?xindex>) {
+  scf.parallel (%i0, %i1) = (%lb_0, %lb_1) to (%ub_0, %ub_1) step (%s_0, %s_1) {
+    %sum_elem = arith.addi %i0, %i1 : index
+    memref.store %sum_elem, %buf[%i0, %i1] : memref<?x?xindex>
+  }
+  func.return
+}
+// CHECK: scf.for [[I_0:%.*]] = [[LB_0:%.*]] to [[UB_0:%.*]] step [[S_0:%.*]]
+// CHECK:   scf.for [[I_1:%.*]] = [[LB_1:%.*]] to [[UB_1:%.*]] step [[S_1:%.*]]
+// CHECK:     [[SUM:%.*]] = arith.addi [[I_0]], [[I_1]] : index
+// CHECK:     memref.store [[SUM]], {{%.*}}{{\[}}[[I_0]], [[I_1]]] : memref<?x?xindex>

tensorflow/compiler/mlir/tools/kernel_gen/tests/rewrite_tf_framework_assert.mlir

@@ -0,0 +1,77 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s -rewrite-tf-framework-assert -split-input-file |\
+// RUN: FileCheck %s
+
+// CHECK-LABEL: func @assert(
+// CHECK-SAME: [[CTX:%.*]]: !tf_framework.op_kernel_context
+func.func @assert(%ctx: !tf_framework.op_kernel_context)
+       -> (memref<*xf32>, memref<*xi32>) attributes {tf_entry} {
+  %true = arith.constant true
+  tf_framework.assert %ctx, %true, INVALID_ARGUMENT, "the one and only"
+  %buf_f32 = tf_framework.alloc(%ctx) : memref<2xf32>
+  %unranked_f32 = memref.cast %buf_f32 : memref<2xf32> to memref<*xf32>
+  %buf_i32 = tf_framework.alloc(%ctx) : memref<3xi32>
+  %unranked_i32 = memref.cast %buf_i32 : memref<3xi32> to memref<*xi32>
+  func.return %unranked_f32, %unranked_i32 : memref<*xf32>, memref<*xi32>
+}
+// CHECK:   [[TRUE:%.*]] = arith.constant true
+// CHECK:   cf.cond_br [[TRUE]], ^bb1, ^bb2
+// CHECK: ^bb1:
+// CHECK:   [[BUF_F32:%.*]] = tf_framework.alloc([[CTX]]) : memref<2xf32>
+// CHECK:   [[OUT_F32:%.*]] = memref.cast [[BUF_F32]]
+// CHECK:   [[BUF_I32:%.*]] = tf_framework.alloc([[CTX]]) : memref<3xi32>
+// CHECK:   [[OUT_I32:%.*]] = memref.cast [[BUF_I32]]
+// CHECK:   return [[OUT_F32]], [[OUT_I32]] : memref<*xf32>, memref<*xi32>
+// CHECK: ^bb2:
+// CHECK:   tf_framework.report_error [[CTX]], INVALID_ARGUMENT,
+// CHECK-SAME: "the one and only"
+// CHECK:   [[NULL_F32:%.*]] = tf_framework.null_memref : memref<*xf32>
+// CHECK:   [[NULL_I32:%.*]] = tf_framework.null_memref : memref<*xi32>
+// CHECK:   return [[NULL_F32]], [[NULL_I32]] : memref<*xf32>, memref<*xi32>
+
+// -----
+
+// CHECK-LABEL: func @double_assert(
+// CHECK-SAME: [[CTX:%.*]]: !tf_framework.op_kernel_context
+func.func @double_assert(%ctx: !tf_framework.op_kernel_context)
+       -> memref<*xf32> attributes {tf_entry} {
+  %true = arith.constant true
+  %false = arith.constant false
+  tf_framework.assert %ctx, %true, INVALID_ARGUMENT, "first assertion"
+  tf_framework.assert %ctx, %false, INVALID_ARGUMENT, "second assertion"
+  %buf = tf_framework.alloc(%ctx) : memref<2xf32>
+  %unranked_buf = memref.cast %buf : memref<2xf32> to memref<*xf32>
+  func.return %unranked_buf : memref<*xf32>
+}
+// CHECK-DAG:   [[TRUE:%.*]] = arith.constant true
+// CHECK-DAG:   [[FALSE:%.*]] = arith.constant false
+// CHECK:   cf.cond_br [[TRUE]], ^bb1, ^bb3
+// CHECK: ^bb1:
+// CHECK:   cf.cond_br [[FALSE]], ^bb2, ^bb4
+// CHECK: ^bb2:
+// CHECK:   [[BUF:%.*]] = tf_framework.alloc([[CTX]]) : memref<2xf32>
+// CHECK:   [[OUT:%.*]] = memref.cast [[BUF]]
+// CHECK:   return [[OUT]] : memref<*xf32>
+// CHECK: ^bb3:
+// CHECK:   tf_framework.report_error [[CTX]], INVALID_ARGUMENT,
+// CHECK-SAME: "first assertion"
+// CHECK:   [[NULL:%.*]] = tf_framework.null_memref : memref<*xf32>
+// CHECK:   return [[NULL]] : memref<*xf32>
+// CHECK: ^bb4:
+// CHECK:   tf_framework.report_error [[CTX]], INVALID_ARGUMENT,
+// CHECK-SAME: "second assertion"
+// CHECK:   [[NULL:%.*]] = tf_framework.null_memref : memref<*xf32>
+// CHECK:   return [[NULL]] : memref<*xf32>

tensorflow/compiler/mlir/tools/kernel_gen/tests/shape_simplification.mlir

@@ -0,0 +1,157 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt -split-input-file -shape-simplification %s | FileCheck %s
+
+// Incompatible shapes. No folding.
+// CHECK-LABEL: func @f
+func.func @f() -> !shape.shape {
+  // CHECK: shape.broadcast
+  %0 = shape.const_shape [2] : !shape.shape
+  %1 = shape.const_shape [7] : !shape.shape
+  %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
+  func.return %2 : !shape.shape
+}
+
+// -----
+
+// Broadcast of partially dynamic shapes yields a static shape.
+// CHECK-LABEL: func @f
+func.func @f(%arg0 : tensor<42x?x42x?xf32>, %arg1 : tensor<42x?x?xf32>) -> !shape.shape {
+  // CHECK: %[[CST:.*]] = shape.const_shape [42, 42, 42, 256] : !shape.shape
+  // CHECK: return %[[CST]]
+  %0 = shape.const_shape [256] : !shape.shape
+  %1 = shape.shape_of %arg0 : tensor<42x?x42x?xf32> -> !shape.shape
+  %2 = shape.shape_of %arg1 : tensor<42x?x?xf32> -> !shape.shape
+  %3 = shape.broadcast %0, %1, %2 : !shape.shape, !shape.shape, !shape.shape -> !shape.shape
+  func.return %3 : !shape.shape
+}
+
+// -----
+
+// Remove operands that don't contribute to the result.
+// CHECK-LABEL: func @f
+func.func @f(%arg0 : tensor<?x?x42x42xf32>, %arg1 : tensor<42x42xf32>) -> tensor<?xindex> {
+  // CHECK: %[[SHAPE0:.*]] = shape.shape_of %arg0 : tensor<?x?x42x42xf32> -> tensor<?xindex>
+  // CHECK: return %[[SHAPE0]]
+  %0 = shape.const_shape [42, 1] : tensor<2xindex>
+  %1 = shape.shape_of %arg0 : tensor<?x?x42x42xf32> -> tensor<?xindex>
+  %2 = shape.shape_of %arg1 : tensor<42x42xf32> -> tensor<2xindex>
+  %3 = shape.broadcast %0, %1, %2 : tensor<2xindex>, tensor<?xindex>, tensor<2xindex> -> tensor<?xindex>
+  func.return %3 : tensor<?xindex>
+}
+
+// -----
+
+// The constant shape needs to stay alive or the result will be smaller.
+// CHECK-LABEL: func @f
+func.func @f(%arg0 : tensor<?xf32>) -> !shape.shape {
+  // CHECK: shape.broadcast
+  %0 = shape.const_shape [1, 1] : !shape.shape
+  %1 = shape.shape_of %arg0 : tensor<?xf32> -> !shape.shape
+  %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
+  func.return %2 : !shape.shape
+}
+
+// -----
+
+// [256] is the only contributor of that constant, keep it.
+// CHECK-LABEL: func @f
+func.func @f(%arg0 : tensor<?x?xf32>) -> !shape.shape {
+  // CHECK: shape.broadcast
+  %0 = shape.const_shape [256] : !shape.shape
+  %1 = shape.shape_of %arg0 : tensor<?x?xf32> -> !shape.shape
+  %2 = shape.broadcast %0, %1 : !shape.shape, !shape.shape -> !shape.shape
+  func.return %2 : !shape.shape
+}
+
+// -----
+
+// CHECK-LABEL: func @static_non1_succeeds
+// CHECK-NEXT:   %[[C2:.*]] = arith.constant 2
+// CHECK-NEXT:   return %[[C2]]
+func.func @static_non1_succeeds(%arg0 : tensor<?x?xf64>, %arg1 : tensor<?x1xf64>,
+                           %arg2: tensor<?x2xf64>) -> index {
+ %c1 = arith.constant 1 : index
+ %1 = shape.shape_of %arg0 : tensor<?x?xf64> -> tensor<2xindex>
+ %2 = shape.shape_of %arg1 : tensor<?x1xf64> -> tensor<2xindex>
+ %3 = shape.shape_of %arg2 : tensor<?x2xf64> -> tensor<2xindex>
+ %4 = shape.broadcast %1, %2, %3 : tensor<2xindex>, tensor<2xindex>,
+                                   tensor<2xindex> -> tensor<2xindex>
+ %result = tensor.extract %4[%c1] : tensor<2xindex>
+ func.return %result : index
+}
+
+// -----
+
+// CHECK-LABEL: func @all_static_1s_succeeds
+// CHECK-NEXT:   %[[C1:.*]] = arith.constant 1
+// CHECK-NEXT:   return %[[C1]]
+func.func @all_static_1s_succeeds(%arg0 : tensor<?x1xf64>, %arg1 : tensor<?x1xf64>)
+                            -> index {
+ %c1 = arith.constant 1 : index
+ %1 = shape.shape_of %arg0 : tensor<?x1xf64> -> tensor<2xindex>
+ %2 = shape.shape_of %arg1 : tensor<?x1xf64> -> tensor<2xindex>
+ %3 = shape.broadcast %1, %2 : tensor<2xindex>, tensor<2xindex>
+                               -> tensor<2xindex>
+ %result = tensor.extract %3[%c1] : tensor<2xindex>
+ func.return %result : index
+}
+
+// -----
+
+// CHECK-LABEL: func @single_non_static_1_succeeds
+// CHECK-SAME:    %[[ARG0:.*]]: tensor<?x?xf64>
+// CHECK:        %[[C1:.*]] = arith.constant 1
+// CHECK:        %[[DIM:.*]] = tensor.dim %[[ARG0]], %[[C1]]
+// CHECK:        return %[[DIM]]
+func.func @single_non_static_1_succeeds(%arg0 : tensor<?x?xf64>,
+                                   %arg1 : tensor<?x1xf64>) -> index {
+ %c0 = arith.constant 1 : index
+ %1 = shape.shape_of %arg0 : tensor<?x?xf64> -> tensor<2xindex>
+ %2 = shape.shape_of %arg1 : tensor<?x1xf64> -> tensor<2xindex>
+ %3 = shape.broadcast %1, %2 : tensor<2xindex>, tensor<2xindex>
+                               -> tensor<2xindex>
+ %result = tensor.extract %3[%c0] : tensor<2xindex>
+ func.return %result : index
+}
+
+// -----
+
+// CHECK-LABEL: func @multiple_non_static_1_fails
+// CHECK-NEXT:   constant 0
+// CHECK-NEXT:   shape.shape_of
+// CHECK-NEXT:   shape.shape_of
+// CHECK-NEXT:   shape.broadcast
+// CHECK-NEXT:   %[[RESULT:.*]] = tensor.extract
+// CHECK-NEXT:   return %[[RESULT]]
+func.func @multiple_non_static_1_fails(%arg0 : tensor<?x?xf64>,
+                                  %arg1 : tensor<?x1xf64>) -> index {
+ %c0 = arith.constant 0 : index
+ %1 = shape.shape_of %arg0 : tensor<?x?xf64> -> tensor<2xindex>
+ %2 = shape.shape_of %arg1 : tensor<?x1xf64> -> tensor<2xindex>
+ %3 = shape.broadcast %1, %2 : tensor<2xindex>, tensor<2xindex>
+                               -> tensor<2xindex>
+ %result = tensor.extract %3[%c0] : tensor<2xindex>
+ func.return %result : index
+}
+
+// -----
+
+// CHECK-LABEL: func @extract_no_crash
+// CHECK-NEXT:   tensor.extract
+func.func @extract_no_crash(%arg0 : tensor<index>) -> index {
+ %result = tensor.extract %arg0[] : tensor<index>
+ func.return %result : index
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_abi_knowledge.mlir

@@ -0,0 +1,369 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s -allow-unregistered-dialect -propagate-tf-abi-knowledge-to-kernels -split-input-file | FileCheck %s --check-prefixes=CHECK,ABI
+// RUN: kernel-gen-opt %s -allow-unregistered-dialect -propagate-shape-knowledge-to-kernels -split-input-file | FileCheck %s --check-prefixes=CHECK,SHAPE
+
+// The input is taken from what is actually used in kernel generator lowering
+// for unary operations.
+
+// CHECK-LABEL: module attributes {gpu.container_module}
+module attributes {gpu.container_module} {
+  // CHECK-LABEL: func @abs
+  func.func @abs(%ctx: !tf_framework.op_kernel_context, %arg0: memref<*xf32>, %size: index)
+      attributes {tf_entry} {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %13 = memref.alloca() : memref<1xindex>
+    memref.store %size, %13[%c0] : memref<1xindex>
+    %14 = memref.reshape %arg0(%13) : (memref<*xf32>, memref<1xindex>) -> memref<?xf32>
+    %15 = memref.dim %14, %c0 : memref<?xf32>
+    %16 = tf_framework.alloc(%ctx, %15) : memref<?xf32>
+    gpu.launch_func @abs_kernel::@abs_kernel
+        blocks in (%c1, %c1, %c1) threads in (%c1, %c1, %c1)
+        args(%14 : memref<?xf32>, %16 : memref<?xf32>)
+    func.return
+  }
+
+  // CHECK-LABEL: gpu.module @abs_kernel
+  gpu.module @abs_kernel {
+    // CHECK-LABEL: llvm.func @abs_kernel
+    // ABI-SAME: %[[ARG0:.*]]: !llvm.ptr, %[[ARG1:.*]]: !llvm.ptr {llvm.align = 16 : index},
+    // ABI-SAME: %[[ARG2:.*]]: i64, %[[ARG3:.*]]: i64, %[[ARG4:.*]]: i64, %[[ARG5:.*]]: !llvm.ptr, %[[ARG6:.*]]: !llvm.ptr {llvm.align = 16 : index, llvm.noalias},
+    // ABI-SAME: %[[ARG7:.*]]: i64, %[[ARG8:.*]]: i64, %[[ARG9:.*]]: i64
+    // SHAPE-SAME: %[[ARG0:.*]]: !llvm.ptr, %[[ARG1:.*]]: !llvm.ptr, %[[ARG2:.*]]: i64, %[[ARG3:.*]]: i64, %[[ARG4:.*]]: i64, %[[ARG5:.*]]: !llvm.ptr, %[[ARG6:.*]]: !llvm.ptr, %[[ARG7:.*]]: i64, %[[ARG8:.*]]: i64, %[[ARG9:.*]]: i64
+    llvm.func @abs_kernel(%arg0: !llvm.ptr, %arg1: !llvm.ptr, %arg2: i64, %arg3: i64, %arg4: i64, %arg5: !llvm.ptr, %arg6: !llvm.ptr, %arg7: i64, %arg8: i64, %arg9: i64) attributes {gpu.kernel} {
+      // ABI: %[[ZERO:.*]] = llvm.mlir.constant(0 : index)
+      // ABI: %[[ONE:.*]] = llvm.mlir.constant(1 : index)
+      // CHECK: llvm.mlir.undef
+      %0 = llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[ARG1]]
+      // SHAPE-NEXT: llvm.insertvalue %[[ARG0]]
+      %1 = llvm.insertvalue %arg0, %0[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // CHECK-NEXT: llvm.insertvalue %[[ARG1]]
+      %2 = llvm.insertvalue %arg1, %1[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[ZERO]]
+      // SHAPE-NEXT: llvm.insertvalue %[[ARG2]]
+      %3 = llvm.insertvalue %arg2, %2[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // CHECK-NEXT: llvm.insertvalue %[[ARG3]]
+      %4 = llvm.insertvalue %arg3, %3[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[ONE]]
+      // SHAPE-NEXT: llvm.insertvalue %[[ARG4]]
+      %5 = llvm.insertvalue %arg4, %4[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // CHECK-NEXT: llvm.mlir.undef
+      %6 = llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[ARG6]]
+      // SHAPE-NEXT: llvm.insertvalue %[[ARG5]]
+      %7 = llvm.insertvalue %arg5, %6[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // CHECK-NEXT: llvm.insertvalue %[[ARG6]]
+      %8 = llvm.insertvalue %arg6, %7[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[ZERO]]
+      // SHAPE-NEXT: llvm.insertvalue %[[ARG7]]
+      %9 = llvm.insertvalue %arg7, %8[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[ARG8]]
+      // SHAPE-NEXT: llvm.insertvalue %[[ARG3]]
+      %10 = llvm.insertvalue %arg8, %9[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[ONE]]
+      // SHAPE-NEXT: llvm.insertvalue %[[ARG4]]
+      %11 = llvm.insertvalue %arg9, %10[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      llvm.return
+      // CHECK-NEXT: llvm.return
+    }
+  }
+}
+
+// -----
+
+// Binary op without broadcasting (same shape).
+
+// CHECK-LABEL: module attributes {gpu.container_module}
+module attributes {gpu.container_module} {
+  // CHECK-LABEL: func @add_same_shape
+  func.func @add_same_shape(%arg0: !tf_framework.op_kernel_context, %arg1: memref<*xf32>, %arg2: memref<*xf32>, %size: index)
+      attributes {tf_entry} {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %82 = memref.reinterpret_cast %arg1 to offset: [0], sizes: [%size], strides: [%c1]: memref<*xf32> to memref<?xf32, strided<[?], offset: 0>>
+    %83 = memref.reinterpret_cast %arg2 to offset: [0], sizes: [%size], strides: [%c1]: memref<*xf32> to memref<?xf32, strided<[?], offset: 0>>
+    %84 = tf_framework.alloc(%arg0, %size) : memref<?xf32>
+    gpu.launch_func  @AddV2_kernel_1::@AddV2_kernel blocks in (%c1, %c1, %c1) threads in (%c1, %c1, %c1) args(%size : index, %82 : memref<?xf32, strided<[?], offset: 0>>, %83 : memref<?xf32, strided<[?], offset: 0>>, %84 : memref<?xf32>)
+    func.return
+  }
+
+  // CHECK-LABEL: gpu.module @AddV2_kernel_1
+  gpu.module @AddV2_kernel_1 {
+    // CHECK-LABEL: llvm.func @AddV2_kernel
+    // ABI-SAME: {llvm.align = 16 : index}
+    // ABI-SAME: {llvm.align = 16 : index}
+    // ABI-SAME: {llvm.align = 16 : index, llvm.noalias}
+    llvm.func @AddV2_kernel(%arg0: i64, %arg1: !llvm.ptr, %arg2: !llvm.ptr, %arg3: i64, %arg4: i64, %arg5: i64, %arg6: !llvm.ptr, %arg7: !llvm.ptr, %arg8: i64, %arg9: i64, %arg10: i64, %arg11: !llvm.ptr, %arg12: !llvm.ptr, %arg13: i64, %arg14: i64, %arg15: i64) attributes {gpu.kernel} {
+      // ABI: %[[C0:.*]] = llvm.mlir.constant(0 : index) : i64
+      // ABI: %[[C1:.*]] = llvm.mlir.constant(1 : index) : i64
+      %0 = llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %1 = llvm.insertvalue %arg1, %0[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %2 = llvm.insertvalue %arg2, %1[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %3 = llvm.insertvalue %arg3, %2[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %4 = llvm.insertvalue %arg4, %3[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %5 = llvm.insertvalue %arg5, %4[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR0:.*]], %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR0]], %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C0]], %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C1]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[SHP:.*]], %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[STR:.*]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %6 = llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %7 = llvm.insertvalue %arg6, %6[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %8 = llvm.insertvalue %arg7, %7[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %9 = llvm.insertvalue %arg8, %8[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %10 = llvm.insertvalue %arg9, %9[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %11 = llvm.insertvalue %arg10, %10[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR1:.*]], %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR1]], %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C0]], %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C1]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[SHP]], %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NOT: llvm.insertvalue %[[STR]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %12 = llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %13 = llvm.insertvalue %arg11, %12[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %14 = llvm.insertvalue %arg12, %13[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %15 = llvm.insertvalue %arg13, %14[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %16 = llvm.insertvalue %arg14, %15[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %17 = llvm.insertvalue %arg15, %16[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR2:.*]], %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR2]], %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C0]], %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C1]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[SHP]], %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[STR1:.*]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NOT: llvm.insertvalue %[[STR]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      llvm.return
+      // CHECK-NEXT: llvm.return
+    }
+  }
+}
+
+// -----
+
+// Test binary op with broadcasting - 2d case.
+
+// CHECK-LABEL: module attributes {gpu.container_module}
+module attributes {gpu.container_module} {
+  // CHECK-LABEL: func @add_same_shape
+  func.func @add_same_shape(%arg0: !tf_framework.op_kernel_context, %arg1: memref<*xf32>, %arg2: memref<*xf32>, %size0: index, %size1: index, %stride0: index, %stride1: index)
+      attributes {tf_entry} {
+    %c1 = arith.constant 1 : index
+    %216 = memref.reinterpret_cast %arg1 to offset: [0], sizes: [%size1, %size0], strides: [%size0, %c1]: memref<*xf32> to memref<?x?xf32, strided<[?, ?], offset: 0>>
+    %241 = memref.reinterpret_cast %arg2 to offset: [0], sizes: [%size1, %size0], strides: [%size0, %c1]: memref<*xf32> to memref<?x?xf32, strided<[?, ?], offset: 0>>
+    %304 = memref.reinterpret_cast %216 to offset: [0], sizes: [%size1, %size0], strides: [%stride1, %stride0]: memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, affine_map<(d0, d1)[s0, s1] -> (d0 * s0 + d1 * s1)>>
+    %309 = memref.reinterpret_cast %241 to offset: [0], sizes: [%size1, %size0], strides: [%stride0, %stride1]: memref<?x?xf32, strided<[?, ?], offset: 0>> to memref<?x?xf32, affine_map<(d0, d1)[s0, s1] -> (d0 * s0 + d1 * s1)>>
+    %310 = tf_framework.alloc(%arg0, %size1, %size0) : memref<?x?xf32>
+    gpu.launch_func  @AddV2_kernel_3::@AddV2_kernel blocks in (%c1, %c1, %c1) threads in (%c1, %c1, %c1) args(%size0 : index, %size1 : index, %310 : memref<?x?xf32>, %304 : memref<?x?xf32, affine_map<(d0, d1)[s0, s1] -> (d0 * s0 + d1 * s1)>>, %309 : memref<?x?xf32, affine_map<(d0, d1)[s0, s1] -> (d0 * s0 + d1 * s1)>>)
+    func.return
+  }
+
+  // CHECK-LABEL: gpu.module @AddV2_kernel_3
+  gpu.module @AddV2_kernel_3 {
+    // CHECK-LABEL: llvm.func @AddV2_kernel
+    // ABI-SAME: {llvm.align = 16 : index, llvm.noalias}
+    // ABI-SAME: {llvm.align = 16 : index}
+    // ABI-SAME: {llvm.align = 16 : index}
+    llvm.func @AddV2_kernel(%arg0: i64, %arg1: i64, %arg2: !llvm.ptr, %arg3: !llvm.ptr {llvm.align = 16 : index, llvm.noalias}, %arg4: i64, %arg5: i64, %arg6: i64, %arg7: i64, %arg8: i64, %arg9: !llvm.ptr, %arg10: !llvm.ptr {llvm.align = 16 : index}, %arg11: i64, %arg12: i64, %arg13: i64, %arg14: i64, %arg15: i64, %arg16: !llvm.ptr, %arg17: !llvm.ptr {llvm.align = 16 : index}, %arg18: i64, %arg19: i64, %arg20: i64, %arg21: i64, %arg22: i64) attributes {gpu.kernel} {
+      // ABI: %[[C0:.*]] = llvm.mlir.constant(0 : index) : i64
+      // ABI: %[[C1:.*]] = llvm.mlir.constant(1 : index) : i64
+      %0 = llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %1 = llvm.insertvalue %arg2, %0[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %2 = llvm.insertvalue %arg3, %1[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %3 = llvm.insertvalue %arg4, %2[2] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %4 = llvm.insertvalue %arg5, %3[3, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %5 = llvm.insertvalue %arg7, %4[4, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %6 = llvm.insertvalue %arg6, %5[3, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %7 = llvm.insertvalue %arg8, %6[4, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR0:.*]], %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR0]], %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C0]], %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[3, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C1]], %{{.*}}[4, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[SHP0:.*]], %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[STR0:.*]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[SHP1:.*]], %{{.*}}[3, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[STR1:.*]], %{{.*}}[4, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %8 = llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %9 = llvm.insertvalue %arg9, %8[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %10 = llvm.insertvalue %arg10, %9[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %11 = llvm.insertvalue %arg11, %10[2] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %12 = llvm.insertvalue %arg12, %11[3, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %13 = llvm.insertvalue %arg14, %12[4, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %14 = llvm.insertvalue %arg13, %13[3, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %15 = llvm.insertvalue %arg15, %14[4, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR0:.*]], %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR0]], %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C0]], %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NOT: llvm.insertvalue %[[C1]], %{{.*}}[3, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[SHP0]], %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NOT: llvm.insertvalue %[[STR0]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE: llvm.insertvalue %[[SHP1]], %{{.*}}[3, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NOT: llvm.insertvalue %[[STR1]], %{{.*}}[4, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %16 = llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %17 = llvm.insertvalue %arg16, %16[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %18 = llvm.insertvalue %arg17, %17[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %19 = llvm.insertvalue %arg18, %18[2] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %20 = llvm.insertvalue %arg19, %19[3, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %21 = llvm.insertvalue %arg21, %20[4, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %22 = llvm.insertvalue %arg20, %21[3, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      %23 = llvm.insertvalue %arg22, %22[4, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR0:.*]], %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR0]], %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C0]], %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[3, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // ABI-NOT: llvm.insertvalue %[[C1]], %{{.*}}[4, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[SHP0]], %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NOT: llvm.insertvalue %[[STR0]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE: llvm.insertvalue %[[SHP1]], %{{.*}}[3, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      // SHAPE-NOT: llvm.insertvalue %[[STR1]], %{{.*}}[4, 1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+      llvm.return
+      // CHECK: llvm.return
+    }
+  }
+}
+
+// -----
+
+
+// Test binary op with broadcasting a scalar.
+
+#map0 = affine_map<(d0)[s0] -> (d0 * s0)>
+
+// CHECK-LABEL: module attributes {gpu.container_module}
+module attributes {gpu.container_module} {
+  // CHECK-LABEL: func @add_one_scalar
+  func.func @add_one_scalar(%arg0: !tf_framework.op_kernel_context, %arg1: memref<*xf32>, %arg2: memref<*xf32>, %size: index)
+      attributes {tf_entry} {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %13 = memref.cast %arg1 : memref<*xf32> to memref<f32>
+    %26 = memref.reinterpret_cast %arg2 to offset: [0], sizes: [%size], strides: [%c1]: memref<*xf32> to memref<?xf32, strided<[?], offset: 0>>
+    %27 = memref.reinterpret_cast %13 to offset: [0], sizes: [%size], strides: [%c0]: memref<f32> to memref<?xf32, #map0>
+    %28 = memref.reinterpret_cast %26 to offset: [0], sizes: [%size], strides: [%c1]: memref<?xf32, strided<[?], offset: 0>> to memref<?xf32, #map0>
+    %29 = tf_framework.alloc(%arg0, %size) : memref<?xf32>
+    gpu.launch_func  @AddV2_kernel::@AddV2_kernel blocks in (%c1, %c1, %c1) threads in (%c1, %c1, %c1) args(%size : index, %29 : memref<?xf32>, %27 : memref<?xf32, #map0>, %28 : memref<?xf32, #map0>)
+    func.return
+  }
+  // CHECK-LABEL: gpu.module @AddV2_kernel
+  gpu.module @AddV2_kernel {
+    // CHECK-LABEL: llvm.func @AddV2_kernel
+    // ABI-SAME: {llvm.align = 16 : index, llvm.noalias}
+    // ABI-SAME: {llvm.align = 16 : index}
+    // ABI-SAME: {llvm.align = 16 : index}
+    llvm.func @AddV2_kernel(%arg0: i64, %arg1: !llvm.ptr, %arg2: !llvm.ptr, %arg3: i64, %arg4: i64, %arg5: i64, %arg6: !llvm.ptr, %arg7: !llvm.ptr, %arg8: i64, %arg9: i64, %arg10: i64, %arg11: !llvm.ptr, %arg12: !llvm.ptr, %arg13: i64, %arg14: i64, %arg15: i64) attributes {gpu.kernel} {
+      // ABI: %[[C0:.*]] = llvm.mlir.constant(0 : index) : i64
+      // ABI: %[[C1:.*]] = llvm.mlir.constant(1 : index) : i64
+      %0 = llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %1 = llvm.insertvalue %arg1, %0[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %2 = llvm.insertvalue %arg2, %1[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %3 = llvm.insertvalue %arg3, %2[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %4 = llvm.insertvalue %arg4, %3[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %5 = llvm.insertvalue %arg5, %4[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR0:.*]], %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR0]], %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C0]], %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C1]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[SHP:.*]], %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[STR:.*]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %6 = llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %7 = llvm.insertvalue %arg6, %6[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %8 = llvm.insertvalue %arg7, %7[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %9 = llvm.insertvalue %arg8, %8[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %10 = llvm.insertvalue %arg9, %9[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %11 = llvm.insertvalue %arg10, %10[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR1:.*]], %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR1]], %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C0]], %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C0]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[SHP]], %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NOT: llvm.insertvalue %[[STR]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %12 = llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %13 = llvm.insertvalue %arg11, %12[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %14 = llvm.insertvalue %arg12, %13[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %15 = llvm.insertvalue %arg13, %14[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %16 = llvm.insertvalue %arg14, %15[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      %17 = llvm.insertvalue %arg15, %16[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR2:.*]], %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[PTR2]], %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C0]], %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // ABI-NEXT: llvm.insertvalue %[[C1]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE: llvm.mlir.undef : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[1] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %{{.*}}, %{{.*}}[2] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NEXT: llvm.insertvalue %[[SHP]], %{{.*}}[3, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      // SHAPE-NOT: llvm.insertvalue %[[STR]], %{{.*}}[4, 0] : !llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>
+      llvm.return
+      // CHECK: llvm.return
+    }
+  }
+}
+

tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_framework_legalize_to_llvm.mlir

@@ -0,0 +1,263 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s -tf-kernel-to-llvm -reconcile-unrealized-casts -split-input-file | FileCheck %s
+
+// CHECK: llvm.func @_mlir_ciface_tf_alloc
+// CHECK-SAME:  (!llvm.ptr, i64, i64, i32, i32, !llvm.ptr) -> !llvm.ptr
+
+// CHECK-LABEL: llvm.func @alloc(
+// CHECK-SAME:    [[TF_CTX:%.*]]: !llvm.ptr,
+// CHECK-SAME:    [[SIZE_0:%.*]]: i64,
+// CHECK-SAME:    [[SIZE_2:%.*]]: i64) -> [[DESC_TY:!.*]] {
+func.func @alloc(%ctx: !tf_framework.op_kernel_context,
+                %size_0 : index , %size_2 : index) -> memref<?x10x?xf32> {
+  %buf = tf_framework.alloc(%ctx, %size_0, %size_2) : memref<?x10x?xf32>
+  func.return %buf : memref<?x10x?xf32>
+}
+// Compute number of elements.
+// CHECK: [[SIZE_1A:%.*]] = llvm.mlir.constant(10 : index) : i64
+// CHECK: [[SIZE_1B:%.*]] = llvm.mlir.constant(10 : index) : i64
+// CHECK: [[NUM_ELEM_0:%.*]] = llvm.mul [[SIZE_0]], [[SIZE_1B]] : i64
+// CHECK: [[NUM_ELEMS:%.*]] = llvm.mul [[NUM_ELEM_0]], [[SIZE_2]] : i64
+
+// Compute the size of an individual element.
+// CHECK: [[NULL:%.*]] = llvm.mlir.zero : !llvm.ptr
+// CHECK: [[GEP:%.*]] = llvm.getelementptr [[NULL]]{{\[}}1]
+// CHECK-SAME:            (!llvm.ptr) -> !llvm.ptr, f32
+// CHECK: [[SIZE_OF_FLOAT:%.*]] = llvm.ptrtoint [[GEP]]
+// CHECK-SAME:            !llvm.ptr to i64
+
+// Compute output index (-1) and candidate indices (0, NULL).
+// CHECK: [[OUTPUT_INDEX:%.*]] = llvm.mlir.constant(-1 : i32) : i32
+// CHECK-NEXT: [[NUM_CANDIDATES:%.*]] = llvm.mlir.constant(0 : i32) : i32
+// CHECK-NEXT: [[CANDIDATES_PTR:%.*]] = llvm.mlir.zero : !llvm.ptr
+
+// Allocate memory.
+// CHECK: [[BYTES_PTR:%.*]] = llvm.call @{{.*}}([[TF_CTX]], [[NUM_ELEMS]],
+// CHECK-SAME: [[SIZE_OF_FLOAT]], [[OUTPUT_INDEX]], [[NUM_CANDIDATES]],
+// CHECK-SAME: [[CANDIDATES_PTR]])
+
+// Build memref descriptor.
+// CHECK: [[DESC_0:%.*]] = llvm.mlir.poison : [[DESC_TY]]
+
+// Set pointers and offset.
+// CHECK: [[DESC_1:%.*]] = llvm.insertvalue [[BYTES_PTR]], [[DESC_0]][0]
+// CHECK: [[DESC_2:%.*]] = llvm.insertvalue [[BYTES_PTR]], [[DESC_1]][1]
+// CHECK: [[C0:%.*]] = llvm.mlir.constant(0 : index) : i64
+// CHECK: [[DESC_3:%.*]] = llvm.insertvalue [[C0]], [[DESC_2]][2] : [[DESC_TY]]
+
+// Set sizes and strides.
+// CHECK: [[STRIDE_2:%.*]] = llvm.mlir.constant(1 : index) : i64
+// CHECK: [[DESC_4:%.*]] = llvm.insertvalue [[SIZE_2]], [[DESC_3]][3, 2]
+// CHECK: [[DESC_5:%.*]] = llvm.insertvalue [[STRIDE_2]], [[DESC_4]][4, 2]
+// CHECK: [[STRIDE_1:%.*]] = llvm.mul [[STRIDE_2]], [[SIZE_2]] : i64
+// CHECK: [[DESC_6:%.*]] = llvm.insertvalue [[SIZE_1A]], [[DESC_5]][3, 1]
+// CHECK: [[DESC_7:%.*]] = llvm.insertvalue [[STRIDE_1]], [[DESC_6]][4, 1]
+// CHECK: [[STRIDE_0:%.*]] = llvm.mul [[STRIDE_1]], [[SIZE_1A]] : i64
+// CHECK: [[DESC_8:%.*]] = llvm.insertvalue [[SIZE_0]], [[DESC_7]][3, 0]
+// CHECK: [[DESC_9:%.*]] = llvm.insertvalue [[STRIDE_0]], [[DESC_8]][4, 0]
+// CHECK: llvm.return [[DESC_9]] : [[DESC_TY]]
+
+// -----
+
+// CHECK: llvm.func @_mlir_ciface_tf_dealloc(!llvm.ptr, !llvm.ptr)
+
+// CHECK-LABEL: llvm.func @dealloc(
+// CHECK-SAME:    [[TF_CTX:%[a-z0-9]*]]: !llvm.ptr
+func.func @dealloc(%ctx: !tf_framework.op_kernel_context,
+                  %memref : memref<?x10xf32>) {
+  tf_framework.dealloc(%ctx, %memref) : memref<?x10xf32>
+  func.return
+}
+// Extract allocated ptr from the memref descriptor.
+// CHECK: %{{.*}} = llvm.mlir.poison : [[DESC_TY:!.*]]
+// CHECK: [[FLOAT_PTR:%.*]] = llvm.extractvalue %{{.*}}[0] : [[DESC_TY]]
+
+// Deallocate.
+// CHECK: llvm.call @_mlir_ciface_tf_dealloc(
+// CHECK-SAME: [[TF_CTX]], [[FLOAT_PTR]]) : (!llvm.ptr, !llvm.ptr) -> ()
+
+// -----
+
+// CHECK-LABEL: llvm.func @_mlir_ciface_tf_report_error(!llvm.ptr, i32, !llvm.ptr)
+// CHECK: llvm.mlir.global internal constant [[MSG_CONST:@error_message_[0-9]+]]("Everything is awesome\00")
+
+func.func @report_error(%ctx: !tf_framework.op_kernel_context) {
+  tf_framework.report_error %ctx, "INVALID_ARGUMENT", "Everything is awesome" loc(unknown)
+  func.return
+}
+// CHECK:     llvm.func @report_error([[CTX:%.*]]: !llvm.ptr)
+// CHECK-NEXT:  [[ADDR:%.*]] = llvm.mlir.addressof [[MSG_CONST]]
+// CHECK:       [[MSG:%.*]] = llvm.getelementptr [[ADDR]]
+// CHECK:       [[CODE:%.*]] = llvm.mlir.constant({{.*}}) : i32
+// CHECK:       llvm.call @{{.*}}_tf_report_error([[CTX]], [[CODE]], [[MSG]])
+
+// -----
+
+// CHECK-LABEL: llvm.func @unranked_null_memref()
+func.func @unranked_null_memref() {
+  %null = tf_framework.null_memref : memref<*xf32>
+  func.return
+}
+// CHECK: [[C0:%.*]] = llvm.mlir.constant(0 : index) : i64
+// CHECK: [[DESC_0:%.*]] = llvm.mlir.poison : !llvm.struct<(i64, ptr)>
+// CHECK: [[DESC_1:%.*]] = llvm.insertvalue [[C0]], [[DESC_0]][0]
+// CHECK: [[PTR:%.*]] = llvm.alloca {{.*}} x i8
+// CHECK: [[DESC_2:%.*]] = llvm.insertvalue [[PTR]], [[DESC_1]][1]
+
+// -----
+
+// CHECK-LABEL: llvm.func @ranked_null_memref()
+func.func @ranked_null_memref() {
+  %null = tf_framework.null_memref : memref<2x?xf32>
+  func.return
+}
+// CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : index) : i64
+// CHECK-NEXT: %[[C1:.*]] = llvm.mlir.constant(1 : index) : i64
+// CHECK-NEXT: %[[C2:.*]] = llvm.mlir.constant(2 : index) : i64
+// CHECK-NEXT: %[[C1_:.*]] = llvm.mlir.constant(1 : index) : i64
+
+// CHECK: llvm.mlir.zero
+// CHECK: %[[NULL:.*]] = llvm.mlir.zero : !llvm.ptr
+// CHECK-NEXT: %[[DESC_0:.*]] = llvm.mlir.poison :
+// CHECK-SAME:   !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
+// CHECK-NEXT: %[[DESC_1:.*]] = llvm.insertvalue %[[NULL]], %[[DESC_0]][0]
+// CHECK-NEXT: %[[DESC_2:.*]] = llvm.insertvalue %[[NULL]], %[[DESC_1]][1]
+// CHECK-NEXT: %[[DESC_3:.*]] = llvm.insertvalue %[[C0]], %[[DESC_2]][2]
+// CHECK-NEXT: %[[DESC_4:.*]] = llvm.insertvalue %[[C2]], %[[DESC_3]][3, 0]
+// CHECK-NEXT: %[[DESC_5:.*]] = llvm.insertvalue %[[C1]], %[[DESC_4]][4, 0]
+// CHECK-NEXT: %[[DESC_6:.*]] = llvm.insertvalue %[[C1]], %[[DESC_5]][3, 1]
+// CHECK-NEXT: %[[DESC_7:.*]] = llvm.insertvalue %[[C1_]], %[[DESC_6]][4, 1]
+
+// -----
+
+// CHECK-LABEL: llvm.func @is_valid_memref
+func.func @is_valid_memref(%buf: memref<?xf32>) -> i1 {
+  %pred = tf_framework.is_valid_memref(%buf) : memref<?xf32> -> i1
+  func.return %pred : i1
+}
+// CHECK: %[[MEMREF:.*]] = llvm.insertvalue %{{.*}}, %{{.*}}[4, 0]
+
+// CHECK-NEXT: %[[IS_EMPTY:.*]] = llvm.mlir.constant(false) : i1
+// CHECK-NEXT: %[[C0:.*]] = llvm.mlir.constant(0 : index) : i64
+// CHECK-NEXT: %[[SIZE:.*]] = llvm.extractvalue %[[MEMREF]][3, 0]
+// CHECK-NEXT: %[[IS_ZERO:.*]] = llvm.icmp "eq" %[[SIZE]], %[[C0]] : i64
+// CHECK-NEXT: %[[IS_EMPTY_:.*]] =  llvm.or %[[IS_EMPTY]], %[[IS_ZERO]] : i1
+
+// CHECK-NEXT: %[[PTR_F32:.*]] = llvm.extractvalue %[[MEMREF]][0]
+// CHECK-NEXT: %[[NULL_PTR:.*]] = llvm.mlir.zero : !llvm.ptr
+// CHECK-NEXT: %[[NOT_NULL:.*]] = llvm.icmp "ne" %[[PTR_F32]], %[[NULL_PTR]]
+
+// CHECK-NEXT: %[[PRED:.*]] = llvm.or %[[NOT_NULL]], %[[IS_EMPTY_]]  : i1
+// CHECK-NEXT: llvm.return %[[PRED]]
+
+// -----
+
+// CHECK-LABEL: llvm.func @_mlir_ciface_tf_jit_compile(!llvm.ptr, !llvm.ptr, i64, !llvm.ptr, i64, !llvm.ptr, i1, i1, i1) -> !llvm.ptr
+// CHECK: llvm.mlir.global internal constant @[[CODE:jit_module_code_[0-9]+]]("placeholder\00")
+
+// CHECK: @jit_compile_from_str(%[[CTX:.*]]: !llvm.ptr)
+func.func @jit_compile_from_str(%ctx: !tf_framework.op_kernel_context)
+    -> !tf_framework.jit_callable {
+  // CHECK: %[[ADDR:.*]] = llvm.mlir.addressof @[[CODE]]
+  // CHECK: %[[CODE_PTR:.*]] = llvm.getelementptr %[[ADDR]][0, 0]
+
+  // Create stack-allocated array for the tile sizes.
+  // CHECK: %[[NUM_TILE_SIZES:.*]] = llvm.mlir.constant(3 : i64)
+  // CHECK: %[[TILE_SIZES:.*]] = llvm.alloca %[[NUM_TILE_SIZES]] x i64
+  // CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : i64)
+  // CHECK: %[[PTR:.*]] = llvm.getelementptr %[[TILE_SIZES]][%[[C0]]]
+  // CHECK: %[[C1:.*]] = llvm.mlir.constant(1 : i64)
+  // CHECK: llvm.store %[[C1]], %[[PTR]]
+  // CHECK: %[[C1:.*]] = llvm.mlir.constant(1 : i64)
+  // CHECK: %[[PTR:.*]] = llvm.getelementptr %[[TILE_SIZES]][%[[C1]]]
+  // CHECK: %[[C2:.*]] = llvm.mlir.constant(2 : i64)
+  // CHECK: llvm.store %[[C2]], %[[PTR]]
+  // CHECK: %[[C2:.*]] = llvm.mlir.constant(2 : i64)
+  // CHECK: %[[PTR:.*]] = llvm.getelementptr %[[TILE_SIZES]][%[[C2]]]
+  // CHECK: %[[C3:.*]] = llvm.mlir.constant(3 : i64)
+  // CHECK: llvm.store %[[C3]], %[[PTR]]
+
+  // Create stack-allocated array for the unroll factors.
+  // CHECK: %[[NUM_UNROLL_FACTORS:.*]] = llvm.mlir.constant(1 : i64) : i64
+  // CHECK: %[[UNROLL_FACTORS:.*]] = llvm.alloca %[[NUM_UNROLL_FACTORS]] x i64
+  // CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : i64)
+  // CHECK: %[[PTR:.*]] = llvm.getelementptr %[[UNROLL_FACTORS]][%[[C0]]]
+  // CHECK: %[[C4:.*]] = llvm.mlir.constant(4 : i64)
+  // CHECK: llvm.store %[[C4]], %[[PTR]]
+
+  // CHECK-DAG: %[[ENABLE_FTZ:.*]] = llvm.mlir.constant(false)
+  // CHECK-DAG: %[[CPU_CODEGEN:.*]] = llvm.mlir.constant(false)
+  // CHECK: %[[RES:.*]] = llvm.call @_mlir_ciface_tf_jit_compile
+  // CHECK-SAME: %[[CTX]], %[[CODE_PTR]],
+  // CHECK-SAME: %[[NUM_TILE_SIZES]], %[[TILE_SIZES]],
+  // CHECK-SAME: %[[NUM_UNROLL_FACTORS]], %[[UNROLL_FACTORS]],
+  // CHECK-SAME: %[[ENABLE_FTZ]], %[[CPU_CODEGEN]]
+  // CHECK: llvm.return %[[RES]]
+  %0 = tf_framework.jit_compile_from_str %ctx, "placeholder" {
+      tileSizes = [1, 2, 3], unrollFactors = [4],
+      enableFtz = false, index64Bit = false, cpuCodegen = false }
+  func.return %0 : !tf_framework.jit_callable
+}
+
+// -----
+
+// CHECK-LABEL: llvm.func @_mlir_ciface_tf_jit_execute(!llvm.ptr, !llvm.ptr, !llvm.ptr, i64, !llvm.ptr)
+
+// CHECK:      @jit_execute
+// CHECK-SAME: (%[[CTX:.*]]: !llvm.ptr, %[[CALLABLE:.*]]: !llvm.ptr, %[[RANK:.*]]: i64, %[[ARG_DESCR:.*]]: !llvm.ptr)
+func.func @jit_execute(%ctx: !tf_framework.op_kernel_context,
+    %callable : !tf_framework.jit_callable, %arg : memref<*xf32>)
+    -> memref<*xf32> {
+  // CHECK: %[[T0:.*]] = llvm.mlir.poison
+  // CHECK: %[[T1:.*]] = llvm.insertvalue %[[RANK]], %[[T0]][0]
+  // CHECK: %[[ARG:.*]] = llvm.insertvalue %[[ARG_DESCR]], %[[T1]][1]
+  // CHECK: %[[C1:.*]] = llvm.mlir.constant(1 : i64)
+  // CHECK: %[[RESULT_PTR:.*]] = llvm.alloca %[[C1]] x !llvm.struct<(i64, ptr)>
+
+  // Copy argument(s) to stack-allocated buffer.
+  // CHECK: %[[NUM_ARGS:.*]] = llvm.mlir.constant(1 : i64)
+  // CHECK: %[[ARGS_PTR:.*]] = llvm.alloca %[[NUM_ARGS]] x !llvm.struct<(i64, ptr)>
+  // CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : i64)
+  // CHECK: %[[ARGS0_PTR:.*]] = llvm.getelementptr %[[ARGS_PTR]][%[[C0]]]
+  // CHECK: llvm.store %[[ARG]], %[[ARGS0_PTR]]
+  // CHECK: llvm.call @_mlir_ciface_tf_jit_execute(%[[CTX]], %[[CALLABLE]], %[[RESULT_PTR]], %[[NUM_ARGS]], %[[ARGS_PTR]])
+  // CHECK: %[[RESULT:.*]] = llvm.load %[[RESULT_PTR]]
+
+  // Copy unranked memref descriptor to stack-allocated memory.
+  // ...
+  // CHECK: %[[STACK_RESULT_DESCR:.*]] = llvm.alloca %[[RESULT_DESCR_SIZE:[0-9]*]] x i8
+  // CHECK: %[[RESULT_DESCR:.*]] = llvm.extractvalue %[[RESULT]][1]
+  // CHECK: "llvm.intr.memcpy"(%[[STACK_RESULT_DESCR]], %[[RESULT_DESCR]], %[[RESULT_DESCR_SIZE]]) <{isVolatile = false}>
+  // CHECK: llvm.call @free(%[[RESULT_DESCR]])
+  // CHECK: %[[T0:.*]] = llvm.mlir.poison
+  // CHECK: %[[RANK:.*]] = llvm.extractvalue %[[RESULT]][0]
+  // CHECK: %[[T1:.*]] = llvm.insertvalue %[[RANK]], %[[T0]][0]
+  // CHECK: %[[RESULT:.*]] = llvm.insertvalue %[[STACK_RESULT_DESCR]], %[[T1]][1]
+
+  // Copy unranked memref descriptor to heap-allocated memory for return.
+  // ...
+  // CHECK: %[[HEAP_RESULT_DESCR:.*]] = llvm.call @malloc(%[[RESULT_DESCR_SIZE:[0-9]*]])
+  // CHECK: %[[STACK_RESULT_DESCR:.*]] = llvm.extractvalue %[[RESULT]][1]
+  // CHECK: "llvm.intr.memcpy"(%[[HEAP_RESULT_DESCR]], %[[STACK_RESULT_DESCR]], %[[RESULT_DESCR_SIZE]]) <{isVolatile = false}>
+  // CHECK: %[[T0:.*]] = llvm.mlir.poison
+  // CHECK: %[[RANK:.*]] = llvm.extractvalue %[[RESULT]][0]
+  // CHECK: %[[T1:.*]] = llvm.insertvalue %[[RANK]], %[[T0]][0]
+  // CHECK: %[[RESULT:.*]] = llvm.insertvalue %[[HEAP_RESULT_DESCR]], %[[T1]][1]
+  // CHECK: llvm.return %[[RESULT]]
+  %0 = tf_framework.jit_execute ctx(%ctx) %callable(%arg)
+      : memref<*xf32> -> memref<*xf32>
+  func.return %0 : memref<*xf32>
+}

tensorflow/compiler/mlir/tools/kernel_gen/tests/tf_kernel_gpu_launch_to_llvm.mlir

@@ -0,0 +1,50 @@
+// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ==============================================================================
+// RUN: kernel-gen-opt %s -tf-kernel-to-llvm -reconcile-unrealized-casts -split-input-file | FileCheck %s --dump-input=always
+
+// CHECK-LABEL: module @main
+module @main attributes {gpu.container_module} {
+
+// CHECK-NOT: gpu.module @kernel_module
+gpu.module @kernel_module attributes {gpu.binary_blob = "BLOB!"} {
+  llvm.func @the_kernel() attributes {gpu.kernel} {
+    llvm.return
+  }
+}
+
+// CHECK: llvm.func @_mlir_ciface_tf_launch_kernel(!llvm.ptr, !llvm.ptr, !llvm.ptr, i64, i64, i64, i64, i64, i64, !llvm.ptr)
+// CHECK-DAG: llvm.mlir.global internal constant @kernel_module_the_kernel_kernel_name("the_kernel\00")
+// CHECK-DAG: llvm.mlir.global internal constant @kernel_module_blob("BLOB!")
+
+// CHECK-LABEL: llvm.func @launch
+// CHECK-SAME: (%[[CTX:.*]]: !llvm.ptr, %{{.*}}: !llvm.ptr, %{{.*}}: !llvm.ptr, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64, %{{.*}}: i64
+func.func @launch(%ctx: !tf_framework.op_kernel_context, %memref: memref<?x10xf32>) {
+  // CHECK: %[[C1:.*]] = llvm.mlir.constant(1 : index) : i64
+  // CHECK: %[[BLOB:.*]] = llvm.mlir.addressof @kernel_module_blob : !llvm.ptr
+  // CHECK: %[[BLOB_PTR:.*]] = llvm.getelementptr %[[BLOB]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<5 x i8>
+  // CHECK: %[[NAME:.*]] = llvm.mlir.addressof @kernel_module_the_kernel_kernel_name : !llvm.ptr
+  // CHECK: %[[NAME_PTR:.*]] = llvm.getelementptr %[[NAME]][0, 0] : (!llvm.ptr) -> !llvm.ptr, !llvm.array<11 x i8>
+  // CHECK: %[[C7:.*]] = llvm.mlir.constant(7 : i32) : i32
+  // CHECK: %[[ARGS:.*]] = llvm.alloca %22 x !llvm.ptr : (i32) -> !llvm.ptr
+  // CHECK: llvm.call @_mlir_ciface_tf_launch_kernel(%[[CTX]], %[[BLOB_PTR]], %[[NAME_PTR]], %[[C1]], %[[C1]], %[[C1]], %[[C1]], %[[C1]], %[[C1]], %[[ARGS]])
+  %c1 = arith.constant 1 : index
+  gpu.launch_func  @kernel_module::@the_kernel
+      blocks in (%c1, %c1, %c1)
+      threads in (%c1, %c1, %c1)
+      args(%memref: memref<?x10xf32>)
+  func.return
+}
+
+}

tensorflow/compiler/mlir/tools/kernel_gen/tf_framework_c_interface.cc

@@ -0,0 +1,371 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tools/kernel_gen/tf_framework_c_interface.h"
+
+#include <cstddef>
+#include <string>
+#include <utility>
+
+#include "absl/log/check.h"
+#include "absl/log/log.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ExecutionEngine/JITSymbol.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/Mangling.h"
+#include "llvm/ExecutionEngine/Orc/Shared/ExecutorAddress.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/TargetSelect.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/MemRef/Transforms/AllocationOpInterfaceImpl.h"  // from @llvm-project
+#include "mlir/ExecutionEngine/CRunnerUtils.h"  // from @llvm-project
+#include "mlir/ExecutionEngine/ExecutionEngine.h"  // from @llvm-project
+#include "mlir/ExecutionEngine/OptUtils.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OwningOpRef.h"  // from @llvm-project
+#include "mlir/Parser/Parser.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/compile_cache_item.pb.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/kernel_creator.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/tf_jit_cache.h"
+#include "xla/stream_executor/stream.h"
+#include "xla/tsl/framework/allocator.h"
+#include "tensorflow/core/framework/allocator.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/env.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/refcount.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/statusor.h"
+
+#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
+#include <optional>
+
+#include "tensorflow/compiler/mlir/tools/kernel_gen/tf_gpu_runtime_wrappers.h"
+#endif
+
+static constexpr absl::string_view kTFJitCacheDirEnvVar = "TF_JIT_CACHE_DIR";
+
+namespace mlir {
+namespace kernel_gen {
+namespace tf_framework {
+namespace {
+
+using tensorflow::Allocator;
+using tensorflow::AllocatorAttributes;
+
+Allocator* GetAllocator(void* op_kernel_ctx) {
+  auto* ctx = static_cast<tensorflow::OpKernelContext*>(op_kernel_ctx);
+  // TODO(pifon): Figure out how to set AllocatorAttributes correctly.
+  AllocatorAttributes attrs;
+  return ctx->get_allocator(attrs);
+}
+
+}  // namespace
+
+extern "C" void* _mlir_ciface_tf_alloc(void* op_kernel_ctx, size_t num_elements,
+                                       size_t element_size,
+                                       int32_t output_index,
+                                       int32_t num_candidates,
+                                       int32_t* candidate_input_indices) {
+  static constexpr int kAmbiguousOutputIndex = -1;
+  auto* ctx = static_cast<tensorflow::OpKernelContext*>(op_kernel_ctx);
+  if (output_index != kAmbiguousOutputIndex) {
+    // Create a 1D shape, because the shapes don't have to match exactly for
+    // input forwarding. Only the number of elements must be the same.
+    tensorflow::TensorShape output_shape;
+    output_shape.AddDim(num_elements);
+
+    // Iterate over indices of all inputs that can potentially be used for
+    // forwarding.
+    for (int i = 0; i < num_candidates; ++i) {
+      auto tensor = ctx->forward_input(candidate_input_indices[i], output_index,
+                                       ctx->expected_output_dtype(output_index),
+                                       output_shape,
+                                       ctx->output_memory_type(output_index),
+                                       ctx->output_alloc_attr(output_index));
+      if (tensor != nullptr) {
+        return tensor->data();
+      }
+    }
+
+    CHECK(!ctx->output_expects_forwarding(output_index));
+  }
+
+  // If no forwarding happened, allocate a chunk of memory.
+  return GetAllocator(op_kernel_ctx)
+      ->AllocateRaw(Allocator::kAllocatorAlignment,
+                    num_elements * element_size);
+}
+
+extern "C" void _mlir_ciface_tf_dealloc(void* op_kernel_ctx, void* ptr) {
+  GetAllocator(op_kernel_ctx)->DeallocateRaw(ptr);
+}
+
+extern "C" void _mlir_ciface_tf_report_error(void* op_kernel_ctx,
+                                             int32_t error_code, char* msg) {
+  std::optional<ErrorCode> symbol = symbolizeErrorCode(error_code);
+  if (!symbol.has_value()) {
+    LOG(ERROR) << "No valid conversion from integer value = " << error_code
+               << "to ErrorCode attribute";
+    return;
+  }
+  auto* ctx = static_cast<tensorflow::OpKernelContext*>(op_kernel_ctx);
+  ctx->CtxFailureWithWarning(
+      absl::Status{ConvertAttrToEnumValue(symbol.value()), msg});
+}
+
+static void ReportError(void* op_kernel_ctx, ErrorCode error_code,
+                        const char* msg) {
+  _mlir_ciface_tf_report_error(op_kernel_ctx, static_cast<uint32_t>(error_code),
+                               const_cast<char*>(msg));
+}
+
+namespace {
+
+std::string GetFileCachePath(const std::string cache_dir,
+                             const std::string& code) {
+  size_t hash = llvm::hash_value(code);
+  return tensorflow::io::JoinPath(cache_dir, std::to_string(hash));
+}
+
+// A callback to register all externally defined symbols needed by the kernel.
+llvm::orc::SymbolMap TFFrameworkSymbolMap(llvm::orc::MangleAndInterner mangle) {
+  llvm::orc::SymbolMap symbol_map;
+  auto bind = [&](llvm::StringRef name, auto symbol_ptr) {
+    symbol_map[mangle(name)] = {llvm::orc::ExecutorAddr::fromPtr(symbol_ptr),
+                                llvm::JITSymbolFlags()};
+  };
+
+  // Register TF framework symbols.
+  bind("_mlir_ciface_tf_alloc", &_mlir_ciface_tf_alloc);
+  bind("_mlir_ciface_tf_dealloc", &_mlir_ciface_tf_dealloc);
+  bind("_mlir_ciface_tf_report_error", &_mlir_ciface_tf_report_error);
+#if defined(GOOGLE_CUDA) || defined(TENSORFLOW_USE_ROCM)
+  bind("_mlir_ciface_tf_launch_kernel", &_mlir_ciface_tf_launch_kernel);
+#endif
+
+  // Register malloc/free to avoid unexpected implementations from shared libs.
+  bind("malloc", &malloc);
+  bind("free", &free);
+
+  return symbol_map;
+}
+
+void InitializeLlvmCompiler() {
+  static const bool initialized = ([] {
+    llvm::InitializeNativeTarget();
+    llvm::InitializeNativeTargetAsmPrinter();
+    return true;
+  })();
+  (void)initialized;
+}
+
+llvm::Expected<std::unique_ptr<ExecutionEngine>> Compile(
+    const std::string code, llvm::SmallVectorImpl<std::string>& architectures,
+    llvm::SmallVectorImpl<int64_t>& tile_sizes,
+    llvm::SmallVectorImpl<int64_t>& unroll_factors, bool enable_ftz,
+    bool index_64bit) {
+  std::string cache_dir;
+  if (const char* dir = getenv(kTFJitCacheDirEnvVar.data())) {
+    cache_dir = dir;
+  }
+
+  // Check if we already have a partially compiled module in the filesystem
+  // based cache.
+  CompilationCacheItem item;
+  auto tenv = tensorflow::Env::Default();
+  if (!cache_dir.empty() && tenv->RecursivelyCreateDir(cache_dir).ok()) {
+    std::string data;
+    if (tensorflow::ReadFileToString(tenv, GetFileCachePath(cache_dir, code),
+                                     &data)
+            .ok()) {
+      item.ParseFromString(data);
+      if (item.original_module() != code) {
+        item.Clear();
+      }
+    }
+  }
+
+  // Create the kernel.
+  mlir::DialectRegistry registry;
+  mlir::memref::registerAllocationOpInterfaceExternalModels(registry);
+  mlir::MLIRContext context(registry);
+
+  mlir::OwningOpRef<mlir::ModuleOp> module;
+
+  if (item.result_module().empty()) {
+    // Otherwise, compile the module now.
+    absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> status_or_module =
+        tensorflow::kernel_gen::GenerateKernelForHloCode(
+            context, code, architectures, tile_sizes, unroll_factors,
+            /*print_ptx=*/false, /*print_llvmir=*/false, enable_ftz,
+            index_64bit,
+            /*jit_compile=*/false,
+            /*jit_i64_indexed_for_large_tensors=*/false,
+            /*apply_cl_options=*/false);
+    if (!status_or_module.ok()) {
+      LOG(ERROR) << status_or_module.status();
+      return nullptr;
+    }
+    module = std::move(status_or_module.value());
+
+    if (!cache_dir.empty() && tenv->RecursivelyCreateDir(cache_dir).ok()) {
+      // Save the compilation result here for future processes to use.
+      item.set_original_module(code);
+      llvm::raw_string_ostream stream(*item.mutable_result_module());
+      module.get().print(stream);
+      stream.flush();
+
+      tensorflow::WriteStringToFile(tenv, GetFileCachePath(cache_dir, code),
+                                    item.SerializeAsString())
+          .IgnoreError();
+    }
+  } else {
+    module = tensorflow::kernel_gen::SetupContextAndParseModule(
+                 context, item.result_module())
+                 .value();
+  }
+
+  // Initialize LLVM targets.
+  InitializeLlvmCompiler();
+
+  // Create execution engine with an inner optimization pipeline.
+  auto opt_pipeline = mlir::makeOptimizingTransformer(
+      /*optLevel=*/2, /*sizeLevel=*/0, /*targetMachine=*/nullptr);
+  mlir::ExecutionEngineOptions engine_options;
+  engine_options.transformer = opt_pipeline;
+  llvm::Expected<std::unique_ptr<ExecutionEngine>> engine =
+      mlir::ExecutionEngine::create(module.get(), engine_options);
+  if (!engine) {
+    LOG(ERROR) << "Failed to create ExecutionEngine: "
+               << toString(engine.takeError());
+    return nullptr;
+  }
+
+  // Finally, register the missing symbols.
+  engine.get()->registerSymbols(TFFrameworkSymbolMap);
+  return engine;
+}
+
+template <typename T, typename U = T>
+llvm::SmallVector<T, 8> SmallVectorFromCArray(int64_t num_elements,
+                                              U* elements_ptr) {
+  llvm::SmallVector<T, 8> result;
+  result.reserve(num_elements);
+  for (int i = 0; i < num_elements; ++i) result.push_back(elements_ptr[i]);
+  return result;
+}
+
+}  // namespace
+
+extern "C" void* _mlir_ciface_tf_jit_compile(
+    void* op_kernel_ctx, char* code, int64_t num_tile_sizes,
+    int64_t* tile_sizes_ptr, int64_t num_unroll_factors,
+    int64_t* unroll_factors_ptr, bool enable_ftz, bool index_64bit) {
+  // Get the resource manager.
+  auto* ctx = static_cast<tensorflow::OpKernelContext*>(op_kernel_ctx);
+  tensorflow::ResourceMgr* rm = ctx->resource_manager();
+  if (!rm) {
+    ReportError(op_kernel_ctx, ErrorCode::UNKNOWN, "No resource manager.");
+    return nullptr;
+  }
+
+  // Get the JIT cache.
+  JITCache* jit_cache = nullptr;
+  auto status = rm->LookupOrCreate<JITCache>(rm->default_container(),
+                                             JITCache::kDefaultResourceName,
+                                             &jit_cache, JITCache::Create);
+  tensorflow::core::ScopedUnref jit_cache_ref(jit_cache);
+  if (!status.ok()) {
+    ReportError(op_kernel_ctx, ErrorCode::UNKNOWN,
+                "Failed to find or create JIT cache.");
+    return nullptr;
+  }
+
+  // Determine the unique architecture for the current GPU, if any.
+  SmallVector<std::string, 1> architectures;
+#if defined(GOOGLE_CUDA)
+  stream_executor::CudaComputeCapability cc =
+      ctx->op_device_context()->stream()->GetCudaComputeCapability();
+  architectures.push_back(absl::StrCat("sm_", cc.major, cc.minor));
+#elif defined(TENSORFLOW_USE_ROCM)
+  stream_executor::RocmComputeCapability cc =
+      ctx->op_device_context()->stream()->GetRocmComputeCapability();
+  architectures.push_back(cc.gcn_arch_name());
+#endif
+
+  // Construct `SmallVector`s from arguments.
+  llvm::SmallVector<int64_t, 8> tile_sizes =
+      SmallVectorFromCArray<int64_t>(num_tile_sizes, tile_sizes_ptr);
+  llvm::SmallVector<int64_t, 8> unroll_factors =
+      SmallVectorFromCArray<int64_t>(num_unroll_factors, unroll_factors_ptr);
+
+  // Lookup or compile the execution module.
+  ExecutionEngine* engine = jit_cache->LookupOrCompile(code, [&]() {
+    return Compile(code, architectures, tile_sizes, unroll_factors, enable_ftz,
+                   index_64bit);
+  });
+  if (engine == nullptr) {
+    ReportError(op_kernel_ctx, ErrorCode::UNKNOWN, "JIT compilation failed.");
+    return nullptr;
+  }
+  return engine;
+}
+
+extern "C" void _mlir_ciface_tf_jit_execute(void* op_kernel_ctx, void* callable,
+                                            void* result, int64_t num_args,
+                                            void* args_ptr) {
+  // JIT compilation must have failed earlier if there is no callable ptr.
+  // Return some empty memory descriptor to prevent a crash.
+  if (callable == nullptr) {
+    auto* desc = static_cast<::UnrankedMemRefType<void>*>(result);
+    desc->rank = 0;
+    auto* inner_desc = static_cast<StridedMemRefType<int8_t, 0>*>(
+        malloc(sizeof(StridedMemRefType<int8_t, 0>)));
+    inner_desc->basePtr = nullptr;
+    inner_desc->data = nullptr;
+    inner_desc->offset = 0;
+    desc->descriptor = inner_desc;
+    return;
+  }
+
+  // Build the argument array according to `ExecutionEngine`'s calling
+  // convention.
+  auto* typed_args_ptr = static_cast<::UnrankedMemRefType<void>*>(args_ptr);
+  llvm::SmallVector<void*, 8> args_array = {&op_kernel_ctx};
+  for (int i = 0; i < num_args; i++) {
+    auto& desc = typed_args_ptr[i];
+    args_array.push_back(&desc.rank);
+    args_array.push_back(&desc.descriptor);
+  }
+  args_array.push_back(result);
+
+  llvm::Error invocation_result =
+      static_cast<ExecutionEngine*>(callable)->invokePacked("main", args_array);
+  if (invocation_result)
+    ReportError(op_kernel_ctx, ErrorCode::UNKNOWN, "JIT invocation failed.");
+}
+
+}  // namespace tf_framework
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/tf_framework_c_interface.h

@@ -0,0 +1,52 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TF_FRAMEWORK_C_INTERFACE_H_
+#define TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TF_FRAMEWORK_C_INTERFACE_H_
+
+#include <cstddef>
+#include <cstdint>
+
+#include "mlir/ExecutionEngine/RunnerUtils.h"  // from @llvm-project
+
+namespace mlir {
+namespace kernel_gen {
+namespace tf_framework {
+
+extern "C" MLIR_RUNNERUTILS_EXPORT void* _mlir_ciface_tf_alloc(
+    void* op_kernel_ctx, size_t num_elements, size_t element_size,
+    int32_t output_index, int32_t num_candidates,
+    int32_t* candidate_input_indices);
+
+extern "C" MLIR_RUNNERUTILS_EXPORT void _mlir_ciface_tf_dealloc(
+    void* op_kernel_ctx, void* ptr);
+
+extern "C" MLIR_RUNNERUTILS_EXPORT void _mlir_ciface_tf_report_error(
+    void* op_kernel_ctx, int32_t error_code, char* msg);
+
+extern "C" MLIR_RUNNERUTILS_EXPORT void* _mlir_ciface_tf_jit_compile(
+    void* op_kernel_ctx, char* code, int64_t num_tile_sizes,
+    int64_t* tile_sizes_ptr, int64_t num_unroll_factors,
+    int64_t* unroll_factors_ptr, bool enable_ftz, bool index_64bit);
+
+extern "C" MLIR_RUNNERUTILS_EXPORT void _mlir_ciface_tf_jit_execute(
+    void* op_kernel_ctx, void* callable, void* result, int64_t num_args,
+    void* args_ptr);
+
+}  // namespace tf_framework
+}  // namespace kernel_gen
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TF_FRAMEWORK_C_INTERFACE_H_

tensorflow/compiler/mlir/tools/kernel_gen/tf_gpu_runtime_wrappers.cc

@@ -0,0 +1,187 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tools/kernel_gen/tf_gpu_runtime_wrappers.h"
+
+#include <cassert>
+#include <cstdint>
+#include <string>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "absl/strings/str_cat.h"
+#include "xla/stream_executor/stream.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/op_requires.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/refcount.h"
+#include "tensorflow/core/platform/status.h"
+
+#if GOOGLE_CUDA
+#include "third_party/gpus/cuda/include/cuda.h"
+#endif
+
+static void ReportInternalError(tensorflow::OpKernelContext *ctx,
+                                const std::string msg) {
+  if (ctx == nullptr) {
+    LOG(WARNING) << msg << "\n";
+    return;
+  }
+  ctx->CtxFailureWithWarning(absl::Status{absl::StatusCode::kInternal, msg});
+}
+
+#if GOOGLE_CUDA
+using GPUResult = CUresult;
+#endif
+#if TENSORFLOW_USE_ROCM
+using GPUResult = hipError_t;
+#endif
+
+void GPUReportIfError(GPUResult result, tensorflow::OpKernelContext *ctx,
+                      const char *expr_str) {
+  if (!result) return;
+  const char *name = nullptr;
+
+#if GOOGLE_CUDA
+  cuGetErrorName(result, &name);
+#endif
+#if TENSORFLOW_USE_ROCM
+  name = hipGetErrorName(result);
+#endif
+
+  if (!name) name = "<unknown>";
+  std::string msg = absl::StrCat("'", expr_str, "' failed with '", name, "'");
+  ReportInternalError(ctx, msg);
+}
+
+#define GPU_REPORT_IF_ERROR_WITH_CTX(expr, ctx) \
+  GPUReportIfError(expr, ctx, #expr)
+#define GPU_REPORT_IF_ERROR(expr) GPU_REPORT_IF_ERROR_WITH_CTX(expr, nullptr)
+
+// Implement the GPU module cache and share what can be shared.
+
+namespace mlir {
+namespace kernel_gen {
+namespace tf_framework {
+
+GPURuntimeCache::~GPURuntimeCache() {
+  tensorflow::mutex_lock lock(mu_);
+  for (auto it : gpu_module_by_data_ptr_) {
+#if GOOGLE_CUDA
+    GPU_REPORT_IF_ERROR(cuModuleUnload(it.second));
+#endif
+#if TENSORFLOW_USE_ROCM
+    GPU_REPORT_IF_ERROR(hipModuleUnload(it.second));
+#endif
+  }
+}
+
+absl::Status GPURuntimeCache::Create(GPURuntimeCache **dst) {
+  *dst = new GPURuntimeCache;
+  return absl::OkStatus();
+}
+
+std::string GPURuntimeCache::DebugString() const { return "GPU runtime cache"; }
+
+GPURuntimeCache::GPUModule GPURuntimeCache::LookupOrLoadModule(void *data) {
+  tensorflow::mutex_lock lock(mu_);
+  GPUModule &module = gpu_module_by_data_ptr_[data];
+
+#if GOOGLE_CUDA
+  if (!module) GPU_REPORT_IF_ERROR(cuModuleLoadData(&module, data));
+#endif
+#if TENSORFLOW_USE_ROCM
+  if (!module) GPU_REPORT_IF_ERROR(hipModuleLoadData(&module, data));
+#endif
+
+  return module;
+}
+
+GPURuntimeCache::GPUFunction GPURuntimeCache::LookupOrGetFunction(
+    GPUModule module, const char *kernel_name) {
+  tensorflow::mutex_lock lock(mu_);
+  GPUFunction &function =
+      gpu_function_by_module_and_name_[{module, kernel_name}];
+
+  if (!function) {
+#if GOOGLE_CUDA
+    GPU_REPORT_IF_ERROR(cuModuleGetFunction(&function, module, kernel_name));
+#endif
+#if TENSORFLOW_USE_ROCM
+    GPU_REPORT_IF_ERROR(hipModuleGetFunction(&function, module, kernel_name));
+#endif
+  }
+
+  return function;
+}
+
+// Implements a C wrapper around the TensorFlow runtime and CUDA (or ROCm)
+// library that allows launching a kernel on the current device and stream from
+// a binary blob for the module and function name.
+// The wrapper uses intptr_t instead of CUDA's unsigned int (or ROCm's unsigned
+// int) to match the type of MLIR's index type. This avoids the need for casts
+// in the generated MLIR code.
+extern "C" void _mlir_ciface_tf_launch_kernel(void *ctx, void *module_blob,
+                                              char *kernel_name, intptr_t gridX,
+                                              intptr_t gridY, intptr_t gridZ,
+                                              intptr_t blockX, intptr_t blockY,
+                                              intptr_t blockZ, void **params) {
+  // For empty grids, we don't need to do anything.
+  if (!gridX || !gridY || !gridZ) return;
+
+  // Get the GPU module cache.
+  auto *op_kernel_ctx = static_cast<tensorflow::OpKernelContext *>(ctx);
+  auto *rm = op_kernel_ctx->resource_manager();
+  if (rm == nullptr) {
+    ReportInternalError(op_kernel_ctx, "expected resource_manager");
+    return;
+  }
+  GPURuntimeCache *cache = nullptr;
+  OP_REQUIRES_OK(op_kernel_ctx, rm->LookupOrCreate<GPURuntimeCache>(
+                                    rm->default_container(),
+                                    GPURuntimeCache::kDefaultResourceName,
+                                    &cache, GPURuntimeCache::Create));
+  assert(cache != nullptr && "cache creation must not fail");
+  tensorflow::core::ScopedUnref ref(cache);
+
+  // Get the GPU module.
+  stream_executor::Stream *se_stream =
+      op_kernel_ctx->op_device_context()->stream();
+  void *stream = se_stream->platform_specific_handle().stream;
+  GPURuntimeCache::GPUModule module = cache->LookupOrLoadModule(module_blob);
+  GPURuntimeCache::GPUFunction function =
+      cache->LookupOrGetFunction(module, kernel_name);
+
+#if GOOGLE_CUDA
+  GPU_REPORT_IF_ERROR_WITH_CTX(
+      cuLaunchKernel(function, gridX, gridY, gridZ, blockX, blockY, blockZ,
+                     /*sharedMemBytes=*/0, reinterpret_cast<CUstream>(stream),
+                     params, nullptr),
+      op_kernel_ctx);
+#endif
+#if TENSORFLOW_USE_ROCM
+  GPU_REPORT_IF_ERROR_WITH_CTX(
+      hipModuleLaunchKernel(
+          function, gridX, gridY, gridZ, blockX, blockY, blockZ,
+          /*sharedMemBytes=*/0, reinterpret_cast<hipStream_t>(stream), params,
+          nullptr),
+      op_kernel_ctx);
+#endif
+}
+
+}  // namespace tf_framework
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/tf_gpu_runtime_wrappers.h

@@ -0,0 +1,98 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TF_GPU_RUNTIME_WRAPPERS_H_
+#define TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TF_GPU_RUNTIME_WRAPPERS_H_
+
+#include "absl/container/flat_hash_map.h"
+#include "mlir/ExecutionEngine/RunnerUtils.h"  // from @llvm-project
+#include "tensorflow/core/framework/resource_base.h"
+#include "tensorflow/core/framework/resource_op_kernel.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/status.h"
+#include "tsl/platform/hash.h"
+#include "tsl/platform/thread_annotations.h"
+
+#if GOOGLE_CUDA
+#include "third_party/gpus/cuda/include/cuda.h"
+#endif
+#if TENSORFLOW_USE_ROCM
+#include "rocm/include/hip/hip_runtime.h"
+#endif
+
+namespace mlir {
+namespace kernel_gen {
+namespace tf_framework {
+
+class GPURuntimeCache : public tensorflow::ResourceBase {
+ public:
+#if GOOGLE_CUDA
+  using GPUModule = CUmodule;
+  using GPUFunction = CUfunction;
+#endif
+#if TENSORFLOW_USE_ROCM
+  using GPUModule = hipModule_t;
+  using GPUFunction = hipFunction_t;
+#endif
+
+  ~GPURuntimeCache() override;
+  static constexpr const char* kDefaultResourceName = "mlir-gpu-runtime-cache";
+  static absl::Status Create(GPURuntimeCache** dst);
+  std::string DebugString() const override;
+
+  // Assumes that no two modules are loaded from the same memory location over
+  // the lifetime of this cache. This allows to use the pointer as a key. All
+  // modules are unloaded on destruction of this cache.
+  GPUModule LookupOrLoadModule(void* data);
+
+  GPUFunction LookupOrGetFunction(GPUModule module, const char* kernel_name);
+
+ private:
+  struct FunctionKey {
+    GPUModule module;
+    const char* kernel_name;
+
+    friend bool operator==(const FunctionKey& lhs, const FunctionKey& rhs) {
+      return lhs.module == rhs.module && lhs.kernel_name == rhs.kernel_name;
+    }
+
+    struct Hash {
+      size_t operator()(const FunctionKey& key) const {
+        return tsl::Hash64Combine(tsl::hash<GPUModule>()(key.module),
+                                  tsl::Hash64(key.kernel_name));
+      }
+    };
+  };
+
+  tensorflow::mutex mu_;
+  absl::flat_hash_map<void*, GPUModule> gpu_module_by_data_ptr_
+      TF_GUARDED_BY(mu_);
+  absl::flat_hash_map<FunctionKey, GPUFunction, FunctionKey::Hash>
+      gpu_function_by_module_and_name_ TF_GUARDED_BY(mu_);
+};
+
+// Implements a C wrapper around the TensorFlow runtime and CUDA (or ROCm)
+// library that allows launching a kernel on the current device and stream from
+// a binary blob for the module and function name.
+extern "C" MLIR_RUNNERUTILS_EXPORT void _mlir_ciface_tf_launch_kernel(
+    void* ctx, void* module_blob, char* kernel_name, intptr_t gridX,
+    intptr_t gridY, intptr_t gridZ, intptr_t blockX, intptr_t blockY,
+    intptr_t blockZ, void** params);
+
+}  // namespace tf_framework
+}  // namespace kernel_gen
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TF_GPU_RUNTIME_WRAPPERS_H_

tensorflow/compiler/mlir/tools/kernel_gen/tf_jit_cache.cc

@@ -0,0 +1,73 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tools/kernel_gen/tf_jit_cache.h"
+
+#include <cstddef>
+#include <functional>
+#include <memory>
+#include <string>
+#include <utility>
+
+#include "absl/status/status.h"
+#include "llvm/Support/Error.h"
+#include "mlir/ExecutionEngine/ExecutionEngine.h"  // from @llvm-project
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace tf_framework {
+
+absl::Status JITCache::Create(JITCache** dst) {
+  *dst = new JITCache;
+  return absl::OkStatus();
+}
+
+std::string JITCache::DebugString() const { return "JIT cache"; }
+
+ExecutionEngine* JITCache::LookupOrCompile(
+    const std::string code,
+    std::function<llvm::Expected<std::unique_ptr<ExecutionEngine>>()>
+        compile_callback) {
+  // Check if we already have a compiled module in the cache.
+  {
+    tensorflow::mutex_lock lock(mu_);
+    if (execution_engine_by_key_.contains(code))
+      return execution_engine_by_key_[code].get();
+  }
+
+  // Otherwise, compile the module now.
+  llvm::Expected<std::unique_ptr<ExecutionEngine>> engine = compile_callback();
+  if (!engine) return nullptr;
+
+  // Insert the compiled module into our cache and return a raw pointer.
+  {
+    tensorflow::mutex_lock lock(mu_);
+    // Check again whether we already have a compiled module in the cache. It
+    // may have been added during the time we ran compile_callback().
+    return execution_engine_by_key_.try_emplace(code, std::move(engine.get()))
+        .first->second.get();
+  }
+}
+
+size_t JITCache::Size() {
+  tensorflow::mutex_lock lock(mu_);
+  return execution_engine_by_key_.size();
+}
+
+}  // namespace tf_framework
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/tf_jit_cache.h

@@ -0,0 +1,59 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TF_JIT_CACHE_H_
+#define TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TF_JIT_CACHE_H_
+
+#include <cstddef>
+#include <functional>
+#include <memory>
+#include <string>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/status/status.h"
+#include "mlir/ExecutionEngine/ExecutionEngine.h"  // from @llvm-project
+#include "tensorflow/core/framework/resource_base.h"
+#include "tensorflow/core/framework/resource_op_kernel.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/core/platform/status.h"
+#include "tsl/platform/thread_annotations.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace tf_framework {
+
+class JITCache : public tensorflow::ResourceBase {
+ public:
+  static constexpr const char* kDefaultResourceName = "mlir-jit-cache";
+  static absl::Status Create(JITCache** dst);
+
+  std::string DebugString() const override;
+  ExecutionEngine* LookupOrCompile(
+      std::string code,
+      std::function<llvm::Expected<std::unique_ptr<ExecutionEngine>>()>
+          compile_callback);
+  size_t Size();
+
+ private:
+  tensorflow::mutex mu_;
+  absl::flat_hash_map<std::string, std::unique_ptr<ExecutionEngine>>
+      execution_engine_by_key_ TF_GUARDED_BY(mu_);
+};
+
+}  // namespace tf_framework
+}  // namespace kernel_gen
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TF_JIT_CACHE_H_

tensorflow/compiler/mlir/tools/kernel_gen/tools/kernel-gen-opt/kernel-gen-opt.cc

@@ -0,0 +1,44 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "llvm/Support/LogicalResult.h"
+#include "mlir/InitAllDialects.h"  // from @llvm-project
+#include "mlir/InitAllPasses.h"  // from @llvm-project
+#include "mlir/Tools/mlir-opt/MlirOptMain.h"  // from @llvm-project
+#include "stablehlo/dialect/Register.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/tensorflow/dialect_registration.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+#include "xla/mlir_hlo/deallocation/transforms/passes.h"
+#include "xla/mlir_hlo/mhlo/IR/register.h"
+#include "xla/mlir_hlo/mhlo/transforms/passes.h"
+
+int main(int argc, char **argv) {
+  mlir::registerAllPasses();
+  mlir::mhlo::registerAllMhloPasses();
+  mlir::kernel_gen::registerKernelGenPasses();
+  mlir::deallocation::registerDeallocationPasses();
+
+  mlir::DialectRegistry registry;
+  mlir::registerAllDialects(registry);
+  mlir::mhlo::registerAllMhloDialects(registry);
+  mlir::stablehlo::registerAllDialects(registry);
+  mlir::RegisterAllTensorFlowDialects(registry);
+
+  registry.insert<mlir::kernel_gen::tf_framework::TFFrameworkDialect>();
+
+  return failed(
+      mlir::MlirOptMain(argc, argv, "MLIR HLO pass driver\n", registry));
+}

tensorflow/compiler/mlir/tools/kernel_gen/transforms/BUILD

@@ -0,0 +1,235 @@
+load("@llvm-project//mlir:tblgen.bzl", "gentbl_cc_library")
+load(
+    "@local_config_rocm//rocm:build_defs.bzl",
+    "if_rocm_is_configured",
+)
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load(
+    "@xla//xla/tsl/platform/default:cuda_build_defs.bzl",
+    "if_cuda_is_configured",
+)
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow/compiler/mlir/tools/kernel_gen:friends"],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "utils",
+    srcs = ["utils.cc"],
+    hdrs = ["utils.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LLVMDialect",
+    ],
+)
+
+cc_library(
+    name = "tf_framework_legalize_to_llvm",
+    srcs = ["tf_framework_legalize_to_llvm.cc"],
+    hdrs = ["rewriters.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":utils",
+        "//tensorflow/compiler/mlir/tools/kernel_gen/ir:tf_framework_ops",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LLVMCommonConversion",
+        "@llvm-project//mlir:LLVMDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
+
+cc_library(
+    name = "bufferize",
+    srcs = ["bufferize.cc"],
+    hdrs = ["rewriters.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/tools/kernel_gen/ir:tf_framework_ops",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:BufferizationTransforms",
+        "@llvm-project//mlir:ComplexDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@stablehlo//:chlo_ops",
+    ],
+)
+
+cc_library(
+    name = "embed_tf_framework",
+    srcs = ["embed_tf_framework.cc"],
+    hdrs = ["rewriters.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/tools/kernel_gen/ir:tf_framework_ops",
+        "@llvm-project//mlir:ControlFlowDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+    ],
+)
+
+gentbl_cc_library(
+    name = "kernel_gen_passes_inc_gen",
+    compatible_with = get_compatible_with_portable(),
+    tbl_outs = {"kernel_gen_passes.h.inc": [
+        "-gen-pass-decls",
+        "-name=KernelGen",
+    ]},
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "passes.td",
+    deps = ["@llvm-project//mlir:PassBaseTdFiles"],
+)
+
+cc_library(
+    name = "gpu_passes",
+    srcs = [
+        "gpu_kernel_to_blob_pass.cc",
+        "tf_kernel_to_llvm_pass.cc",
+    ],
+    hdrs = ["passes.h"],
+    copts = if_cuda_is_configured(["-DGOOGLE_CUDA=1"]) + if_rocm_is_configured(["-DTENSORFLOW_USE_ROCM=1"]),
+    deps = [
+        ":embed_tf_framework",
+        ":kernel_gen_passes_inc_gen",
+        "//tensorflow/compiler/mlir/tensorflow",
+        "//tensorflow/compiler/mlir/tools/kernel_gen/ir:tf_framework_ops",
+        "//tensorflow/core:lib",
+        "//tensorflow/core/platform:errors",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//llvm:TargetParser",
+        "@llvm-project//llvm:TransformUtils",
+        "@llvm-project//mlir:AffineDialect",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:ArithToLLVM",
+        "@llvm-project//mlir:ArithTransforms",
+        "@llvm-project//mlir:ComplexDialect",
+        "@llvm-project//mlir:ComplexToLLVM",
+        "@llvm-project//mlir:ControlFlowDialect",
+        "@llvm-project//mlir:ControlFlowToLLVM",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:FuncToLLVM",
+        "@llvm-project//mlir:FuncTransforms",
+        "@llvm-project//mlir:GPUDialect",
+        "@llvm-project//mlir:GPUToGPURuntimeTransforms",
+        "@llvm-project//mlir:GPUTransforms",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:LLVMCommonConversion",
+        "@llvm-project//mlir:LLVMDialect",
+        "@llvm-project//mlir:LinalgDialect",
+        "@llvm-project//mlir:LinalgTransforms",
+        "@llvm-project//mlir:MathDialect",
+        "@llvm-project//mlir:MathToLLVM",
+        "@llvm-project//mlir:MathToLibm",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:MemRefToLLVM",
+        "@llvm-project//mlir:MemRefTransforms",
+        "@llvm-project//mlir:NVVMToLLVMIRTranslation",  # buildcleaner: keep
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:ROCDLToLLVMIRTranslation",  # buildcleaner: keep
+        "@llvm-project//mlir:ReconcileUnrealizedCasts",
+        "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:SCFToControlFlow",
+        "@llvm-project//mlir:SCFTransforms",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:ShapeToStandard",
+        "@llvm-project//mlir:ShapeTransforms",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:TensorTransforms",
+        "@llvm-project//mlir:ToLLVMIRTranslation",
+        "@llvm-project//mlir:TransformUtils",
+        "@llvm-project//mlir:Transforms",
+        "@llvm-project//mlir:VectorDialect",
+        "@llvm-project//mlir:VectorToLLVM",
+        "@llvm-project//mlir:VectorTransforms",
+        "@xla//xla:debug_options_flags",
+        "@xla//xla:xla_proto_cc",
+        "@xla//xla/mlir_hlo",
+        "@xla//xla/mlir_hlo:mhlo_passes",
+        "@xla//xla/mlir_hlo:type_conversion",
+        "@xla//xla/service/gpu:gpu_asm_opts_util",
+        "@xla//xla/service/gpu:target_constants",
+        "@xla//xla/stream_executor:device_description",
+    ] + if_cuda_is_configured([
+        "@tsl//tsl/platform:cuda_root_path",
+        "@xla//xla/stream_executor/cuda:cuda_asm_compiler",
+        "@xla//xla/service/gpu/llvm_gpu_backend:nvptx_backend",
+    ]) + if_rocm_is_configured([
+        "@xla//xla/stream_executor/gpu:asm_compiler",
+        "@xla//xla/service/gpu/llvm_gpu_backend:amdgpu_backend",
+        "//tensorflow/core/platform:rocm_rocdl_path",
+    ]),
+)
+
+cc_library(
+    name = "passes",
+    srcs = [
+        "broadcast_propagation_pass.cc",
+        "buffer_reuse_pass.cc",
+        "bufferize_pass.cc",
+        "copy_cleanup_pass.cc",
+        "embed_tf_framework_pass.cc",
+        "func_to_jit_invocations.cc",
+        "fuse_inner_parallel_loops_pass.cc",
+        "legalize_tensor_reshape_pass.cc",
+        "merge_assuming_ops_pass.cc",
+        "parallel_loops_to_sequential.cc",
+        "rewrite_tf_framework_assert.cc",
+        "same_shape_propagation.cc",
+        "shape_simplification_pass.cc",
+        "shape_to_descriptors_pass.cc",
+        "tensorflow_abi_knowledge_propagation.cc",
+    ],
+    hdrs = ["passes.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        ":bufferize",  # buildcleaner: keep
+        ":embed_tf_framework",  # buildcleaner: keep
+        ":kernel_gen_passes_inc_gen",
+        ":tf_framework_legalize_to_llvm",  # buildcleaner: keep
+        "//tensorflow/compiler/mlir/tools/kernel_gen/ir:tf_framework_ops",
+        "@llvm-project//llvm:Support",
+        "@llvm-project//mlir:Analysis",
+        "@llvm-project//mlir:ArithDialect",
+        "@llvm-project//mlir:BufferizationTransforms",
+        "@llvm-project//mlir:ControlFlowDialect",
+        "@llvm-project//mlir:FuncDialect",
+        "@llvm-project//mlir:GPUDialect",
+        "@llvm-project//mlir:IR",
+        "@llvm-project//mlir:InferTypeOpInterface",
+        "@llvm-project//mlir:LLVMDialect",
+        "@llvm-project//mlir:LinalgDialect",
+        "@llvm-project//mlir:MathDialect",
+        "@llvm-project//mlir:MemRefDialect",
+        "@llvm-project//mlir:Pass",
+        "@llvm-project//mlir:SCFDialect",
+        "@llvm-project//mlir:SCFToControlFlow",
+        "@llvm-project//mlir:SCFTransforms",
+        "@llvm-project//mlir:ShapeDialect",
+        "@llvm-project//mlir:ShapeToStandard",
+        "@llvm-project//mlir:ShapeTransforms",
+        "@llvm-project//mlir:SideEffectInterfaces",
+        "@llvm-project//mlir:Support",
+        "@llvm-project//mlir:TensorDialect",
+        "@llvm-project//mlir:TransformUtils",
+        "@stablehlo//:base",
+        "@xla//xla/mlir_hlo",
+        "@xla//xla/mlir_hlo:mhlo_passes",
+        "@xla//xla/mlir_hlo:transforms_passes",
+    ],
+)

tensorflow/compiler/mlir/tools/kernel_gen/transforms/broadcast_propagation_pass.cc

@@ -0,0 +1,449 @@
+/* Copyright 2021 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+==============================================================================*/
+
+#include <algorithm>
+#include <cassert>
+#include <cstdint>
+#include <utility>
+
+#include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/Base.h"  // from @stablehlo
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+
+namespace mlir {
+namespace kernel_gen {
+
+using mhlo::DynamicBroadcastInDimOp;
+
+#define GEN_PASS_DEF_BROADCASTPROPAGATIONPASS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+namespace {
+
+// To avoid duplicate broadcasts, we collect all the intended broadcasts ahead
+// of realizing any broadcasts in the IR. These are broadcasted versions of
+// values that we are interested in, and they are uniquely characterized by a
+// `BroadcastIntent` value.
+struct BroadcastIntent {
+  RankedTensorType resultType;
+  Value targetValue;
+  Value outputDimensions;
+  Attribute broadcastDimensions;
+  bool operator==(BroadcastIntent rhs) const {
+    return resultType == rhs.resultType && targetValue == rhs.targetValue &&
+           outputDimensions == rhs.outputDimensions &&
+           broadcastDimensions == rhs.broadcastDimensions;
+  }
+  bool operator!=(BroadcastIntent rhs) const { return !(*this == rhs); }
+};
+
+}  // namespace
+}  // namespace kernel_gen
+}  // namespace mlir
+
+namespace llvm {
+
+using mlir::kernel_gen::BroadcastIntent;
+
+template <>
+struct DenseMapInfo<BroadcastIntent> {
+  static unsigned getHashValue(const BroadcastIntent &intent) {
+    return hash_combine(
+        DenseMapInfo<mlir::RankedTensorType>::getHashValue(intent.resultType),
+        DenseMapInfo<mlir::Value>::getHashValue(intent.targetValue),
+        DenseMapInfo<mlir::Value>::getHashValue(intent.outputDimensions),
+        DenseMapInfo<mlir::Attribute>::getHashValue(
+            intent.broadcastDimensions));
+  }
+  static bool isEqual(const BroadcastIntent &lhs, const BroadcastIntent &rhs) {
+    return lhs == rhs;
+  }
+};
+
+}  // namespace llvm
+
+namespace mlir {
+namespace kernel_gen {
+namespace {
+
+bool allowsForElementwiseBroadcastPropagation(Operation *op) {
+  if (op && op->hasTrait<mlir::OpTrait::SameOperandsAndResultShape>() &&
+      op->hasTrait<mlir::OpTrait::Elementwise>() && op->getNumResults() == 1) {
+    return true;
+  }
+  if (op && op->hasTrait<hlo::OpTrait::BroadcastingElementwise>() &&
+      op->getNumResults() == 1) {
+    return true;
+  }
+  return false;
+}
+
+bool allowsForBroadcastPropagation(Operation *op) {
+  return llvm::isa_and_nonnull<DynamicBroadcastInDimOp>(op) ||
+         allowsForElementwiseBroadcastPropagation(op);
+}
+
+DenseIntElementsAttr composeBroadcastDimensionsAttr(OpBuilder &builder,
+                                                    DenseIntElementsAttr a,
+                                                    DenseIntElementsAttr b) {
+  SmallVector<int64_t> bVec =
+      llvm::to_vector(llvm::map_range(b, [](const APInt &it) {
+        return static_cast<int64_t>(it.getLimitedValue());
+      }));
+  SmallVector<int64_t> composedVec = llvm::to_vector(llvm::map_range(
+      a, [bVec](const APInt &it) { return bVec[it.getLimitedValue()]; }));
+  return builder.getI64TensorAttr(composedVec);
+}
+
+// Find all the broadcast intents and their dependencies. Start analyzing from
+// the root an collect all broadcast intents that can help broadcast propagation
+// from there.
+void findBroadcastIntents(
+    DynamicBroadcastInDimOp root, Block *parentBlock,
+    BroadcastIntent &rootBcastIntent,
+    SmallVector<BroadcastIntent> &bcastIntents,
+    DenseMap<BroadcastIntent, SmallVector<BroadcastIntent>>
+        &bcastIntentDependencies) {
+  OpBuilder builder(root.getContext());
+
+  // Use the result vector of broadcast intents as a worklist. The set of
+  // broadcast intents helps to ensure their uniqueness.
+  DenseSet<BroadcastIntent> bcastIntentsSet;
+  auto addToWorklistIfNew = [&](BroadcastIntent bcastIntent) {
+    if (!bcastIntentsSet.count(bcastIntent)) {
+      bcastIntentsSet.insert(bcastIntent);
+      bcastIntents.push_back(bcastIntent);
+    }
+  };
+
+  // Derive the broadcast intent associated with the root broadcast operation.
+  // Add it to the worklist to seed the analysis.
+  rootBcastIntent = {mlir::cast<RankedTensorType>(root.getResult().getType()),
+                     root.getOperand(), root.getOutputDimensions(),
+                     root.getBroadcastDimensions()};
+  addToWorklistIfNew(rootBcastIntent);
+
+  // We use result vector of broadcast intents as a worklist, the first `i`
+  // intents of which have been processed.
+  for (int64_t i = 0; i < static_cast<int64_t>(bcastIntents.size()); ++i) {
+    BroadcastIntent it = bcastIntents[i];
+    Operation *producerOp = it.targetValue.getDefiningOp();
+
+    // We can propagate broadcasts over (broadcasting) element-wise operations
+    // and dynamic_broadcast_in_dim ops with the restriction that they must be
+    // in the same block as they may depend on assuming regions.
+    if (!producerOp || producerOp->getBlock() != parentBlock ||
+        !allowsForBroadcastPropagation(producerOp)) {
+      continue;
+    }
+
+    // We can skip broadcasting producers (dynamic_broadcast_in_dim ops) if we
+    // compose their broadcasting dimensions.
+    if (auto producerBcastOp =
+            llvm::dyn_cast<DynamicBroadcastInDimOp>(producerOp)) {
+      DenseIntElementsAttr composedBcastDims = composeBroadcastDimensionsAttr(
+          builder, producerBcastOp.getBroadcastDimensions(),
+          mlir::cast<DenseIntElementsAttr>(it.broadcastDimensions));
+      BroadcastIntent bcastedOperandIntent = {
+          it.resultType, producerBcastOp.getOperand(), it.outputDimensions,
+          composedBcastDims};
+
+      // Record dependency and "recur".
+      bcastIntentDependencies[it] = {bcastedOperandIntent};
+      addToWorklistIfNew(bcastedOperandIntent);
+      continue;
+    }
+
+    // We can propagate broadcasts over (broadcasting) element-wise operations.
+    // Instead of broadcasting the result of such an op, we can broadcast the
+    // operands and apply the element-wise operation to them.
+    assert(allowsForElementwiseBroadcastPropagation(producerOp));
+    bcastIntentDependencies[it] = {};
+    for (auto operand : producerOp->getOperands()) {
+      auto operandTy = mlir::cast<RankedTensorType>(operand.getType());
+      auto operandBcastDims = operandTy.getRank() == 0
+                                  ? builder.getI64TensorAttr({})
+                                  : it.broadcastDimensions;
+      auto bcastedOperandTy = RankedTensorType::get(it.resultType.getShape(),
+                                                    operandTy.getElementType());
+      BroadcastIntent bcastedOperandIntent = {
+          bcastedOperandTy, operand, it.outputDimensions, operandBcastDims};
+
+      // Record dependency and "recur".
+      bcastIntentDependencies[it].push_back(bcastedOperandIntent);
+      addToWorklistIfNew(bcastedOperandIntent);
+    }
+  }
+}
+
+void sortBroadcastIntentsInReverseTopologicalOrder(
+    SmallVector<BroadcastIntent> &bcastIntentsVec, Block *parentBlock) {
+  // Sort broadcast intents in reverse topological order of the producer ops. We
+  // can use the positions in the block for this. All broadcast intents outside
+  // the block (e.g. arguments) will be sorted towards the front.
+  // This ordering is independent of the output dimensions as dependencies can
+  // only occur between broadcast intents of the same output dimension.
+  std::sort(bcastIntentsVec.begin(), bcastIntentsVec.end(),
+            [parentBlock](const BroadcastIntent &a, const BroadcastIntent &b) {
+              Operation *producerOpA = a.targetValue.getDefiningOp();
+              Operation *producerOpB = b.targetValue.getDefiningOp();
+              bool aInBlock = producerOpA != nullptr &&
+                              producerOpA->getBlock() == parentBlock;
+              bool bInBlock = producerOpB != nullptr &&
+                              producerOpB->getBlock() == parentBlock;
+              if (aInBlock && bInBlock) {
+                return producerOpA->isBeforeInBlock(producerOpB);
+              }
+              return !aInBlock && bInBlock;
+            });
+}
+
+void setInsertionPointToEarliestPointWithAllValuesAvailable(
+    PatternRewriter &rewriter, Block *block, ValueRange values) {
+  Operation *lastDef = nullptr;
+  for (Value v : values) {
+    Operation *def = v.getDefiningOp();
+    if (def && def->getBlock() == block) {
+      if (!lastDef || lastDef->isBeforeInBlock(def)) lastDef = def;
+    }
+  }
+  if (lastDef) {
+    rewriter.setInsertionPointAfter(lastDef);
+  } else {
+    rewriter.setInsertionPointToStart(block);
+  }
+}
+
+DenseMap<BroadcastIntent, Value> realizeBroadcastIntents(
+    SmallVector<BroadcastIntent> &sortedBcastIntents,
+    DenseMap<BroadcastIntent, SmallVector<BroadcastIntent>>
+        &bcastIntentDependencies,
+    Block *parentBlock, PatternRewriter &rewriter) {
+  // Realize broadcast intents in order. They must be sorted so that their
+  // dependencies are realized before them.
+  DenseMap<BroadcastIntent, Value> realizations;
+  for (auto it : sortedBcastIntents) {
+    Operation *producerOp = it.targetValue.getDefiningOp();
+    assert(!realizations.count(it) && "expect unrealized broadcast intent");
+    auto deps = bcastIntentDependencies.find(it);
+
+    // If we cannot propagate broadcasts further, materialize them as a
+    // dynamic_broadcast_in_dim op.
+    if (!producerOp || producerOp->getBlock() != parentBlock ||
+        !allowsForBroadcastPropagation(producerOp)) {
+      assert(deps == bcastIntentDependencies.end() && "expect no dependencies");
+      setInsertionPointToEarliestPointWithAllValuesAvailable(
+          rewriter, parentBlock,
+          ValueRange{it.targetValue, it.outputDimensions});
+      realizations[it] = DynamicBroadcastInDimOp::create(
+          rewriter, it.targetValue.getLoc(), it.resultType, it.targetValue,
+          it.outputDimensions,
+          mlir::cast<DenseIntElementsAttr>(it.broadcastDimensions));
+      continue;
+    }
+
+    // For broadcast propagation across dynamic_broadcast_in_dim ops, the
+    // broadcasted value is already materialized. Forward it.
+    if (auto producerBcastOp =
+            llvm::dyn_cast_or_null<DynamicBroadcastInDimOp>(producerOp)) {
+      assert(deps != bcastIntentDependencies.end() &&
+             deps->second.size() == 1 && "expect one dependency");
+      auto bcastedOperand = realizations.find(deps->second.front());
+      assert(bcastedOperand != realizations.end());
+      realizations[it] = Value(bcastedOperand->second);
+      continue;
+    }
+
+    // Othwerwise, realize broadcast intent for a (broadcasting) element-wise
+    // operation based on the broadcasted operands.
+    assert(allowsForElementwiseBroadcastPropagation(producerOp) &&
+           "expect broadcast propagation over an (broadcasting) element-wise "
+           "operation");
+    assert(deps != bcastIntentDependencies.end() &&
+           deps->second.size() == producerOp->getNumOperands() &&
+           "expect one dependency per operand");
+    auto bcastedOperands = llvm::to_vector(
+        llvm::map_range(deps->second, [&](BroadcastIntent operandIntent) {
+          auto bcastedOperand = realizations.find(operandIntent);
+          assert(bcastedOperand != realizations.end() &&
+                 "expect dependencies to be realized earlier");
+          return bcastedOperand->second;
+        }));
+    setInsertionPointToEarliestPointWithAllValuesAvailable(
+        rewriter, parentBlock, bcastedOperands);
+    OperationState newProducerOpState(
+        producerOp->getLoc(), producerOp->getName().getStringRef(),
+        bcastedOperands, it.resultType, producerOp->getAttrs());
+    Operation *newProducerOp = rewriter.create(newProducerOpState);
+    assert(newProducerOp->getNumResults() == 1 && "expect exactly one result");
+    realizations[it] = newProducerOp->getResults().front();
+  }
+
+  return realizations;
+}
+
+void transitivelyEraseUnusedSideEffectFreeOps(Operation *root,
+                                              PatternRewriter &rewriter) {
+  // Find ops to erase.
+  SmallPtrSet<Operation *, 16> opsToEraseSet;
+  SmallVector<Operation *, 16> opsToErase;
+  SmallVector<Operation *, 16> worklist = {root};
+  while (!worklist.empty()) {
+    Operation *op = worklist.pop_back_val();
+
+    // Erase ops only once.
+    if (opsToEraseSet.count(op)) continue;
+
+    // Erase only operations that are unused and free of side effects.
+    if (!isMemoryEffectFree(op) ||
+        !llvm::all_of(op->getUsers(), [opsToEraseSet](Operation *user) {
+          return opsToEraseSet.count(user);
+        })) {
+      continue;
+    }
+
+    // Erase and "recur".
+    opsToEraseSet.insert(op);
+    opsToErase.push_back(op);
+    for (Value operand : op->getOperands()) {
+      if (Operation *def = operand.getDefiningOp()) worklist.push_back(def);
+    }
+  }
+
+  // Finally, erase the ops in the order of their uses.
+  for (Operation *op : opsToErase) rewriter.eraseOp(op);
+}
+
+LogicalResult propagateBroadcast(DynamicBroadcastInDimOp root,
+                                 Block *parentBlock,
+                                 PatternRewriter &rewriter) {
+  // We can move broadcasts up over (i) (broadcasting) element-wise operations
+  // and (i) dynamic_broadcast_in_dim ops. This way, we propagate them through
+  // the IR to perform them early. Instead of broadcasting the result of such an
+  // op, we can broadcast the operands and apply the element-wise operation to
+  // them.
+  //
+  // To avoid exponential growth of the IR, we will do this in two phases:
+  //   1) First, we collect all the unique broadcast intents. These are
+  //      broadcasted versions of values that we are interested in. They may
+  //      later be materialized as an explicit broadcast or they can be the
+  //      direct result of an operation over which a broadcast was propagated.
+  //   2) Then, we fulfill every broadcast intent in reverse topological order
+  //      to ensure that their dependencies (the broadcasted operands) are
+  //      available.
+
+  // Find the unique broadcast intents.
+  BroadcastIntent rootBcastIntent;
+  SmallVector<BroadcastIntent> bcastIntents;
+  DenseMap<BroadcastIntent, SmallVector<BroadcastIntent>>
+      bcastIntentDependencies;
+  findBroadcastIntents(root, parentBlock, rootBcastIntent, bcastIntents,
+                       bcastIntentDependencies);
+
+  // Fail if there is nothing but the root intent, i.e. if there is nothing to
+  // rewrite here.
+  if (bcastIntents.size() <= 1) {
+    assert(bcastIntents.front() == rootBcastIntent && "expect root intent");
+    return failure();
+  }
+
+  // Sort the broadcast intents in reverse topological order so that they can be
+  // materialized and every depency is available when needed.
+  sortBroadcastIntentsInReverseTopologicalOrder(bcastIntents, parentBlock);
+
+  // Realize broadcast intents.
+  DenseMap<BroadcastIntent, Value> realizations = realizeBroadcastIntents(
+      bcastIntents, bcastIntentDependencies, parentBlock, rewriter);
+
+  // Find the operations that may become redundant after replacing the root
+  // operation. This allows us to transitively erase unused side effect-free
+  // operations that result from this rewrite (after the root operation is no
+  // longer accessible).
+  SmallVector<Operation *> possiblyUnused;
+  for (auto operand : root->getOperands()) {
+    if (Operation *def = operand.getDefiningOp()) possiblyUnused.push_back(def);
+  }
+
+  // Replace the root operation with its broadcast intent's realization.
+  rewriter.replaceOp(root, realizations[rootBcastIntent]);
+
+  // Erase all the operations that have become redundant as a result of this
+  // rewrite.
+  for (Operation *op : possiblyUnused) {
+    transitivelyEraseUnusedSideEffectFreeOps(op, rewriter);
+  }
+
+  return success();
+}
+
+struct BroadcastPropagationPattern
+    : public OpRewritePattern<DynamicBroadcastInDimOp> {
+  using OpRewritePattern<DynamicBroadcastInDimOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(DynamicBroadcastInDimOp op,
+                                PatternRewriter &rewriter) const override {
+    return propagateBroadcast(op, op->getBlock(), rewriter);
+  }
+};
+
+struct BroadcastPropagationPass
+    : public impl::BroadcastPropagationPassBase<BroadcastPropagationPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<mhlo::MhloDialect>();
+  }
+
+  void runOnOperation() override {
+    MLIRContext *ctx = &getContext();
+
+    // Collect patterns.
+    RewritePatternSet patterns(ctx);
+    patterns.add<BroadcastPropagationPattern>(ctx);
+
+    // Apply broadcast propagation in reverse order to start propagation at
+    // the root of broadcast chains. This avoids duplicate work.
+    GreedyRewriteConfig config;
+    config.setUseTopDownTraversal(false);
+
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns),
+                                     config))) {
+      return signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/buffer_reuse_pass.cc

@@ -0,0 +1,275 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cassert>
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <vector>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Analysis/Liveness.h"  // from @llvm-project
+#include "mlir/Dialect/Bufferization/Transforms/BufferViewFlowAnalysis.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Linalg/IR/Linalg.h"  // from @llvm-project
+#include "mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
+#include "mlir/IR/AffineMap.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewriters.h"
+
+constexpr llvm::StringRef
+    mlir::kernel_gen::tf_framework::TFAllocOp::kReuseOutputAttrName;
+constexpr llvm::StringRef
+    mlir::kernel_gen::tf_framework::TFAllocOp::kReuseInputCandidatesAttrName;
+constexpr llvm::StringRef
+    mlir::kernel_gen::tf_framework::TFFrameworkDialect::kTFEntryAttrName;
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+namespace {
+
+class BufferReuseAnalysis {
+ public:
+  explicit BufferReuseAnalysis(func::FuncOp f) { build(f); }
+
+  static constexpr int32_t kIndexAmbiguous = -1;
+
+  std::optional<SmallVector<int32_t, 2>> get_reuse_candiates(
+      memref::AllocOp op) {
+    auto it = reuse_candidates_.find(op);
+    if (it == reuse_candidates_.end()) return std::nullopt;
+    return it->second;
+  }
+
+  std::optional<int32_t> get_output_index(memref::AllocOp op) {
+    auto it = output_indices_.find(op);
+    if (it == output_indices_.end()) return std::nullopt;
+    return it->second;
+  }
+
+ private:
+  void build(func::FuncOp &f) {
+    BufferViewFlowAnalysis aliases(f);
+    find_output_indices(f, aliases);
+    find_reuse_candiates(f, aliases);
+  }
+
+  void find_output_indices(func::FuncOp &f, BufferViewFlowAnalysis &aliases) {
+    f.walk([&](memref::AllocOp alloc_op) {
+      int32_t output_index = kIndexAmbiguous;
+      int count_return_uses = 0;
+      auto buffer_aliases = aliases.resolve(alloc_op.getResult());
+      for (Value alias : buffer_aliases) {
+        for (auto &use : alias.getUses()) {
+          if (isa<func::ReturnOp>(use.getOwner())) {
+            int32_t index = use.getOperandNumber();
+            if (count_return_uses++ == 0)
+              output_index = index;
+            else if (output_index != index)
+              output_index = kIndexAmbiguous;
+          }
+        }
+      }
+      output_indices_[alloc_op] = output_index;
+    });
+  }
+
+  void find_reuse_candiates(func::FuncOp &f, BufferViewFlowAnalysis &aliases) {
+    Liveness liveness(f);
+    f.walk([&](Block *block) {
+      find_reuse_candiates(block, aliases, liveness.getLiveness(block),
+                           f.getArguments());
+    });
+  }
+
+  void find_reuse_candiates(Block *block, BufferViewFlowAnalysis &aliases,
+                            const LivenessBlockInfo *liveness,
+                            ArrayRef<BlockArgument> arguments) {
+    for (Operation &op : *block) {
+      auto alloc_op = dyn_cast<memref::AllocOp>(op);
+      if (!alloc_op) continue;
+
+      // Find first use of the newly allocated buffer within this block.
+      Value new_buffer = alloc_op.getResult();
+      Operation *first_reuse = find_first_use_in_block(new_buffer, block);
+      assert((first_reuse == nullptr || first_reuse->getBlock() == block) &&
+             "Expected first use in same block if found.");
+
+      // Find reuse candidates for the regarded allocation.
+      SmallVector<int32_t, 2> local_reuse_candidates;
+      for (BlockArgument old_buffer : arguments) {
+        if (!mlir::isa<BaseMemRefType>(old_buffer.getType())) continue;
+
+        // Lifetime criterion: Only reuse buffers that are no longer used on
+        // first reuse, i.e. they are no longer alive.
+        bool lifetimes_compatible = true;
+        for (Value old_buffer_alias : aliases.resolve(old_buffer)) {
+          if (first_reuse == nullptr) {
+            // If the first use is beyond the end of this block we look at the
+            // block end. An argument buffer that is already reusable there is
+            // certainly reusable at any later actual use. Otherwise, lifetimes
+            // are incompatible.
+            if (liveness->isLiveOut(old_buffer_alias)) {
+              lifetimes_compatible = false;
+              break;
+            }
+          } else {
+            // A buffer is reusable if
+            //   i)  its last use is before the point of reuse, or
+            //   ii) its last use is also its first reuse and the operation
+            //       allows for local reuse.
+            // Otherwise, lifetimes are incompatible.
+            Operation *last_use =
+                liveness->getEndOperation(old_buffer_alias, &block->front());
+            assert(last_use != nullptr && last_use->getBlock() == block &&
+                   "Expected last use in same block.");
+            if (first_reuse->isBeforeInBlock(last_use)) {
+              lifetimes_compatible = false;
+              break;
+            }
+            if (first_reuse == last_use &&
+                !can_reuse_locally(first_reuse, old_buffer_alias, new_buffer)) {
+              lifetimes_compatible = false;
+              break;
+            }
+          }
+        }
+
+        if (lifetimes_compatible) {
+          // All criteria are fulfilled 🙂.
+          int32_t old_buffer_index = old_buffer.getArgNumber();
+          local_reuse_candidates.push_back(old_buffer_index);
+        }
+      }
+
+      reuse_candidates_[&op] = local_reuse_candidates;
+    }
+  }
+
+  Operation *find_first_use_in_block(Value value, Block *block) {
+    Operation *first_use = nullptr;
+    for (Operation *op : value.getUsers()) {
+      Operation *ancestor_op = block->findAncestorOpInBlock(*op);
+      if (ancestor_op == nullptr) continue;
+      if (first_use == nullptr || ancestor_op->isBeforeInBlock(first_use))
+        first_use = ancestor_op;
+    }
+    return first_use;
+  }
+
+  std::vector<Value> get_buffer_arguments(func::FuncOp &f) {
+    std::vector<Value> buffer_arguments;
+    for (BlockArgument arg : f.getArguments()) {
+      if (mlir::isa<BaseMemRefType>(arg.getType()))
+        buffer_arguments.push_back(arg);
+    }
+    return buffer_arguments;
+  }
+
+  bool can_reuse_locally(Operation *op, Value old_buffer, Value new_buffer) {
+    // For now, we support only memrefs with the same memory layout.
+    auto old_buffer_ty = mlir::dyn_cast<MemRefType>(old_buffer.getType());
+    auto new_buffer_ty = mlir::dyn_cast<MemRefType>(old_buffer.getType());
+    if (!old_buffer_ty || !new_buffer_ty ||
+        old_buffer_ty.getLayout() != new_buffer_ty.getLayout())
+      return false;
+
+    if (auto generic_op = dyn_cast<linalg::GenericOp>(op)) {
+      auto op_operands = op->getOpOperands();
+      auto old_it = llvm::find_if(op_operands, [&](OpOperand &op_operand) {
+        return op_operand.get() == old_buffer;
+      });
+      auto new_it = llvm::find_if(op_operands, [&](OpOperand &op_operand) {
+        return op_operand.get() == new_buffer;
+      });
+      assert(old_it != op_operands.end() && new_it != op_operands.end() &&
+             "Expect `old/new_buffer` to be operand of `op`.");
+
+      auto is_projection = [](AffineMap map) {
+        // Allow dropping dimensions but no permutations.
+        int64_t i = -1;
+        for (AffineExpr expr : map.getResults()) {
+          auto dim_expr = mlir::dyn_cast<AffineDimExpr>(expr);
+          if (!dim_expr || dim_expr.getPosition() <= i) return false;
+          i = dim_expr.getPosition();
+        }
+        return true;
+      };
+
+      // If `linalg.generic` indexing maps are the same for input and output
+      // buffer then the last use of the input buffer happens before its first
+      // reuse (per memory location). Since we know that the inputs and outputs
+      // have the same size we also know that when one side has an identity map
+      // and the other side only drops dimensions, these dimensions have to be
+      // of size 1.
+      AffineMap old_indexing_map = generic_op.getMatchingIndexingMap(old_it);
+      AffineMap new_indexing_map = generic_op.getMatchingIndexingMap(new_it);
+      return (old_indexing_map == new_indexing_map &&
+              old_indexing_map.isProjectedPermutation()) ||
+             (old_indexing_map.isIdentity() &&
+              is_projection(new_indexing_map)) ||
+             (is_projection(old_indexing_map) && new_indexing_map.isIdentity());
+    }
+    return false;
+  }
+
+  DenseMap<Operation *, SmallVector<int32_t, 2>> reuse_candidates_;
+  DenseMap<Operation *, int32_t> output_indices_;
+};
+
+#define GEN_PASS_DEF_BUFFERREUSEPASS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+struct BufferReusePass : public impl::BufferReusePassBase<BufferReusePass> {
+  void runOnOperation() override {
+    if (!getOperation()->getAttrOfType<UnitAttr>(
+            tf_framework::TFFrameworkDialect::kTFEntryAttrName))
+      return;
+
+    BufferReuseAnalysis analysis(getOperation());
+
+    // Annotate IR with reuse candidates and output indices per allocation.
+    Builder builder(&getContext());
+    getOperation().walk([&](memref::AllocOp op) {
+      if (auto output_index = analysis.get_output_index(op)) {
+        auto attr = builder.getI32IntegerAttr(*output_index);
+        op.getOperation()->setAttr(
+            tf_framework::TFAllocOp::kReuseOutputAttrName, attr);
+      }
+      if (auto reuse_candiates = analysis.get_reuse_candiates(op)) {
+        auto attr = builder.getI32ArrayAttr(*reuse_candiates);
+        op.getOperation()->setAttr(
+            tf_framework::TFAllocOp::kReuseInputCandidatesAttrName, attr);
+      }
+    });
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> CreateBufferReusePass() {
+  return std::make_unique<BufferReusePass>();
+}
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize.cc

@@ -0,0 +1,78 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements logic for translating mixed IR to buffer form.
+
+#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"  // from @llvm-project
+
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Complex/IR/Complex.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
+#include "mlir/Dialect/SCF/IR/SCF.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/ImplicitLocOpBuilder.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "stablehlo/dialect/ChloOps.h"  // from @stablehlo
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewriters.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+namespace {
+
+struct BufferizeJITExecuteOp
+    : public OpConversionPattern<tf_framework::JITExecuteOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      tf_framework::JITExecuteOp op, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    Type result_ty = getTypeConverter()->convertType(op.getType());
+    rewriter.replaceOpWithNewOp<tf_framework::JITExecuteOp>(
+        op, result_ty, adaptor.getOperands(), op->getAttrs());
+    return success();
+  }
+};
+
+}  // namespace
+
+void populateExtraBufferizePatterns(ConversionTarget &target,
+                                    MLIRContext *context,
+                                    TypeConverter *converter,
+                                    RewritePatternSet *patterns) {
+  target.addLegalDialect<tf_framework::TFFrameworkDialect>();
+  auto typesAreLegal = [converter](Operation *op) {
+    return converter->isLegal(op->getOperandTypes()) &&
+           converter->isLegal(op->getResultTypes());
+  };
+  target.addDynamicallyLegalOp<tf_framework::JITExecuteOp>(typesAreLegal);
+  // clang-format off
+  patterns->add<
+      BufferizeJITExecuteOp
+  >(*converter, context);
+  // clang-format on
+}
+
+void populateExtraBufferizeDialects(DialectRegistry &registry) {
+  registry.insert<tf_framework::TFFrameworkDialect>();
+}
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/bufferize_pass.cc

@@ -0,0 +1,58 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file implements logic for translating mixed IR to buffer form.
+
+#include <memory>
+
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Bufferization/Transforms/Bufferize.h"  // from @llvm-project
+#include "mlir/Pass/PassManager.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewriters.h"
+#include "xla/mlir_hlo/transforms/passes.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+namespace {
+
+#define GEN_PASS_DEF_KERNELGENFINALBUFFERIZEPASS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+struct KernelgenFinalBufferizePass
+    : public impl::KernelgenFinalBufferizePassBase<
+          KernelgenFinalBufferizePass> {
+  // Default alignment_ specified in passes.td
+  KernelgenFinalBufferizePass() = default;
+
+  void runOnOperation() override {
+    mlir::PassManager pm(&getContext());
+    pm.addPass(mlir::createFinalBufferizePass(/*alignment=*/64,
+                                              populateExtraBufferizeDialects,
+                                              populateExtraBufferizePatterns));
+    (void)runPipeline(pm, getOperation());
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<ModuleOp>> CreateKernelgenFinalBufferizePass() {
+  return std::make_unique<KernelgenFinalBufferizePass>();
+}
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/copy_cleanup_pass.cc

@@ -0,0 +1,177 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+namespace {
+#define GEN_PASS_DEF_COPYCLEANUPPASS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+// A pass to remove memref::AllocOps and memref::CopyOps ops.
+//
+// The idea behind this pass is to collect all patterns we are interested in in
+// a single place. Eventually, this should be replaced by a generalized copy
+// removal pass.
+
+// Handles the pattern where an input operand of a linalg generic is copied
+// even though the producer is not mutated.
+void RemoveCopyIfTargetOnlyRead(func::FuncOp func) {
+  llvm::SmallVector<memref::AllocOp, 8> allocs_to_remove;
+  llvm::SmallVector<memref::CopyOp, 8> copies_to_remove;
+
+  // Gather all allocs and copies which are only read and have an immutable
+  // source.
+  func->walk([&](memref::AllocOp op) {
+    memref::CopyOp copy;
+    MemoryEffectOpInterface reader;
+    bool at_most_one_copy = true;
+    bool at_most_one_read = true;
+    for (auto user : op->getUsers()) {
+      if (auto copy_user = dyn_cast<memref::CopyOp>(user)) {
+        if (copy) {
+          at_most_one_copy = false;
+        } else {
+          copy = copy_user;
+        }
+        continue;
+      }
+      if (auto effect_interface = dyn_cast<MemoryEffectOpInterface>(user)) {
+        if (reader) {
+          at_most_one_read = false;
+        } else {
+          reader = effect_interface;
+        }
+        SmallVector<MemoryEffects::EffectInstance, 2> effects;
+        effect_interface.getEffectsOnValue(op.getResult(), effects);
+        if (llvm::any_of(effects, [](MemoryEffects::EffectInstance it) {
+              return !isa<MemoryEffects::Read>(it.getEffect());
+            })) {
+          at_most_one_read = false;
+        }
+        continue;
+      }
+      // We don't understand this use, be conservative.
+      at_most_one_read = false;
+    }
+    if (!copy || !at_most_one_copy) return;
+    if (!reader || !at_most_one_read) return;
+    // The copy should have the alloc op as target.
+    if (copy.getTarget() != op.getResult()) return;
+
+    // The copy should be before the reading use.
+    if (copy->getBlock() != reader->getBlock() ||
+        !copy->isBeforeInBlock(reader)) {
+      return;
+    }
+
+    // No write effects between copy and use. With aliasing information, this
+    // could be made more precise but for now we have to be conservative. The
+    // only thing we allow are writes to values that are allocated after the
+    // copy, as the aliasing is clear in those cases.
+    bool source_is_mutated = false;
+    for (Operation *pos = copy->getNextNode(), *end = reader; pos != end;
+         pos = pos->getNextNode()) {
+      auto effect_interface = dyn_cast<MemoryEffectOpInterface>(pos);
+      if (!effect_interface) {
+        continue;
+      }
+      SmallVector<MemoryEffects::EffectInstance, 2> effects;
+      effect_interface.getEffects<MemoryEffects::Write>(effects);
+      for (auto effect : effects) {
+        if (effect.getValue() != nullptr) {
+          if (auto alloc = effect.getValue().getDefiningOp<memref::AllocOp>()) {
+            if (alloc->getBlock() == copy->getBlock() &&
+                copy->isBeforeInBlock(alloc)) {
+              continue;
+            }
+          }
+        }
+        source_is_mutated = true;
+        break;
+      }
+    }
+    if (source_is_mutated) return;
+
+    op->replaceAllUsesWith(ValueRange{copy.getSource()});
+    allocs_to_remove.push_back(op);
+    copies_to_remove.push_back(copy);
+  });
+  llvm::for_each(allocs_to_remove, [](Operation *op) { op->erase(); });
+  llvm::for_each(copies_to_remove, [](Operation *op) { op->erase(); });
+}
+
+// Handles the case where the last instructions of a function implements a copy
+// back to a function argument.
+void RemoveCopyIfTargetIsFunctionArg(func::FuncOp func) {
+  // For now only support this on functions with a single block.
+  if (!func.getBody().hasOneBlock()) return;
+
+  llvm::SmallVector<memref::AllocOp> allocs_to_remove;
+  llvm::SmallVector<memref::CopyOp> copies_to_remove;
+
+  Block &body = func.getBody().front();
+  for (auto &op : llvm::reverse(body.without_terminator())) {
+    if (auto copy = dyn_cast<memref::CopyOp>(op)) {
+      auto block_arg = mlir::dyn_cast<BlockArgument>(copy.getTarget());
+      if (!block_arg) break;
+      if (!isa<func::FuncOp>(block_arg.getOwner()->getParentOp()) ||
+          !block_arg.hasOneUse())
+        break;
+      auto alloc = copy.getSource().getDefiningOp<memref::AllocOp>();
+      if (!alloc) break;
+      alloc->replaceAllUsesWith(ValueRange{block_arg});
+      allocs_to_remove.push_back(alloc);
+      copies_to_remove.push_back(copy);
+      continue;
+    }
+    break;
+  }
+  llvm::for_each(allocs_to_remove, [](Operation *op) { op->erase(); });
+  llvm::for_each(copies_to_remove, [](Operation *op) { op->erase(); });
+}
+
+}  // namespace
+
+struct CopyCleanupPass : public impl::CopyCleanupPassBase<CopyCleanupPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<memref::MemRefDialect>();
+  }
+
+  void runOnOperation() override {
+    RemoveCopyIfTargetOnlyRead(getOperation());
+    RemoveCopyIfTargetIsFunctionArg(getOperation());
+  }
+};
+
+std::unique_ptr<OperationPass<func::FuncOp>> CreateCopyCleanupPass() {
+  return std::make_unique<CopyCleanupPass>();
+}
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/embed_tf_framework.cc

@@ -0,0 +1,199 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <optional>
+
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/TypeRange.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewriters.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace tf_framework {
+namespace {
+
+// Prepends argument type list of the function with an OpKernelContextType arg.
+class FuncOpConverter : public OpConversionPattern<func::FuncOp> {
+ public:
+  using OpConversionPattern<func::FuncOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      func::FuncOp func, OpAdaptor /*adaptor*/,
+      ConversionPatternRewriter &rewriter) const override {
+    // Convert function arguments using the provided TypeConverter.
+    auto func_type = func.getFunctionType();
+    TypeConverter::SignatureConversion conversion(func_type.getNumInputs());
+
+    conversion.addInputs(OpKernelContextType::get(rewriter.getContext()));
+    for (auto arg_type : llvm::enumerate(func_type.getInputs())) {
+      conversion.addInputs(arg_type.index(), arg_type.value());
+    }
+
+    rewriter.applySignatureConversion(&func.getBody().front(), conversion);
+
+    // Update the signature of the function.
+    rewriter.modifyOpInPlace(func, [&] {
+      func.setType(rewriter.getFunctionType(conversion.getConvertedTypes(),
+                                            func_type.getResults()));
+    });
+    return success();
+  }
+};
+
+std::optional<Value> FindOpKernelContext(Operation *op) {
+  auto func = op->getParentOfType<func::FuncOp>();
+  if (func.getNumArguments() == 0) {
+    return std::nullopt;
+  }
+  Value ctx = func.getArgument(0);
+  if (!mlir::isa<OpKernelContextType>(ctx.getType())) {
+    return std::nullopt;
+  }
+  return ctx;
+}
+
+// Converts std.alloc to tf_framework.alloc_raw using OpKernelContextType arg of
+// the parent function.
+struct AllocOpConverter : public OpConversionPattern<memref::AllocOp> {
+  using OpConversionPattern<memref::AllocOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      memref::AllocOp alloc, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    std::optional<Value> ctx = FindOpKernelContext(alloc);
+    if (!ctx) return failure();
+
+    // Symbolic operands that bind to the symbols of the memref's layout map are
+    // not supported by TFAllocOp.
+    if (!alloc.getSymbolOperands().empty()) {
+      return failure();
+    }
+    auto reuse_input_candidates = alloc->getAttrOfType<ArrayAttr>(
+        TFAllocOp::kReuseInputCandidatesAttrName);
+    auto reuse_output_index =
+        alloc->getAttrOfType<IntegerAttr>(TFAllocOp::kReuseOutputAttrName);
+    Value buffer = rewriter.replaceOpWithNewOp<TFAllocOp>(
+        alloc, alloc.getType(), *ctx, adaptor.getOperands(),
+        reuse_input_candidates, reuse_output_index);
+    Location loc = buffer.getLoc();
+    Value cond = IsValidMemRefOp::create(rewriter, loc,
+                                         rewriter.getIntegerType(1), buffer);
+    TFAssertOp::create(rewriter, loc, *ctx, cond, ErrorCode::RESOURCE_EXHAUSTED,
+                       "failed to allocate memory");
+    return success();
+  }
+};
+
+// Converts std.dealloc to tf_framework.dealloc_raw using OpKernelContextType
+// arg of the parent function.
+struct DeallocOpConverter : public OpConversionPattern<memref::DeallocOp> {
+  using OpConversionPattern<memref::DeallocOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      memref::DeallocOp dealloc, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    std::optional<Value> ctx = FindOpKernelContext(dealloc);
+    if (!ctx) return failure();
+
+    // Operand with no layout is expected.
+    auto operand_memref_type =
+        mlir::cast<MemRefType>(dealloc.getMemref().getType());
+    if (!operand_memref_type.getLayout().isIdentity()) {
+      return failure();
+    }
+    rewriter.replaceOpWithNewOp<TFDeallocOp>(dealloc, *ctx,
+                                             adaptor.getMemref());
+    return success();
+  }
+};
+
+// Converts std.assert to tf_framework.assert with using OpKernelContextType
+// arg of the parent function.
+struct AssertOpConverter : public OpConversionPattern<cf::AssertOp> {
+ public:
+  using OpConversionPattern<cf::AssertOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      cf::AssertOp op, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    std::optional<Value> ctx = FindOpKernelContext(op);
+    if (!ctx) return failure();
+    rewriter.replaceOpWithNewOp<TFAssertOp>(op, *ctx, adaptor.getArg(),
+                                            ErrorCode::INVALID_ARGUMENT,
+                                            adaptor.getMsg());
+    return success();
+  }
+};
+
+// Amends `tf_framework.jit_execute` with the newly introduced OpKernelContext.
+struct JITExecuteOpConverter : public OpConversionPattern<JITExecuteOp> {
+  using OpConversionPattern<JITExecuteOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      JITExecuteOp op, OpAdaptor /*adaptor*/,
+      ConversionPatternRewriter &rewriter) const override {
+    std::optional<Value> ctx = FindOpKernelContext(op);
+    if (!ctx) return failure();
+    rewriter.replaceOpWithNewOp<JITExecuteOp>(
+        op, op.getResult().getType(), *ctx, op.getCallable(), op.getInputs());
+    return success();
+  }
+};
+
+// Amends `tf_framework.jit_compile_from_str` with the newly introduced
+// OpKernelContext.
+struct JITCompileFromStrOpConverter
+    : public OpConversionPattern<JITCompileFromStrOp> {
+  using OpConversionPattern<JITCompileFromStrOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      JITCompileFromStrOp op, OpAdaptor /*adaptor*/,
+      ConversionPatternRewriter &rewriter) const override {
+    std::optional<Value> ctx = FindOpKernelContext(op);
+    if (!ctx) return failure();
+    rewriter.replaceOpWithNewOp<JITCompileFromStrOp>(
+        op, rewriter.getType<JITCallableType>(), *ctx, op->getAttrs());
+    return success();
+  }
+};
+
+}  // namespace
+
+void PopulateEmbedTFFrameworkAssertPattern(RewritePatternSet *patterns) {
+  patterns->add<AssertOpConverter>(patterns->getContext());
+}
+
+void PopulateEmbedTFFrameworkPatterns(RewritePatternSet *patterns) {
+  // clang-format off
+  patterns->add<
+      AllocOpConverter,
+      AssertOpConverter,
+      DeallocOpConverter,
+      FuncOpConverter,
+      JITCompileFromStrOpConverter,
+      JITExecuteOpConverter>(patterns->getContext());
+  // clang-format on
+}
+
+}  // namespace tf_framework
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/embed_tf_framework_pass.cc

@@ -0,0 +1,98 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewriters.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace tf_framework {
+namespace {
+
+#define GEN_PASS_DEF_EMBEDTFFRAMEWORKPASS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+bool IsNotInsideTfEntryFunction(Operation* op) {
+  auto func = op->getParentOfType<func::FuncOp>();
+  return !func->hasAttrOfType<UnitAttr>(TFFrameworkDialect::kTFEntryAttrName);
+}
+
+template <typename OpTy>
+bool HasInitializedOpKernelContextOperand(OpTy op) {
+  return op.getCtx() != nullptr;
+}
+
+// The pass rewrites the function marked with `tf_entry` attribute.
+// * adds tf_framework::OpKernelContextType argument to the function,
+// * std.alloc becomes tf_framework.alloc_raw,
+// * std.dealloc becomes tf_framework.dealloc_raw.
+class EmbedTFFrameworkPass
+    : public impl::EmbedTFFrameworkPassBase<EmbedTFFrameworkPass> {
+  void getDependentDialects(DialectRegistry& registry) const override {
+    registry.insert<mlir::kernel_gen::tf_framework::TFFrameworkDialect>();
+  }
+
+ public:
+  void runOnOperation() override {
+    ModuleOp m = getOperation();
+
+    // Populate patterns.
+    RewritePatternSet patterns(&getContext());
+    PopulateEmbedTFFrameworkPatterns(&patterns);
+
+    // Set target.
+    ConversionTarget target(getContext());
+    target.addLegalDialect<tf_framework::TFFrameworkDialect>();
+
+    target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
+      if (!op->hasAttrOfType<UnitAttr>(TFFrameworkDialect::kTFEntryAttrName)) {
+        return true;
+      }
+      FunctionType func_type = op.getFunctionType();
+      return func_type.getNumInputs() > 0 &&
+             mlir::isa<OpKernelContextType>(func_type.getInput(0));
+    });
+    target.addDynamicallyLegalOp<cf::AssertOp, memref::AllocOp,
+                                 memref::DeallocOp>(IsNotInsideTfEntryFunction);
+    target.addDynamicallyLegalOp<JITExecuteOp>(
+        &HasInitializedOpKernelContextOperand<JITExecuteOp>);
+    target.addDynamicallyLegalOp<JITCompileFromStrOp>(
+        &HasInitializedOpKernelContextOperand<JITCompileFromStrOp>);
+
+    if (failed(applyPartialConversion(m, target, std::move(patterns)))) {
+      signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<ModuleOp> > CreateEmbedTFFrameworkPass() {
+  return std::make_unique<EmbedTFFrameworkPass>();
+}
+
+}  // namespace tf_framework
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/func_to_jit_invocations.cc

@@ -0,0 +1,275 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cassert>
+#include <cstdint>
+#include <memory>
+#include <string>
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
+#include "mlir/IR/Block.h"  // from @llvm-project
+#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/IR/TypeRange.h"  // from @llvm-project
+#include "mlir/IR/Visitors.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LogicalResult.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+
+constexpr llvm::StringRef
+    mlir::kernel_gen::tf_framework ::JITCompileFromStrOp::kJITEntryFunctionName;
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+namespace {
+
+constexpr int64_t i32Limit = 2147483647;
+using func::FuncOp;
+using shape::ShapeOfOp;
+
+LogicalResult RewriteToFullJit(func::FuncOp op) {
+  // Rewrite all functions that have a single result.
+  if (op.getResultTypes().size() != 1) return failure();
+
+  IRRewriter rewriter(op.getContext());
+
+  // Insert a new block at the front of the function.
+  Block *old_body = &op.getFunctionBody().front();
+  Location loc = op.getLoc();
+  llvm::SmallVector<Location> locs(old_body->getNumArguments(), loc);
+  Block *new_body = rewriter.createBlock(&op.getBody(), op.getBody().begin(),
+                                         old_body->getArgumentTypes(), locs);
+
+  // Create the JIT compile op.
+  auto jit_compile_op = tf_framework::JITCompileOp::create(
+      rewriter, loc, rewriter.getType<tf_framework::JITCallableType>(),
+      /*ctx=*/mlir::Value());
+
+  // Move the original functions operations into the body.
+  {
+    OpBuilder::InsertionGuard guard(rewriter);
+    Block *jit_block = rewriter.createBlock(&jit_compile_op.getBody(), {},
+                                            old_body->getArgumentTypes(), locs);
+
+    rewriter.inlineBlockBefore(old_body, jit_block, jit_block->begin(),
+                               jit_block->getArguments());
+
+    Operation *terminator = jit_block->getTerminator();
+    rewriter.setInsertionPointAfter(terminator);
+    tf_framework::JITCompileYieldOp::create(rewriter, loc,
+                                            terminator->getOperands().front());
+    terminator->erase();
+  }
+
+  // Create JIT execute op.
+  auto execute = tf_framework::JITExecuteOp::create(
+      rewriter, loc, op.getResultTypes().front(), /*ctx=*/Value(),
+      jit_compile_op.getResult(), new_body->getArguments());
+
+  // Create a return.
+  func::ReturnOp::create(rewriter, loc, execute.getResult());
+  return success();
+}
+
+LogicalResult RewriteToLargeSizeJit(FuncOp op) {
+  // Rewrite all functions that have at most two arguments and a single result.
+  if (op.getArgumentTypes().size() > 2 || op.getResultTypes().size() != 1)
+    return failure();
+
+  IRRewriter rewriter(op.getContext());
+  Location loc = op.getLoc();
+
+  // Insert a new block at the front of the function.
+  Block *old_body = &op.getFunctionBody().front();
+  llvm::SmallVector<Location> locs(old_body->getNumArguments(), loc);
+  Block *new_body = rewriter.createBlock(&op.getBody(), op.getBody().begin(),
+                                         old_body->getArgumentTypes(), locs);
+
+  // Create large argument condition.
+  auto arg_1 = new_body->getArgument(0);
+  auto shape_1 = shape::ShapeOfOp::create(rewriter, loc, arg_1);
+  auto num_elems_1 = shape::NumElementsOp::create(rewriter, loc, shape_1);
+  Value cst_i32_limit = arith::ConstantIndexOp::create(rewriter, loc, i32Limit);
+  Value large_tensor_predicate = arith::CmpIOp::create(
+      rewriter, loc, arith::CmpIPredicate::sgt, num_elems_1, cst_i32_limit);
+  if (new_body->getNumArguments() > 1) {
+    auto arg_2 = new_body->getArgument(1);
+    auto shape_2 = shape::ShapeOfOp::create(rewriter, loc, arg_2);
+    auto num_elems_2 = shape::NumElementsOp::create(rewriter, loc, shape_2);
+    large_tensor_predicate = arith::OrIOp::create(
+        rewriter, loc, large_tensor_predicate,
+        // Compare op to check size of the second op
+        arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::sgt,
+                              num_elems_2, cst_i32_limit));
+  }
+
+  // Create dispatch code.
+  auto jit_body_builder_fn = [&](OpBuilder &b, Location loc) {
+    // Create JIT compile op.
+    auto callable_ty = b.getType<tf_framework::JITCallableType>();
+    auto jit_compile_op = tf_framework::JITCompileOp::create(
+        b, loc, callable_ty, /*ctx=*/Value());
+    {
+      OpBuilder::InsertionGuard g(b);
+      Block *block = b.createBlock(
+          &jit_compile_op.getBody(), {}, new_body->getArgumentTypes(),
+          SmallVector<Location>(new_body->getNumArguments(), loc));
+      b.setInsertionPointToStart(block);
+      IRMapping bvm;
+      bvm.map(old_body->getArguments(), block->getArguments());
+      for (auto &op : old_body->without_terminator()) {
+        b.clone(op, bvm);
+      }
+      tf_framework::JITCompileYieldOp::create(
+          b, loc, block->back().getResults().front());
+    }
+
+    // Create JIT execute op.
+    auto jit_execute_op = tf_framework::JITExecuteOp::create(
+        b, loc, op.getResultTypes().front(), /*ctx=*/Value(),
+        jit_compile_op.getResult(), new_body->getArguments());
+    scf::YieldOp::create(b, loc, jit_execute_op.getResult());
+  };
+  auto aot_body_builder_fn = [&](OpBuilder &b, Location loc) {
+    IRMapping bvm;
+    bvm.map(old_body->getArguments(), new_body->getArguments());
+    Operation *last_clone;
+    for (auto &op : old_body->without_terminator()) {
+      last_clone = b.clone(op, bvm);
+    }
+    scf::YieldOp::create(b, loc, last_clone->getResults().front());
+  };
+
+  // Create the conditional and return operation.
+  auto ifOp = scf::IfOp::create(rewriter, loc, large_tensor_predicate,
+                                jit_body_builder_fn, aot_body_builder_fn);
+  func::ReturnOp::create(rewriter, loc, ifOp.getResults().front());
+
+  // Remove the old body.
+  rewriter.eraseBlock(old_body);
+  return success();
+}
+
+void PackJITCompileOp(tf_framework::JITCompileOp op,
+                      llvm::ArrayRef<int64_t> tile_sizes,
+                      llvm::ArrayRef<int64_t> unroll_factors, bool enable_ftz,
+                      bool index_64bit, bool cpu_codegen) {
+  IRRewriter rewriter(op.getContext());
+  Block *body = op.SingleBlock::getBody();
+  auto yield_op =
+      llvm::cast<tf_framework::JITCompileYieldOp>(body->getTerminator());
+
+  // Temporarily, build the module that would be JIT-compiled. This is only to
+  // obtain the serialized code attribute.
+  auto loc = op->getLoc();
+  auto jit_module = ModuleOp::create(rewriter, loc);
+  {
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPointToStart(jit_module.SingleBlock::getBody());
+    auto jit_function = func::FuncOp::create(
+        rewriter, loc, tf_framework::JITCompileFromStrOp::kJITEntryFunctionName,
+        rewriter.getFunctionType(body->getArgumentTypes(),
+                                 yield_op->getOperandTypes()));
+    jit_function->setAttr(tf_framework::TFFrameworkDialect::kTFEntryAttrName,
+                          rewriter.getUnitAttr());
+    jit_function.getBody().takeBody(op.getBodyRegion());
+    rewriter.setInsertionPointToEnd(&jit_function.getBody().front());
+    func::ReturnOp::create(rewriter, loc, yield_op.getResult());
+    rewriter.eraseOp(yield_op);
+  }
+
+  // Serialize JIT module.
+  std::string code;
+  llvm::raw_string_ostream ss(code);
+  assert(succeeded(jit_module.verify()));
+  mlir::OpPrintingFlags flags;
+  jit_module.print(ss, flags.assumeVerified());
+
+  // Remove temporary module.
+  rewriter.eraseOp(jit_module);
+
+  // Finally, create the new JIT compile op.
+  rewriter.setInsertionPointAfter(op);
+  rewriter.replaceOpWithNewOp<tf_framework::JITCompileFromStrOp>(
+      op, op->getResultTypes(), op.getCtx(), rewriter.getStringAttr(code),
+      rewriter.getI64ArrayAttr(tile_sizes),
+      rewriter.getI64ArrayAttr(unroll_factors),
+      rewriter.getBoolAttr(enable_ftz), rewriter.getBoolAttr(index_64bit),
+      rewriter.getBoolAttr(cpu_codegen));
+}
+
+#define GEN_PASS_DEF_FUNCTOJITINVOCATIONPASS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+struct FuncToJITInvocationPass
+    : public impl::FuncToJITInvocationPassBase<FuncToJITInvocationPass> {
+  explicit FuncToJITInvocationPass(llvm::ArrayRef<int64_t> tile_sizes,
+                                   llvm::ArrayRef<int64_t> unroll_factors,
+                                   bool enable_ftz, bool index_64bit,
+                                   bool cpu_codegen,
+                                   bool jit_i64_indexed_for_large_tensors) {
+    tile_sizes_ = tile_sizes;
+    unroll_factors_ = unroll_factors;
+    enable_ftz_ = enable_ftz;
+    index_64bit_ = index_64bit;
+    cpu_codegen_ = cpu_codegen;
+    jit_i64_indexed_for_large_tensors_ = jit_i64_indexed_for_large_tensors;
+  }
+
+  void runOnOperation() override {
+    if (jit_i64_indexed_for_large_tensors_) {
+      if (failed(RewriteToLargeSizeJit(getOperation()))) {
+        return signalPassFailure();
+      }
+    } else {
+      if (failed(RewriteToFullJit(getOperation()))) {
+        return signalPassFailure();
+      }
+    }
+
+    getOperation().walk([&](tf_framework::JITCompileOp op) {
+      PackJITCompileOp(op, tile_sizes_, unroll_factors_, enable_ftz_,
+                       index_64bit_ || jit_i64_indexed_for_large_tensors_,
+                       cpu_codegen_);
+    });
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>> CreateFuncToJITInvocationPass(
+    llvm::ArrayRef<int64_t> tile_sizes, llvm::ArrayRef<int64_t> unroll_factors,
+    bool enable_ftz, bool index_64bit, bool cpu_codegen,
+    bool jit_i64_indexed_for_large_tensors) {
+  return std::make_unique<FuncToJITInvocationPass>(
+      tile_sizes, unroll_factors, enable_ftz, index_64bit, cpu_codegen,
+      jit_i64_indexed_for_large_tensors);
+}
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/fuse_inner_parallel_loops_pass.cc

@@ -0,0 +1,54 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+
+#include "mlir/Analysis/AliasAnalysis.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/SCF/IR/SCF.h"  // from @llvm-project
+#include "mlir/Dialect/SCF/Transforms/Transforms.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+namespace {
+
+#define GEN_PASS_DEF_FUSEINNERPARALLELLOOPSPASS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+struct FuseInnerParallelLoopsPass
+    : impl::FuseInnerParallelLoopsPassBase<FuseInnerParallelLoopsPass> {
+  void runOnOperation() override {
+    auto &alias_analysis = getAnalysis<AliasAnalysis>();
+    auto may_alias = [&](Value val1, Value val2) -> bool {
+      return !alias_analysis.alias(val1, val2).isNo();
+    };
+    getOperation().walk([&](mlir::scf::ParallelOp op) {
+      mlir::scf::naivelyFuseParallelOps(op.getRegion(), may_alias);
+    });
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateFuseInnerParallelLoopsPass() {
+  return std::make_unique<FuseInnerParallelLoopsPass>();
+}
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/gpu_kernel_to_blob_pass.cc

@@ -0,0 +1,250 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <algorithm>
+#include <iterator>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "llvm/TargetParser/Triple.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Target/LLVMIR/Export.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+#include "xla/debug_options_flags.h"
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+#include "xla/service/gpu/gpu_asm_opts_util.h"
+#include "xla/service/gpu/target_constants.h"
+#include "xla/stream_executor/device_description.h"
+#include "xla/xla.pb.h"
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/path.h"
+#include "tensorflow/core/platform/status.h"
+#include "tensorflow/core/platform/statusor.h"
+
+#if GOOGLE_CUDA
+#include "xla/service/gpu/llvm_gpu_backend/nvptx_backend.h"
+#include "xla/stream_executor/cuda/cuda_asm_compiler.h"
+#elif TENSORFLOW_USE_ROCM
+#include "xla/service/gpu/llvm_gpu_backend/amdgpu_backend.h"
+#include "xla/stream_executor/gpu/asm_compiler.h"
+#include "tensorflow/core/platform/rocm_rocdl_path.h"
+#endif
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+namespace {
+
+#define GEN_PASS_DEF_GPUKERNELTOBLOBPASS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+class GpuKernelToBlobPass
+    : public impl::GpuKernelToBlobPassBase<GpuKernelToBlobPass> {
+ public:
+  GpuKernelToBlobPass(StringRef blob_annotation,
+                      llvm::ArrayRef<std::string> architectures, bool print_ptx,
+                      bool print_llvmir, bool enable_ftz) {
+    if (!blob_annotation.empty()) blob_annotation_ = blob_annotation.str();
+    architectures_ = architectures;
+    print_ptx_ = print_ptx;
+    print_llvmir_ = print_llvmir;
+    enable_ftz_ = enable_ftz;
+  }
+
+  void runOnOperation() override {
+    gpu::GPUModuleOp gpu_module = getOperation();
+    auto blob_or = GetGpuBinaryBlob(gpu_module);
+    if (blob_or.ok()) {
+      const auto& blob = blob_or.value();
+      std::string blob_string(blob.begin(), blob.end());
+      gpu_module->setAttr(blob_annotation_,
+                          StringAttr::get(&getContext(), blob_string));
+      return;
+    }
+    // Forward the error by attaching the message to the gpu module.
+    gpu_module.emitError(blob_or.status().message());
+    return signalPassFailure();
+  }
+
+  absl::StatusOr<std::vector<uint8_t>> GetGpuBinaryBlob(
+      gpu::GPUModuleOp gpu_module) {
+    if (architectures_.empty()) {
+      return absl::InternalError("Expected at least one GPU architecture.");
+    }
+
+    // Lower to LLVM module.
+    llvm::LLVMContext llvmContext;
+    auto llvmModule = translateModuleToLLVMIR(gpu_module, llvmContext);
+    if (!llvmModule) {
+      return absl::InternalError("Could not translate MLIR module to LLVM IR");
+    }
+    llvmModule->setModuleIdentifier(gpu_module.getName());
+
+#if TENSORFLOW_USE_ROCM
+    xla::DebugOptions options = xla::GetDebugOptionsFromFlags();
+    options.set_xla_gpu_ftz(enable_ftz_);
+    options.set_xla_gpu_dump_llvmir(print_llvmir_);
+
+    using AmdGpuHsaco = std::vector<tensorflow::uint8>;
+    std::vector<tensorflow::se::HsacoImage> images;
+    images.reserve(architectures_.size());
+    for (const std::string& arch_str : architectures_) {
+      // Parse ROCm architecture.
+      absl::string_view consumable_arch(arch_str);
+      if (!absl::ConsumePrefix(&consumable_arch, "gfx")) {
+        return tensorflow::errors::Internal(
+            "Could not parse ROCm architecture prefix (expected gfx)");
+      }
+      auto llvm_module_copy = llvm::CloneModule(*llvmModule);
+      auto hsaco_or = xla::gpu::amdgpu::CompileToHsaco(
+          llvm_module_copy.get(),
+          tensorflow::se::RocmComputeCapability{arch_str}, options,
+          options.DebugString());
+      if (!hsaco_or.ok()) {
+        return tensorflow::errors::Internal("Failure when generating HSACO");
+      }
+      auto hsaco = hsaco_or.value();
+      images.push_back({arch_str, std::move(hsaco)});
+    }
+
+    // TODO(b/169870789): Revisit the use of fatbins.
+    // Bundle HSACO images into a single fatbin.
+    if (images.size() == 1) return images.front().bytes;
+    return tensorflow::se::BundleGpuAsm(images, tensorflow::RocmRoot());
+
+#elif GOOGLE_CUDA
+    xla::DebugOptions options = xla::GetDebugOptionsFromFlags();
+    options.set_xla_gpu_ftz(enable_ftz_);
+    options.set_xla_gpu_dump_llvmir(print_llvmir_);
+    // Make sure we use full precision division operations.
+    (*options.mutable_xla_backend_extra_options())["-nvptx-prec-divf32"] = "2";
+    // Disable tail sinking as it interferes with load/store vectorization. If
+    // we have common tails that is intentional.
+    (*options.mutable_xla_backend_extra_options())["-simplifycfg-sink-common"] =
+        "false";
+
+    llvmModule->setDataLayout(xla::gpu::nvptx::DataLayout());
+    llvmModule->setTargetTriple(llvm::Triple(xla::gpu::nvptx::TargetTriple()));
+
+    // Compile and collect requested cubin and PTX images.
+    std::vector<tensorflow::se::CubinOrPTXImage> images;
+    auto gpu_asm_opts = xla::gpu::PtxOptsFromDebugOptions(options);
+    for (const std::string& arch_str : architectures_) {
+      TF_ASSIGN_OR_RETURN(auto arch_pair, ParseCudaArch(arch_str));
+      bool is_compute_profile = arch_pair.first;
+      int arch = arch_pair.second;
+      int cc_major = arch / 10;
+      int cc_minor = arch % 10;
+      tensorflow::se::CudaComputeCapability cc{cc_major, cc_minor};
+
+      // Generate PTX code.
+      // Module may be changed by CompileToPtx.
+      auto llvm_module_copy = llvm::CloneModule(*llvmModule);
+      auto enable_fusion = [](llvm::TargetMachine* target) {
+        target->Options.AllowFPOpFusion =
+            llvm::FPOpFusion::FPOpFusionMode::Fast;
+      };
+      TF_ASSIGN_OR_RETURN(
+          std::string ptx,
+          xla::gpu::nvptx::CompileToPtx(
+              llvm_module_copy.get(), stream_executor::GpuComputeCapability(cc),
+              options, enable_fusion));
+      if (print_ptx_) {
+        llvm::dbgs() << "Generated PTX code for module '"
+                     << gpu_module.getName() << "' on architecture sm_" << arch
+                     << ":\n";
+        llvm::dbgs() << ptx << "\n";
+      }
+
+      // Compile PTX code with ptxas if requested and possible and fall back to
+      // a compute image, otherwise.
+      if (!is_compute_profile) {
+        auto gpu_asm = tensorflow::se::CompileGpuAsm(cc, ptx, gpu_asm_opts);
+        if (gpu_asm.ok()) {
+          images.push_back({/*is_ptx=*/false, cc, std::move(gpu_asm.value())});
+        } else {
+#ifdef PLATFORM_GOOGLE
+          // Require compilation with ptxas.
+          return gpu_asm;
+#else
+          // Fall back to compilation by driver in OSS.
+          LOG(WARNING) << "Failed to compile generated PTX with ptxas. Falling "
+                          "back to compilation by driver.";
+          is_compute_profile = true;
+#endif
+        }
+      }
+      if (is_compute_profile) {
+        std::vector<uint8_t> ptx_bytes;
+        ptx_bytes.reserve(ptx.size() + 1);
+        std::copy(ptx.begin(), ptx.end(), std::back_inserter(ptx_bytes));
+        ptx_bytes.push_back('\0');
+        images.push_back({/*is_ptx=*/true, cc, std::move(ptx_bytes)});
+      }
+    }
+
+    // TODO(b/169870789): Revisit the use of fatbins.
+    // Bundle cubin and PTX images into a single fatbin if needed.
+    if (images.size() == 1) return images.front().bytes;
+    return tensorflow::se::BundleGpuAsm(images, gpu_asm_opts);
+
+#else
+    return tensorflow::errors::Internal(
+        "Neither TENSORFLOW_USE_ROCM nor GOOGLE_CUDA are defined."
+        " Did you specify either --config=rocm or --config=cuda ?");
+#endif
+  }
+
+ private:
+  absl::StatusOr<std::pair<bool, int>> ParseCudaArch(
+      const std::string& arch_str) {
+    absl::string_view consumable_arch(arch_str);
+    bool is_compute_profile;
+    if (absl::ConsumePrefix(&consumable_arch, "compute_")) {
+      is_compute_profile = true;
+    } else if (absl::ConsumePrefix(&consumable_arch, "sm_")) {
+      is_compute_profile = false;
+    } else {
+      return absl::InternalError(
+          "Could not parse cuda architecture prefix (expected sm_ or "
+          "compute_)");
+    }
+    int arch;
+    if (!absl::SimpleAtoi(consumable_arch, &arch)) {
+      return absl::InternalError("Could not parse cuda architecture number");
+    }
+    return std::pair<bool, int>(is_compute_profile, arch);
+  }
+
+  bool enable_ftz_;
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<gpu::GPUModuleOp>> CreateGpuKernelToBlobPass(
+    StringRef blob_annotation, ArrayRef<std::string> architectures,
+    bool print_ptx, bool print_llvmir, bool enable_ftz) {
+  return std::make_unique<GpuKernelToBlobPass>(
+      blob_annotation, architectures, print_ptx, print_llvmir, enable_ftz);
+}
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/legalize_tensor_reshape_pass.cc

@@ -0,0 +1,73 @@
+/* Copyright 2025 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+==============================================================================*/
+
+#include <cassert>
+#include <utility>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project  // IWYU pragma: keep
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+
+namespace mlir {
+namespace kernel_gen {
+
+#define GEN_PASS_DEF_LEGALIZETENSORRESHAPEPASS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+namespace {
+
+struct LegalizeTensorReshapePattern
+    : public OpRewritePattern<tensor::ReshapeOp> {
+  explicit LegalizeTensorReshapePattern(MLIRContext* context)
+      : OpRewritePattern<tensor::ReshapeOp>(context) {}
+
+  LogicalResult matchAndRewrite(tensor::ReshapeOp op,
+                                PatternRewriter& rewriter) const override {
+    rewriter.replaceOpWithNewOp<mhlo::DynamicReshapeOp>(op, op.getResultType(),
+                                                        op.getOperands());
+    return success();
+  }
+};
+
+}  // namespace
+
+struct LegalizeTensorReshapePass
+    : public impl::LegalizeTensorReshapePassBase<LegalizeTensorReshapePass> {
+  void getDependentDialects(DialectRegistry& registry) const override {
+    registry.insert<mhlo::MhloDialect>();
+  }
+
+  void runOnOperation() override {
+    MLIRContext* ctx = &getContext();
+    RewritePatternSet patterns(ctx);
+    patterns.add<LegalizeTensorReshapePattern>(ctx);
+    GreedyRewriteConfig config;
+    config.setMaxIterations(GreedyRewriteConfig::kNoLimit);
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns),
+                                     config))) {
+      return signalPassFailure();
+    }
+  }
+};
+
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/merge_assuming_ops_pass.cc

@@ -0,0 +1,476 @@
+/* Copyright 2021 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+==============================================================================*/
+
+#include <cassert>
+#include <cstddef>
+#include <tuple>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
+#include "mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
+#include "mlir/IR/Block.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/IRMapping.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/OpDefinition.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/OperationSupport.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Interfaces/InferTypeOpInterface.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "stablehlo/dialect/Base.h"  // from @stablehlo
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+
+namespace mlir {
+namespace kernel_gen {
+
+#define GEN_PASS_DEF_MERGEASSUMINGOPSPASS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+namespace {
+
+struct ShapeReificationPattern : public OpRewritePattern<shape::ShapeOfOp> {
+  explicit ShapeReificationPattern(MLIRContext *context)
+      : OpRewritePattern<shape::ShapeOfOp>(context) {
+    // Recursively reify until we hit an op that doesn't support it.
+    setHasBoundedRewriteRecursion();
+  }
+
+  LogicalResult matchAndRewrite(shape::ShapeOfOp op,
+                                PatternRewriter &rewriter) const override {
+    // Only reify shape computation if operand allows for it.
+    auto shapeOrigin = op.getArg().getDefiningOp<InferShapedTypeOpInterface>();
+    if (!shapeOrigin) return failure();
+
+    llvm::SmallVector<Value, 1> reifications;
+    if (failed(shapeOrigin.reifyReturnTypeShapes(
+            rewriter, shapeOrigin->getOperands(), reifications)))
+      return failure();
+    assert(reifications.size() == 1);
+    Value reifiedShape = reifications.front();
+
+    // Insert cast if needed.
+    if (reifiedShape.getType() != op.getType()) {
+      reifiedShape = tensor::CastOp::create(rewriter, op.getLoc(), op.getType(),
+                                            reifiedShape);
+    }
+
+    rewriter.replaceOp(op, reifiedShape);
+    return success();
+  }
+};
+
+template <typename OpTy>
+struct InlineBroadcastedShapeOperandsPattern : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const override {
+    // Find all the shape operands, direct and indirect.
+    SmallVector<Value, 8> inlinedOperands;
+    for (Value direct : op->getOperands()) {
+      if (auto bcastOp = direct.getDefiningOp<shape::BroadcastOp>()) {
+        for (Value indirect : bcastOp->getOperands())
+          inlinedOperands.push_back(indirect);
+      } else {
+        inlinedOperands.push_back(direct);
+      }
+    }
+
+    // Only rewrite if it makes a difference.
+    if (inlinedOperands.size() == op.getNumOperands()) return failure();
+
+    // Inline shape operands.
+    rewriter.replaceOpWithNewOp<OpTy>(op, op->getResultTypes(), inlinedOperands,
+                                      op->getAttrs());
+    return success();
+  }
+};
+
+LogicalResult moveUpIntoAssumingOpMatchAndRewrite(Operation *op,
+                                                  PatternRewriter &rewriter) {
+  // Only implemented for single-result ops.
+  if (op->getNumResults() != 1) return failure();
+
+  // Find a preceding `assuming` op.
+  auto *theBlock = op->getBlock();
+  Operation *prev = op->getPrevNode();
+  while (prev != nullptr && !llvm::isa<shape::AssumingOp>(prev))
+    prev = prev->getPrevNode();
+  auto assumingOp = llvm::dyn_cast_or_null<shape::AssumingOp>(prev);
+  if (!assumingOp) return failure();
+  assert(assumingOp->getBlock() == theBlock && op->getBlock() == theBlock &&
+         "expect assuming op and root op to be in the same block");
+
+  // Make sure that all operands will be available after moving.
+  auto isAvailable = [&](Value v) {
+    Operation *def = v.getDefiningOp();
+    return def == nullptr || def->getBlock() != theBlock ||
+           !assumingOp->isBeforeInBlock(def);
+  };
+  if (!llvm::all_of(op->getOperands(), isAvailable)) return failure();
+
+  Block *body = assumingOp.getBody();
+  auto yieldOp = llvm::cast<shape::AssumingYieldOp>(body->getTerminator());
+
+  // Find the operands to use if the op was within the assuming region. We
+  // will later use their copies, as we copy the assuming op and its body.
+  SmallVector<Value, 8> newOperandsUnmapped =
+      llvm::to_vector<8>(llvm::map_range(op->getOperands(), [&](Value v) {
+        for (const auto &result : llvm::enumerate(assumingOp->getResults())) {
+          if (result.value() == v) return yieldOp->getOperand(result.index());
+        }
+        return v;
+      }));
+
+  // Insert the rewritten assuming op right before the old one.
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPoint(assumingOp);
+  auto newAssumingOp = shape::AssumingOp::create(
+      rewriter, assumingOp.getLoc(), assumingOp.getWitness(),
+      [&](OpBuilder& b, Location) {
+        // Copy body.
+        IRMapping mapping;
+        for (auto &nested : body->without_terminator())
+          b.clone(nested, mapping);
+
+        // Copy op into the new body and use the mapped operands.
+        for (auto it : llvm::zip(op->getOperands(), newOperandsUnmapped)) {
+          Value oldOperand, newOperandUnmapped;
+          std::tie(oldOperand, newOperandUnmapped) = it;
+          mapping.map(oldOperand, mapping.lookupOrDefault(newOperandUnmapped));
+        }
+        Operation *newOp = b.clone(*op, mapping);
+
+        // Yield the previous results and also the new ones.
+        auto mappedResults = llvm::to_vector<8>(llvm::map_range(
+            yieldOp.getOperands(),
+            [&](Value v) { return mapping.lookupOrDefault(v); }));
+        mappedResults.append(newOp->getResults().begin(),
+                             newOp->getResults().end());
+        return mappedResults;
+      });
+
+  // Replace the assuming op and the root op with the corresponding result
+  // values.
+  ValueRange newAssumingOpResults = newAssumingOp->getResults();
+  rewriter.replaceOp(assumingOp, newAssumingOpResults.drop_back());
+  rewriter.replaceOp(op, newAssumingOpResults.back());
+  return success();
+}
+
+/// Move operation into a preceding assuming op. This allows to process
+/// operations that depend on the assuming op's results. It will eventually
+/// allow to make assuming regions' constraints independent from each other.
+template <typename OpTy>
+struct MoveUpIntoAssumingOpPattern : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const override {
+    return moveUpIntoAssumingOpMatchAndRewrite(op.getOperation(), rewriter);
+  }
+};
+
+// Move elementwise operations into a preceding assuming op. This will
+// eventually allow for more fusion opportunities.
+struct MoveElementwiseOpsUpIntoAssumingOpPattern : public RewritePattern {
+  explicit MoveElementwiseOpsUpIntoAssumingOpPattern(MLIRContext *ctx)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, ctx) {}
+
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    // Apply to all elementwise and broadcasting elementwise operations with no
+    // side effects.
+    if (!op->hasTrait<mlir::OpTrait::Elementwise>() &&
+        !op->hasTrait<hlo::OpTrait::BroadcastingElementwise>()) {
+      return failure();
+    }
+    if (!isMemoryEffectFree(op)) return failure();
+
+    return moveUpIntoAssumingOpMatchAndRewrite(op, rewriter);
+  }
+};
+
+// Move operation into an assuming region if all uses are within its body.
+LogicalResult moveDownIntoAssumingOpMatchAndRewrite(Operation *op,
+                                                    PatternRewriter &rewriter) {
+  auto users = op->getUsers();
+  auto it = users.begin();
+  auto end = users.end();
+  if (it == end) return failure();
+
+  // Find candidate assuming op.
+  auto assumingOp = (it++)->getParentOfType<shape::AssumingOp>();
+  if (!assumingOp || assumingOp->isProperAncestor(op)) return failure();
+
+  // Make sure all uses are within the unique assuming op's body.
+  while (it != end) {
+    auto hopefullySameAssumingOp = (it++)->getParentOfType<shape::AssumingOp>();
+    if (!hopefullySameAssumingOp || hopefullySameAssumingOp != assumingOp) {
+      return failure();
+    }
+  }
+
+  // Move op into the assuming region.
+  OpBuilder::InsertionGuard guard(rewriter);
+  rewriter.setInsertionPointToStart(assumingOp.getBody());
+  Operation *newOp = rewriter.clone(*op);
+  rewriter.replaceOp(op, newOp->getResults());
+  return success();
+}
+
+// Move elementwise operations into succeeding assuming regions. This will
+// eventually allow for more fusion opportunities.
+struct MoveElementwiseOpsDownIntoAssumingOpPattern : public RewritePattern {
+  explicit MoveElementwiseOpsDownIntoAssumingOpPattern(MLIRContext *ctx)
+      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, ctx) {}
+
+  LogicalResult matchAndRewrite(Operation *op,
+                                PatternRewriter &rewriter) const override {
+    // Apply to all elementwise and broadcasting elementwise operations with no
+    // side effects.
+    if (!op->hasTrait<mlir::OpTrait::Elementwise>() &&
+        !op->hasTrait<hlo::OpTrait::BroadcastingElementwise>()) {
+      return failure();
+    }
+    if (!isMemoryEffectFree(op)) return failure();
+
+    return moveDownIntoAssumingOpMatchAndRewrite(op, rewriter);
+  }
+};
+
+/// Move operation out of assuming op. This is only valid for
+/// constraint-independent ops, like `cstr_broadcastable` and `shape_of`. It
+/// will eventually allow to make assuming regions' constraints independent from
+/// each other.
+template <typename OpTy>
+struct MoveUpOutOfAssumingOpPattern : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const override {
+    // Must be inside of an assuming op.
+    auto assumingOp = op->template getParentOfType<shape::AssumingOp>();
+    if (!assumingOp) return failure();
+
+    // Operands must not be defined within the assuming op.
+    Block *body = assumingOp.getBody();
+    auto isAvailable = [&](Value v) {
+      Operation *def = v.getDefiningOp();
+      return def == nullptr || def->getBlock() != body;
+    };
+    if (!llvm::all_of(op->getOperands(), isAvailable)) return failure();
+
+    // Move op before the assuming region.
+    OpBuilder::InsertionGuard guard(rewriter);
+    rewriter.setInsertionPoint(assumingOp);
+    Operation *newOp = rewriter.clone(*op);
+    rewriter.replaceOp(op, newOp->getResults());
+
+    // If the assuming region yields none of the new op's results, these values
+    // are exclusively used in the assuming op's body. In these cases there is
+    // no need for further rewrites.
+    auto isNewOpResult = [newOp](Value v) {
+      return llvm::is_contained(newOp->getResults(), v);
+    };
+    auto yieldOp = cast<shape::AssumingYieldOp>(body->getTerminator());
+    if (llvm::none_of(yieldOp.getOperands(), isNewOpResult)) return success();
+
+    // If the assuming region yields any of the new op's results, these values
+    // can instead bypass the assuming region. There is no need to yield them
+    // explicitly as they are assumed to be independent. The assuming op is
+    // rewritten accordingly.
+    SmallVector<Value, 2> replacementValues;
+    auto newAssumingOp = shape::AssumingOp::create(
+        rewriter, assumingOp.getLoc(), assumingOp.getWitness(),
+        [&](OpBuilder& b, Location) {
+          // Copy body.
+          IRMapping mapping;
+          for (Operation &nested : body->without_terminator()) {
+            b.clone(nested, mapping);
+          }
+
+          // Collect new yield operands.
+          SmallVector<Value, 2> newYieldOperands;
+          for (Value result : yieldOp.getOperands()) {
+            if (isNewOpResult(result)) {
+              replacementValues.push_back(result);
+            } else {
+              newYieldOperands.push_back(mapping.lookupOrDefault(result));
+              replacementValues.push_back(nullptr);
+            }
+          }
+          return newYieldOperands;
+        });
+
+    // Use the assuming op's results for the missing replacement values.
+    auto src = newAssumingOp.getResults().begin();
+    for (auto &dst : replacementValues) {
+      if (dst) continue;
+      dst = *src++;
+    }
+
+    rewriter.replaceOp(assumingOp, replacementValues);
+    return success();
+  }
+};
+
+/// Merge assuming regions if their constraints are independent from each other.
+struct MergeAssumingOpsPattern : public OpRewritePattern<shape::AssumingOp> {
+  using OpRewritePattern<shape::AssumingOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(shape::AssumingOp op,
+                                PatternRewriter &rewriter) const override {
+    // Merge assuming op with directly preceding one if both witnesses are
+    // available.
+    auto precedingOp =
+        llvm::dyn_cast_or_null<shape::AssumingOp>(op->getPrevNode());
+    if (!precedingOp) return failure();
+    if (op.getWitness().getDefiningOp() == precedingOp) return failure();
+
+    // Merge witnesses.
+    OpBuilder::InsertionGuard guard(rewriter);
+    rewriter.setInsertionPoint(precedingOp);
+    Value newWitness = shape::AssumingAllOp::create(
+        rewriter, op.getWitness().getDefiningOp()->getLoc(),
+        ValueRange{precedingOp.getWitness(), op.getWitness()});
+
+    // Merge assuming ops.
+    Block *body_a = precedingOp.getBody();
+    Block *body_b = op.getBody();
+    auto newAssumingOp = shape::AssumingOp::create(
+        rewriter, precedingOp.getLoc(), newWitness,
+        [&](OpBuilder& b, Location) {
+          // Copy preceding op's body.
+          IRMapping mapping;
+          for (auto &nested : body_a->without_terminator()) {
+            b.clone(nested, mapping);
+          }
+
+          // Map result values of preceding assuming op.
+          auto yieldOpA =
+              llvm::dyn_cast<shape::AssumingYieldOp>(body_a->getTerminator());
+          for (auto pair :
+               llvm::zip(precedingOp->getResults(), yieldOpA.getOperands())) {
+            mapping.map(std::get<0>(pair),
+                        mapping.lookupOrDefault(std::get<1>(pair)));
+          }
+
+          // Copy op's body.
+          for (auto &nested : body_b->without_terminator()) {
+            b.clone(nested, mapping);
+          }
+
+          // Collect merged assuming op's results.
+          SmallVector<Value, 4> mappedResults;
+          auto yieldOpB =
+              llvm::dyn_cast<shape::AssumingYieldOp>(body_b->getTerminator());
+          for (Value v : yieldOpA.getOperands()) {
+            mappedResults.push_back(mapping.lookupOrDefault(v));
+          }
+          for (Value v : yieldOpB.getOperands()) {
+            mappedResults.push_back(mapping.lookupOrDefault(v));
+          }
+          return mappedResults;
+        });
+
+    // Replace the two assuming ops with the new corresponding results.
+    ValueRange newResults = newAssumingOp->getResults();
+    size_t splitAt = precedingOp->getNumResults();
+    rewriter.replaceOp(precedingOp, newResults.take_front(splitAt));
+    rewriter.replaceOp(op, newResults.drop_front(splitAt));
+    return success();
+  }
+};
+
+struct EliminateDuplicateCstrBroadcastableOps
+    : public OpRewritePattern<shape::CstrBroadcastableOp> {
+  using OpRewritePattern<shape::CstrBroadcastableOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(shape::CstrBroadcastableOp op,
+                                PatternRewriter &rewriter) const override {
+    // Search for previous occurence of the same constraint.
+    Operation *it = op->getPrevNode();
+    while (it != nullptr) {
+      if (auto candidate = llvm::dyn_cast<shape::CstrBroadcastableOp>(it)) {
+        if (candidate.getShapes() == op.getShapes()) {
+          rewriter.replaceOp(op, candidate.getResult());
+          return success();
+        }
+      }
+      it = it->getPrevNode();
+    }
+
+    return failure();
+  }
+};
+
+void populateMergeAssumingOpsPatterns(MLIRContext *context,
+                                      RewritePatternSet *patterns) {
+  patterns->add<
+      EliminateDuplicateCstrBroadcastableOps,
+      InlineBroadcastedShapeOperandsPattern<shape::CstrBroadcastableOp>,
+      MergeAssumingOpsPattern, MoveElementwiseOpsDownIntoAssumingOpPattern,
+      MoveElementwiseOpsUpIntoAssumingOpPattern,
+      MoveUpIntoAssumingOpPattern<shape::AssumingAllOp>,
+      MoveUpIntoAssumingOpPattern<shape::CstrBroadcastableOp>,
+      MoveUpIntoAssumingOpPattern<shape::ShapeOfOp>,
+      MoveUpOutOfAssumingOpPattern<shape::AssumingAllOp>,
+      MoveUpOutOfAssumingOpPattern<shape::CstrBroadcastableOp>,
+      MoveUpOutOfAssumingOpPattern<shape::ShapeOfOp>, ShapeReificationPattern>(
+      context);
+  mhlo::DynamicBroadcastInDimOp::getCanonicalizationPatterns(*patterns,
+                                                             context);
+  mhlo::DynamicReshapeOp::getCanonicalizationPatterns(*patterns, context);
+  shape::AssumingAllOp::getCanonicalizationPatterns(*patterns, context);
+  shape::AssumingOp::getCanonicalizationPatterns(*patterns, context);
+  shape::BroadcastOp::getCanonicalizationPatterns(*patterns, context);
+  shape::CstrBroadcastableOp::getCanonicalizationPatterns(*patterns, context);
+  tensor::CastOp::getCanonicalizationPatterns(*patterns, context);
+}
+
+struct MergeAssumingOpsPass
+    : public impl::MergeAssumingOpsPassBase<MergeAssumingOpsPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<shape::ShapeDialect, mhlo::MhloDialect>();
+  }
+
+  void runOnOperation() override {
+    MLIRContext *ctx = &getContext();
+    RewritePatternSet patterns(ctx);
+    populateMergeAssumingOpsPatterns(ctx, &patterns);
+    GreedyRewriteConfig config;
+    config.setMaxIterations(GreedyRewriteConfig::kNoLimit);
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns),
+                                     config))) {
+      return signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/parallel_loops_to_sequential.cc

@@ -0,0 +1,59 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include <utility>
+
+#include "mlir/Conversion/SCFToControlFlow/SCFToControlFlow.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/SCF/IR/SCF.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+namespace {
+
+#define GEN_PASS_DEF_PARALLELLOOPSTOSEQUENTIAL
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+struct ParallelLoopsToSequentialPass
+    : public impl::ParallelLoopsToSequentialBase<
+          ParallelLoopsToSequentialPass> {
+  void runOnOperation() override {
+    mlir::RewritePatternSet patterns(&getContext());
+    mlir::populateSCFToControlFlowConversionPatterns(patterns);
+
+    mlir::ConversionTarget target(getContext());
+    target.addIllegalOp<mlir::scf::ParallelOp>();
+    target.addLegalOp<mlir::scf::ForOp, mlir::scf::IfOp>();
+    target.markUnknownOpDynamicallyLegal([](mlir::Operation*) { return true; });
+    if (failed(applyPartialConversion(getOperation(), target,
+                                      std::move(patterns))))
+      signalPassFailure();
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<mlir::OperationPass<mlir::func::FuncOp>>
+CreateParallelLoopsToSequential() {
+  return std::make_unique<ParallelLoopsToSequentialPass>();
+}
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h

@@ -0,0 +1,124 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TRANSFORMS_PASSES_H_
+#define TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TRANSFORMS_PASSES_H_
+
+#include <memory>
+#include <string>
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"  // from @llvm-project
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Pass/Pass.h"  // from @llvm-project
+
+#define GEN_PASS_DECL_TFKERNELTOLLVMPASS
+#define GEN_PASS_DECL_EMBEDTFFRAMEWORKPASS
+#define GEN_PASS_DECL_REWRITETFFRAMEWORKASSERT
+#define GEN_PASS_DECL_FUNCTOJITINVOCATIONPASS
+#define GEN_PASS_DECL_BUFFERREUSEPASS
+#define GEN_PASS_DECL_SHAPETODESCRIPTORSPASS
+#define GEN_PASS_DECL_KERNELGENFINALBUFFERIZEPASS
+#define GEN_PASS_DECL_GPUKERNELTOBLOBPASS
+#define GEN_PASS_DECL_PARALLELLOOPSTOSEQUENTIAL
+#define GEN_PASS_DECL_PROPAGATETFABIKNOWLEDGETOKERNELS
+#define GEN_PASS_DECL_PROPAGATESHAPEKNOWLEDGETOKERNELS
+#define GEN_PASS_DECL_FUSEINNERPARALLELLOOPSPASS
+#define GEN_PASS_DECL_COPYCLEANUPPASS
+#define GEN_PASS_DECL_SHAPESIMPLIFICATIONPASS
+#define GEN_PASS_DECL_MERGEASSUMINGOPSPASS
+#define GEN_PASS_DECL_BROADCASTPROPAGATIONPASS
+#define GEN_PASS_DECL_LEGALIZETENSORRESHAPEPASS
+
+namespace mlir {
+namespace kernel_gen {
+namespace tf_framework {
+
+// Pass to replace some of the Standard ops with TF Framework ops.
+// * adds tf_framework::OpKernelContextType argument to the function
+// * std.alloc becomes tf_framework.alloc_raw
+// * std.dealloc becomes tf_framework.dealloc_raw
+// * std.assert becomes tf_framework.assert
+std::unique_ptr<OperationPass<ModuleOp>> CreateEmbedTFFrameworkPass();
+
+// Pass to convert tf_framework.assert operations to calls to
+// tf_framework.report_error and create the required control flow to abort the
+// function on failed execution.
+std::unique_ptr<OperationPass<ModuleOp>> CreateRewriteTFFrameworkAssert();
+
+}  // namespace tf_framework
+
+namespace transforms {
+
+// Pass to find and annotate candidates for buffer reuse.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateBufferReusePass();
+
+// Pass to rewrite all functions to JIT invocations through the TF
+// framework.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateFuncToJITInvocationPass(
+    llvm::ArrayRef<int64_t> tile_sizes = {},
+    llvm::ArrayRef<int64_t> unroll_factors = {}, bool enable_ftz = false,
+    bool index_64bit = false, bool cpu_codegen = false,
+    bool jit_i64_indexed_for_large_tensors = false);
+
+// Pass for applying LLVM legalization patterns.
+std::unique_ptr<OperationPass<ModuleOp>> CreateTFKernelToLLVMPass(
+    mlir::StringRef blob_annotation = {});
+
+// Pass to tranform shape computations in shape dialect to standard and scf
+// using memref descriptors.
+std::unique_ptr<OperationPass<ModuleOp>> CreateShapeToDescriptorsPass();
+
+// Pass to convert scf::ParallelOp to scf::ForOp.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateParallelLoopsToSequential();
+
+// Pass to annotate GPU Module with its PTX.
+std::unique_ptr<OperationPass<gpu::GPUModuleOp>> CreateGpuKernelToBlobPass(
+    mlir::StringRef blob_annotation = {},
+    ArrayRef<std::string> architectures = {}, bool print_ptx = false,
+    bool print_llvmir = false, bool enable_ftz = false);
+
+// Pass to propagate tensorflow runtime ABI knowledge across kernel boundaries.
+std::unique_ptr<OperationPass<func::FuncOp>>
+CreatePropagateTfAbiKnowledgeToKernels();
+
+// Pass to propagate shape equalities across kernel boundaries.
+std::unique_ptr<OperationPass<func::FuncOp>>
+CreatePropagateShapeKnowledgeToKernels();
+
+/// Greedily maps loops to GPU hardware dimensions.
+std::unique_ptr<mlir::OperationPass<func::FuncOp>> CreateMapParallelLoopsPass();
+
+/// We need to direct fusion to the inner loops. This cannot be done with
+/// a passmanager alone ATM, as nested pass managers require operations to
+/// be closed from above.
+std::unique_ptr<mlir::OperationPass<func::FuncOp>>
+CreateFuseInnerParallelLoopsPass();
+
+// Pass to remove copies which are consumed by a GenericOp.
+std::unique_ptr<OperationPass<func::FuncOp>> CreateCopyCleanupPass();
+
+std::unique_ptr<OperationPass<ModuleOp>> CreateKernelgenFinalBufferizePass();
+
+}  // namespace transforms
+
+#define GEN_PASS_REGISTRATION
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+}  // namespace kernel_gen
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TRANSFORMS_PASSES_H_

tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.td

@@ -0,0 +1,164 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TF_KERNEL_GEN_PASSES
+#define TF_KERNEL_GEN_PASSES
+
+include "mlir/Pass/PassBase.td"
+
+def TFKernelToLLVMPass : Pass<"tf-kernel-to-llvm", "ModuleOp"> {
+  let summary = "Pass for applying LLVM legalization patterns.";
+  let constructor = "transforms::CreateTFKernelToLLVMPass()";
+  let options = [
+      Option<"blob_annotation_", "blob-annotation", "std::string",
+           /*default=*/"\"gpu.binary_blob\"", "Blob attribute name">,
+  ];
+}
+
+def EmbedTFFrameworkPass
+    : Pass<"embed-tf-framework", "ModuleOp"> {
+  let summary = "Pass to embed TF Framework for allocation and assertions,";
+  let constructor = "tf_framework::CreateEmbedTFFrameworkPass()";
+}
+
+def RewriteTFFrameworkAssert : Pass<"rewrite-tf-framework-assert", "ModuleOp"> {
+  let dependentDialects = ["cf::ControlFlowDialect"];
+  let summary = "Pass to rewrite TFAssertOps to CondBranchOp";
+  let constructor = "tf_framework::CreateRewriteTFFrameworkAssert()";
+}
+
+def FuncToJITInvocationPass : Pass<"func-to-jit-invocation", "mlir::func::FuncOp"> {
+  let dependentDialects = [
+      "mlir::kernel_gen::tf_framework::TFFrameworkDialect",
+      "scf::SCFDialect",
+      "shape::ShapeDialect"
+  ];
+  let summary = "Pass to rewrite a function to JIT invocations";
+  let constructor = "transforms::CreateFuncToJITInvocationPass()";
+  let options = [
+      ListOption<"tile_sizes_", "tile-sizes", "int64_t", "Tiling sizes",
+                 "llvm::cl::ZeroOrMore">,
+      ListOption<"unroll_factors_", "unroll-factors", "int64_t",
+                 "Unrolling in each tile dimension", "llvm::cl::ZeroOrMore">,
+      Option<"enable_ftz_", "enable-ftz", "bool", /*default=*/"",
+             "Enable the denormal flush to zero mode when generating code">,
+      Option<"index_64bit_", "index_64bit", "bool", /*default=*/"",
+             "Enable the 64 bit indexing for GPU kernels">,
+      Option<"cpu_codegen_", "cpu-codegen", "bool", /*default=*/"",
+             "CPU codegen (false implies GPU)">,
+      Option<"jit_i64_indexed_for_large_tensors_",
+             "jit_i64_indexed_for_large_tensors", "bool", /*default=*/"false",
+             "Enable JIT compilation of i64-indexed kernels for large input "
+             "tensors.">,
+  ];
+}
+
+def BufferReusePass : Pass<"buffer-reuse", "mlir::func::FuncOp"> {
+  let summary = "Pass to find and annotate candidates for buffer reuse.";
+  let constructor = "transforms::CreateBufferReusePass()";
+}
+
+def ShapeToDescriptorsPass : Pass<"shape-to-descriptors", "ModuleOp"> {
+  let summary = "Pass to transform shape computations to descriptors";
+  let constructor = "transforms::CreateShapeToDescriptorsPass()";
+}
+
+def KernelgenFinalBufferizePass : Pass<"kernelgen-final-bufferize", "ModuleOp"> {
+  let summary = "Pass to transform late operations on values to buffer based "
+                "ones.";
+  let constructor = "transforms::CreateKernelgenFinalBufferizePass()";
+}
+
+def GpuKernelToBlobPass : Pass<"gpu-kernel-to-blob", "gpu::GPUModuleOp"> {
+  let summary = "Pass to annotate GPU Module with its PTX";
+  let options = [
+    Option<"blob_annotation_", "blob-annotation", "std::string",
+           /*default=*/"\"gpu.binary_blob\"", "Blob attribute name">,
+    ListOption<"architectures_", "arch", "std::string", "GPU architectures">,
+    Option<"generate_fatbin_", "generate-fatbin", "bool", /*default=*/"true",
+           "Bundle machine code for the different architectures in one "
+           "fatbin.">,
+    Option<"print_ptx_", "print-ptx", "bool", /*default=*/"false",
+           "Print generated PTX code per target architecture.">,
+    Option<"print_llvmir_", "print-llvmir", "bool", /*default=*/"false",
+           "Print llvm ir when lowering code per target architecture.">,
+  ];
+  let constructor = "transforms::CreateGpuKernelToBlobPass()";
+}
+
+def ParallelLoopsToSequential : Pass<"parallel-loops-to-sequential", "mlir::func::FuncOp"> {
+  let summary = "Pass to convert scf::ParallelOp to scf::ForOp";
+  let constructor = "transforms::CreateParallelLoopsToSequential()";
+}
+
+def PropagateTfAbiKnowledgeToKernels
+    : Pass<"propagate-tf-abi-knowledge-to-kernels", "mlir::func::FuncOp"> {
+  let summary = "Pass to propagate tensorflow ABI knowledge to kernels";
+  let constructor = "transforms::CreatePropagateTfAbiKnowledgeToKernels()";
+}
+
+def PropagateShapeKnowledgeToKernels
+    : Pass<"propagate-shape-knowledge-to-kernels", "mlir::func::FuncOp"> {
+  let summary = "Pass to propagate shape information into kernels";
+  let constructor = "transforms::CreatePropagateShapeKnowledgeToKernels()";
+}
+
+def FuseInnerParallelLoopsPass
+    : Pass<"fuse-inner-parallel-loops", "mlir::func::FuncOp"> {
+  let summary = "Limited pass to forward stores to loads.";
+  let constructor = "transforms::CreateFuseInnerParallelLoopsPass()";
+  let description = [{
+    Directs parallel loop fusion to the inner loops. This cannot be done with
+    a passmanager alone ATM, as nested pass managers require operations to
+    be closed from above.
+  }];
+}
+
+def CopyCleanupPass : Pass<"copy-cleanup", "mlir::func::FuncOp"> {
+  let summary = "Pass to remove copies which are consumed by a GenericOp.";
+  let constructor = "transforms::CreateCopyCleanupPass()";
+  let description = [{
+    We can have GenericOps which have operands that are a copy, but the copy is
+    not used by any other op. In this case, the GenericOp can just use the
+    buffer which is the source of the copy, and we can remove the Alloc and the
+    copy.
+  }];
+}
+
+def ShapeSimplificationPass
+    : Pass<"shape-simplification", "mlir::func::FuncOp"> {
+  let summary = "Simplify shape ops";
+}
+
+def MergeAssumingOpsPass : Pass<"mhlo-merge-assuming-ops", "func::FuncOp"> {
+  let summary = "Prepare moving dynamic broadcasts up over element-wise "
+    "operations and broadcast the operands rather than the result. This will "
+    "eventually allow for larger fusions.";
+}
+
+def BroadcastPropagationPass : Pass<"mhlo-broadcast-propagation", "func::FuncOp"> {
+  let summary = "Move dynamic broadcasts up over element-wise operations and "
+    "broadcast the operands rather than the result. This will eventually allow "
+    "for larger fusions.";
+}
+
+def LegalizeTensorReshapePass
+    : Pass<"legalize-tensor-reshape", "mlir::func::FuncOp"> {
+  let summary = "Legalize tensor.reshape to mhlo.dynamic_reshape";
+}
+
+
+
+#endif // TF_KERNEL_GEN_PASSES

tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewrite_tf_framework_assert.cc

@@ -0,0 +1,115 @@
+/* Copyright 2021 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <iterator>
+#include <memory>
+#include <utility>
+
+#include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewriters.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace tf_framework {
+namespace {
+
+// Converts tf_framework.assert to a conditional branch that reports an error to
+// OpKernelContext and creates a fake memref using NullMemRefOp.
+class TFAssertOpConverter : public OpConversionPattern<TFAssertOp> {
+ public:
+  using OpConversionPattern<TFAssertOp>::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(
+      TFAssertOp op, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+
+    // Split the block to insert CondBr.
+    OpBuilder::InsertPoint ip = rewriter.saveInsertionPoint();
+    Block *split_block = rewriter.splitBlock(
+        rewriter.getInsertionBlock(), std::next(rewriter.getInsertionPoint()));
+
+    auto func = op->getParentOfType<func::FuncOp>();
+    Block *error_reporting_block =
+        rewriter.createBlock(&func.getRegion(), {}, {});
+    ReportErrorOp::create(rewriter, loc, adaptor.getCtx(),
+                          adaptor.getErrorCode(), adaptor.getMsg());
+
+    SmallVector<Value, 2> null_memrefs;
+    for (auto type : func.getFunctionType().getResults()) {
+      null_memrefs.push_back(NullMemRefOp::create(rewriter, loc, type));
+    }
+    func::ReturnOp::create(rewriter, loc, null_memrefs);
+
+    rewriter.restoreInsertionPoint(ip);
+    rewriter.replaceOpWithNewOp<cf::CondBranchOp>(
+        op, adaptor.getArg(), split_block, mlir::ValueRange(),
+        error_reporting_block, mlir::ValueRange());
+    return success();
+  }
+};
+
+#define GEN_PASS_DEF_REWRITETFFRAMEWORKASSERT
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+bool IsNotInsideTfEntryFunction(Operation *op) {
+  auto func = op->getParentOfType<func::FuncOp>();
+  return !func->hasAttrOfType<UnitAttr>(TFFrameworkDialect::kTFEntryAttrName);
+}
+// All contained `tf_framework.assert` operations are rewritten into calls to
+// `tf_framework.report_error` and the required control flow to make
+// execution of the function terminate.
+class RewriteTFFrameworkAssertPass
+    : public impl::RewriteTFFrameworkAssertBase<RewriteTFFrameworkAssertPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<mlir::kernel_gen::tf_framework::TFFrameworkDialect>();
+  }
+
+ public:
+  void runOnOperation() override {
+    ModuleOp m = getOperation();
+
+    // Populate patterns.
+    MLIRContext *context = &getContext();
+    RewritePatternSet patterns(context);
+    patterns.add<TFAssertOpConverter>(context);
+    PopulateEmbedTFFrameworkAssertPattern(&patterns);
+
+    // Set target.
+    ConversionTarget target(getContext());
+    target.addLegalDialect<tf_framework::TFFrameworkDialect, func::FuncDialect,
+                           cf::ControlFlowDialect>();
+    target.addIllegalOp<TFAssertOp>();
+    target.addDynamicallyLegalOp<cf::AssertOp>(IsNotInsideTfEntryFunction);
+
+    if (failed(applyPartialConversion(m, target, std::move(patterns)))) {
+      signalPassFailure();
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<ModuleOp> > CreateRewriteTFFrameworkAssert() {
+  return std::make_unique<RewriteTFFrameworkAssertPass>();
+}
+
+}  // namespace tf_framework
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewriters.h

@@ -0,0 +1,59 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TRANSFORMS_REWRITERS_H_
+#define TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TRANSFORMS_REWRITERS_H_
+
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+
+namespace mlir {
+namespace bufferization {
+class BufferizeTypeConverter;
+}
+class ConversionTarget;
+class LLVMTypeConverter;
+class MLIRContext;
+class RewritePatternSet;
+class TypeConverter;
+
+namespace kernel_gen {
+namespace tf_framework {
+
+/// Collects a set of patterns to convert from the TF Framework dialect to LLVM.
+void PopulateTFFrameworkToLLVMConversionPatterns(LLVMTypeConverter *converter,
+                                                 RewritePatternSet *patterns);
+
+/// Collects a set of patterns to rewrite functions for use with TF framework
+/// and also replace `alloc`, `dealloc` and `assert`.
+void PopulateEmbedTFFrameworkPatterns(RewritePatternSet *patterns);
+void PopulateEmbedTFFrameworkAssertPattern(RewritePatternSet *patterns);
+
+}  // namespace tf_framework
+
+namespace transforms {
+
+/// Collects a set of patterns that bufferize operations from the standard and
+/// other dialects.
+void populateExtraBufferizeDialects(DialectRegistry &registry);
+void populateExtraBufferizePatterns(ConversionTarget &target,
+                                    MLIRContext *context,
+                                    TypeConverter *converter,
+                                    RewritePatternSet *patterns);
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TRANSFORMS_REWRITERS_H_

tensorflow/compiler/mlir/tools/kernel_gen/transforms/same_shape_propagation.cc

@@ -0,0 +1,383 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file contains the analysis and transformation to rewrite kernel
+// functions such that they use a single set of arguments for the strides and
+// sizes of operands with equal shapes.
+
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <tuple>
+#include <utility>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/ADT/EquivalenceClasses.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"  // from @llvm-project
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"  // from @llvm-project
+#include "mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
+#include "mlir/IR/AsmState.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Value.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+
+#define DEBUG_TYPE "kernel-gen-shapes"
+
+namespace {
+
+using mlir::ArrayRef;
+using mlir::SmallVector;
+using mlir::Value;
+
+/// Represents a value or constant. Used to unify operands for operations that
+/// take both ssa values and attributes.
+struct ValueOrConst {
+  explicit ValueOrConst(Value v) : value_or_constant(v), is_constant(false) {}
+  explicit ValueOrConst(int64_t c) : value_or_constant(c), is_constant(true) {}
+
+  Value value() const {
+    assert(!is_constant);
+    return value_or_constant.value;
+  }
+
+  int64_t constant() const {
+    assert(is_constant);
+    return value_or_constant.constant;
+  }
+
+  bool isConstant() const { return is_constant; }
+
+ private:
+  union ValueOrConstStorage {
+    explicit ValueOrConstStorage(Value v) : value(v) {}
+    explicit ValueOrConstStorage(size_t c) : constant(c) {}
+
+    Value value;
+    int64_t constant;
+  } value_or_constant;
+
+  bool is_constant;
+};
+
+llvm::hash_code hash_value(ValueOrConst value) {
+  return value.isConstant() ? static_cast<llvm::hash_code>(value.constant())
+                            : mlir::hash_value(value.value());
+}
+
+bool operator==(ValueOrConst lhs, ValueOrConst rhs) {
+  if (lhs.isConstant()) {
+    return rhs.isConstant() && lhs.constant() == rhs.constant();
+  } else {
+    return !rhs.isConstant() && lhs.value() == rhs.value();
+  }
+}
+
+inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
+                                     const ValueOrConst &value) {
+  if (value.isConstant()) {
+    os << value.constant();
+  } else {
+    Value val = value.value();
+    mlir::AsmState asm_state(
+        val.getParentRegion()->getParentOfType<mlir::func::FuncOp>());
+    val.printAsOperand(os, asm_state);
+  }
+  return os;
+}
+
+/// Represents a shape, as either a single SSA value that represents the entire
+/// shape vector or as a vector of SSA values representing scalars.
+struct ShapeValue {
+  explicit ShapeValue(Value vector)
+      : shape({ValueOrConst{vector}}), is_vector(true) {}
+  explicit ShapeValue(ValueOrConst vector) : shape({vector}), is_vector(true) {
+    assert(!vector.isConstant());
+  }
+  template <typename T>
+  explicit ShapeValue(T values)
+      : shape(values.begin(), values.end()), is_vector(false) {}
+
+  ValueOrConst vector() const {
+    assert(is_vector);
+    return shape.front();
+  }
+
+  ArrayRef<ValueOrConst> scalars() const {
+    assert(!is_vector);
+    return llvm::ArrayRef(shape);
+  }
+
+  bool isVector() const { return is_vector; }
+
+ private:
+  SmallVector<ValueOrConst, 4> shape;
+  bool is_vector;
+};
+
+llvm::hash_code hash_value(ShapeValue shape) {
+  return shape.isVector() ? hash_value(shape.vector())
+                          : hash_value(shape.scalars());
+}
+
+bool operator==(ShapeValue lhs, ShapeValue rhs) {
+  if (lhs.isVector()) {
+    return rhs.isVector() && lhs.vector() == rhs.vector();
+  } else {
+    return !rhs.isVector() && lhs.scalars() == rhs.scalars();
+  }
+}
+
+inline llvm::raw_ostream &operator<<(llvm::raw_ostream &os,
+                                     const ShapeValue &shape) {
+  if (shape.isVector()) {
+    os << shape.vector();
+    return os;
+  }
+  os << "[";
+  bool first = true;
+  for (auto scalar : shape.scalars()) {
+    if (!first) {
+      os << ", ";
+    }
+    first = false;
+    os << scalar;
+  }
+  os << "]";
+  return os;
+}
+
+}  // namespace
+
+namespace llvm {
+
+template <>
+struct DenseMapInfo<ShapeValue> {
+  static unsigned getHashValue(ShapeValue shape) { return hash_value(shape); }
+  static bool isEqual(ShapeValue LHS, ShapeValue RHS) { return LHS == RHS; }
+};
+
+}  // namespace llvm
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+
+namespace {
+
+#define GEN_PASS_DEF_PROPAGATESHAPEKNOWLEDGETOKERNELS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+// A basic shape equality inference. This should be superceeded by a proper
+// inference once available. Until then, we just build this out to the needs of
+// the kernel generator project.
+class ShapeEqualityKnowledge {
+ public:
+  /// Checks all operations for potential shape equality of their respective
+  /// results.
+  void build(func::FuncOp function) {
+    function.walk([&](Operation *op) {
+      if (auto reshape = dyn_cast<memref::ReshapeOp>(op)) {
+        registerAssociation(ShapeValue{(Value)reshape.getShape()},
+                            reshape.getResult());
+        return;
+      }
+      if (auto cast = dyn_cast<memref::ReinterpretCastOp>(op)) {
+        // Only support fully dynamic sizes for now.
+        // TODO(herhut): Fix once the op has canonicalizers that break this.
+        for (unsigned int p = 0, e = cast.getResultRank(); p < e; ++p) {
+          if (!cast.isDynamicSize(p)) {
+            return;
+          }
+        }
+        registerAssociation(ShapeValue{cast.getSizes()}, cast.getResult());
+        return;
+      }
+      if (auto alloc = dyn_cast<memref::AllocOp>(op)) {
+        SmallVector<ValueOrConst, 4> shape;
+        ShapedType type = mlir::cast<ShapedType>(alloc.getResult().getType());
+        fillShapeFromAllocLike(alloc.getDynamicSizes(), type, shape);
+        registerAssociation(ShapeValue{shape}, alloc.getResult());
+        return;
+      }
+      if (auto alloc = dyn_cast<tf_framework::TFAllocOp>(op)) {
+        // Construct a symbol representing the allocated shape.
+        SmallVector<ValueOrConst, 4> shape;
+        ShapedType type = mlir::cast<ShapedType>(alloc.getResult().getType());
+        fillShapeFromAllocLike(alloc.getDynSizes(), type, shape);
+        registerAssociation(ShapeValue{shape}, alloc.getResult());
+        return;
+      }
+    });
+  }
+
+  /// Checks whether `one` and `other` are known to have the same shape and
+  /// strides.
+  bool haveSameShape(Value one, Value other) {
+    return equal_shapes_.isEquivalent(one.getAsOpaquePointer(),
+                                      other.getAsOpaquePointer());
+  }
+
+ private:
+  static void fillShapeFromAllocLike(mlir::OperandRange operands,
+                                     ShapedType type,
+                                     SmallVectorImpl<ValueOrConst> &shape) {
+    assert(type.hasRank());
+    auto dynamic_sizes = operands.begin();
+    for (auto extent : type.getShape()) {
+      shape.push_back(ShapedType::isDynamic(extent)
+                          ? ValueOrConst{*(dynamic_sizes++)}
+                          : ValueOrConst{extent});
+    }
+  }
+
+  /// Registers the value `value` to have the shape represented by `shape`. If
+  /// `shape` has been registered before, place `value` into the same
+  /// equivalence class. Otherwise register `value` as an equivalence class of
+  /// its own.
+  void registerAssociation(ShapeValue shape, Value value) {
+    LLVM_DEBUG({ llvm::dbgs() << "Processing " << value << "\n"; });
+    auto insert_symbolic = symbolic_shapes_.insert({shape, value});
+    if (insert_symbolic.second) {
+      LLVM_DEBUG({ llvm::dbgs() << "New symbolic shape " << shape << "\n"; });
+      equal_shapes_.insert(value.getAsOpaquePointer());
+      // We have seen this symbolic shape for the first time. Try to match it
+      // with a vector or shape we already know and alias classes if possible.
+      // This could be based on shape dialect if we weren't late in the
+      // lowering.
+      tryEvaluateShapeToRoot(shape, value);
+    } else {
+      auto rep = insert_symbolic.first->second;
+      LLVM_DEBUG({ llvm::dbgs() << "Aliasing with rep " << rep << "\n"; });
+      equal_shapes_.unionSets(rep.getAsOpaquePointer(),
+                              value.getAsOpaquePointer());
+    }
+  }
+
+  /// Follows the definition chains of the ShapeValue `shape` to identify cases
+  /// where `shape` is derived from some other value's shape. In such case, the
+  /// equivalence classes of that other value and `value` are unioned.
+  /// This is based on pattern matching and not complete.
+  void tryEvaluateShapeToRoot(ShapeValue shape, Value value) {
+    // Just some pattern matching for common cases here.
+    if (!shape.isVector()) {
+      // Patterns that revolve around scalars.
+      // Check whether the scalars are all dim operations for some other memref.
+      Value candidate;
+      bool all_are_dimops =
+          llvm::all_of(llvm::enumerate(shape.scalars()), [&candidate](auto p) {
+            ValueOrConst val = p.value();
+            if (val.isConstant()) return false;
+            auto dimOp = val.value().getDefiningOp<memref::DimOp>();
+            if (!dimOp) return false;
+            if (!candidate) candidate = dimOp.getSource();
+            auto index = dimOp.getConstantIndex();
+            if (!index.has_value()) return false;
+            return candidate == dimOp.getSource() && p.index() == index.value();
+          });
+      if (all_are_dimops && candidate) {
+        equal_shapes_.unionSets(candidate.getAsOpaquePointer(),
+                                value.getAsOpaquePointer());
+      }
+    }
+  }
+
+  // These are values with identical shapes (or rather their opaque pointers).
+  llvm::EquivalenceClasses<void *> equal_shapes_;
+  // A map from a value that encodes a shape to a value that has this shape.
+  llvm::DenseMap<ShapeValue, Value> symbolic_shapes_;
+};
+
+/// For arguments to kernels that have the same shape, use the stride and
+/// shape information of the left-most argument inside of the kernel function.
+/// That way, llvm can CSE index computations on same-shaped inputs.
+struct PropagateShapeKnowledgeToKernels
+    : public impl::PropagateShapeKnowledgeToKernelsBase<
+          PropagateShapeKnowledgeToKernels> {
+  void runOnOperation() override {
+    ShapeEqualityKnowledge knowledge;
+
+    knowledge.build(getOperation());
+
+    getOperation().walk([&](gpu::LaunchFuncOp launch) {
+      auto module = launch->getParentOfType<ModuleOp>();
+      auto kernel = module.lookupSymbol<LLVM::LLVMFuncOp>(launch.getKernel());
+
+      if (!kernel || kernel.isExternal()) return;
+
+      llvm::SmallVector<std::pair<Value, int>, 4> seen_memrefs;
+      // Position of the kernel argument we are currently at.
+      int kernel_p = 0;
+      for (auto operand : launch.getKernelOperands()) {
+        auto memref = mlir::dyn_cast<MemRefType>(operand.getType());
+        if (!memref) {
+          // Scalar argument, advance kernel position by one.
+          kernel_p++;
+          continue;
+        }
+        for (auto previous : seen_memrefs) {
+          if (!knowledge.haveSameShape(operand, previous.first)) {
+            continue;
+          }
+          auto previous_type = mlir::cast<MemRefType>(previous.first.getType());
+          // We use the first equality found and replace uses of corresponding
+          // size and (potentially) stride information here.
+          auto args_to_replace = memref.getRank();
+          // If both memrefs have identity layouts, we can also reuse the
+          // strides here, as they are the identity strides and hence fully
+          // determinded by the shape.
+          if (previous_type.getLayout().isIdentity() &&
+              memref.getLayout().isIdentity()) {
+            args_to_replace *= 2;
+          }
+          int previous_args_pos = previous.second;
+          auto previous_args = kernel.getArguments()
+                                   .drop_front(previous_args_pos + 3)
+                                   .take_front(args_to_replace);
+          auto current_args = kernel.getArguments()
+                                  .drop_front(kernel_p + 3)
+                                  .take_front(args_to_replace);
+          for (auto pair : llvm::zip(previous_args, current_args)) {
+            mlir::BlockArgument prev, curr;
+            std::tie(prev, curr) = pair;
+            curr.replaceAllUsesWith(prev);
+          }
+          break;
+        }
+        seen_memrefs.push_back({operand, kernel_p});
+        // Advance base, aligned, offset, strides and sizes many arguments.
+        kernel_p += memref.getRank() * 2 + 3;
+      }
+    });
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+CreatePropagateShapeKnowledgeToKernels() {
+  return std::make_unique<PropagateShapeKnowledgeToKernels>();
+}
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/shape_simplification_pass.cc

@@ -0,0 +1,255 @@
+/* Copyright 2021 The OpenXLA Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file contains the patterns to simplify shape ops that were deemed not
+// suitable for shape op canonicalization in MLIR Core.
+
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <utility>
+
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
+#include "mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"  // from @llvm-project
+#include "xla/mlir_hlo/mhlo/IR/hlo_ops.h"
+#include "xla/mlir_hlo/mhlo/transforms/passes.h"
+
+namespace mlir {
+namespace kernel_gen {
+
+#define GEN_PASS_DEF_SHAPESIMPLIFICATIONPASS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+namespace {
+
+using shape::BroadcastOp;
+using shape::ConstShapeOp;
+using shape::ShapeOfOp;
+
+// Try to remove operands from broadcasts that don't contribute to the final
+// result.
+struct BroadcastRemoveSubsumedOperandsPattern
+    : public OpRewritePattern<BroadcastOp> {
+  using OpRewritePattern<BroadcastOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(BroadcastOp op,
+                                PatternRewriter &rewriter) const override {
+    // First collect the static components when joining all shapes. The
+    // resulting vector contains a static dimension if any operand has a static
+    // non-1 dimension in that position. The remaining dimensions are set to
+    // dynamic size.
+    SmallVector<int64_t> knownExtents;
+    SmallVector<SmallVector<int64_t, 4>, 4> operandExtents;
+    for (Value shape : op.getShapes()) {
+      auto &extents = operandExtents.emplace_back();
+      if (failed(shape::getShapeVec(shape, extents))) return failure();
+
+      // Prepend dynamic dims if sizes don't match.
+      if (extents.size() > knownExtents.size()) {
+        knownExtents.insert(knownExtents.begin(),
+                            extents.size() - knownExtents.size(),
+                            ShapedType::kDynamic);
+      }
+
+      for (size_t i = 0, e = extents.size(); i != e; ++i) {
+        int64_t extent = extents[e - i - 1];
+        if (extent != ShapedType::kDynamic && extent != 1) {
+          int64_t &knownExtent = knownExtents[knownExtents.size() - i - 1];
+          // A dynamic dimension is subsumed by a static one, but bail out for
+          // known conflicting shapes.
+          if (knownExtent != extent && knownExtent != ShapedType::kDynamic)
+            return failure();
+          knownExtent = extent;
+        }
+      }
+    }
+
+    // If we've figured out all shapes to be constants we're done.
+    if (!llvm::is_contained(knownExtents, ShapedType::kDynamic)) {
+      rewriter.replaceOpWithNewOp<ConstShapeOp>(
+          op, op->getResultTypes(), rewriter.getIndexTensorAttr(knownExtents));
+      return success();
+    }
+
+    // If only some dimensions are known see if any of the operands can be
+    // removed without affecting the result.
+    SmallVector<Value, 4> filteredOperands;
+    for (auto tuple : llvm::zip(op.getShapes(), operandExtents)) {
+      Value shape = std::get<0>(tuple);
+      auto &extents = std::get<1>(tuple);
+
+      // An operand can't be dead if it's the only operand of the maximum rank.
+      // Removing it would reduce the rank of the output.
+      if (llvm::count_if(operandExtents, [&](ArrayRef<int64_t> op) {
+            return op.size() >= extents.size();
+          }) <= 1) {
+        filteredOperands.push_back(shape);
+        continue;
+      }
+
+      for (size_t i = 0, e = extents.size(); i != e; ++i) {
+        int64_t extent = extents[e - i - 1];
+        // A dimension of an operand can be subsumed if it's
+        //   - a 1 dimension. All other operands will have 1 dims or better.
+        if (extent == 1) continue;
+
+        //   - a dynamic dim but the result is known to be constant.
+        int64_t knownExtent = knownExtents[knownExtents.size() - i - 1];
+        assert(knownExtent != 1);
+        if (knownExtent != ShapedType::kDynamic &&
+            extent == ShapedType::kDynamic)
+          continue;
+
+        //   - a constant non-1 dimension equal to the "known" dim.
+        // In this case we also have to check whether this operand is the only
+        // contributor of that constant.
+        if (knownExtent != ShapedType::kDynamic && extent == knownExtent &&
+            llvm::count_if(operandExtents, [&](ArrayRef<int64_t> operandShape) {
+              return i < operandShape.size() &&
+                     operandShape[operandShape.size() - i - 1] == knownExtent;
+            }) > 1)
+          continue;
+
+        filteredOperands.push_back(shape);
+        break;
+      }
+    }
+    if (filteredOperands.size() != op.getShapes().size()) {
+      rewriter.replaceOpWithNewOp<BroadcastOp>(op, op->getResultTypes(),
+                                               filteredOperands);
+      return success();
+    }
+    return failure();
+  }
+};
+
+// Convert cases like:
+// ```
+//  %1 = shape.shape_of %arg0 : tensor<?x?x?xf64> -> tensor<3xindex>
+//  %2 = shape.shape_of %arg1 : tensor<?x?x1xf64> -> tensor<3xindex>
+//  %3 = shape.broadcast %1, %2 : tensor<3xindex>, tensor<3xindex>
+//                                -> tensor<3xindex>
+//  %result = tensor.extract %3[%c2] : tensor<3xindex>
+// ```
+// to
+//
+// ```
+//  %result = tensor.dim %arg0[%c2] : tensor<?x?x2048xf64>
+// ```
+struct ExtractFromBroadcastedTensorCanonicalizationPattern
+    : public OpRewritePattern<tensor::ExtractOp> {
+  using OpRewritePattern<tensor::ExtractOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::ExtractOp op,
+                                PatternRewriter &rewriter) const override {
+    auto broadcastOp = op.getTensor().getDefiningOp<BroadcastOp>();
+    if (!broadcastOp) return failure();
+
+    // Confirm that there is a constant index. This is required, so we can
+    // confirm the DimOp's input will define the resulting broadcasted shape in
+    // that dimension.
+    auto index =
+        op.getIndices().front().getDefiningOp<arith::ConstantIndexOp>();
+    if (!index) return failure();
+    auto idx = index.value();
+
+    // Iterate through the operands with 3 considerations in this order:
+    // 1. If a static, non-1 dimension is seen, we know this to be the
+    // broadcasted result
+    // 2. If a single dynamic dimension is seen, we know this to be the
+    // broadcasted result (with a possibly 1 or non-1 result)
+    // 3. If no dynamic dimensions and no non-1 static dimensions are seen, we
+    // know the result to be 1
+    //
+    // Iterate through all operands, keeping track of dynamic dimensions and
+    // returning immediately if a non-1 static dimension is seen.
+    ShapeOfOp dynamicShape;
+    int64_t numDynamic = 0;
+    for (auto shape : broadcastOp.getShapes()) {
+      auto shapeOfOp = shape.getDefiningOp<ShapeOfOp>();
+      if (!shapeOfOp) return failure();
+      auto shapedType =
+          mlir::cast<ShapedType>(shapeOfOp->getOperandTypes().front());
+
+      // Abort on the existence of unranked shapes as they require more logic.
+      if (!shapedType.hasRank()) return failure();
+      if (shapedType.getRank() <= idx) continue;
+
+      // Only consider dynamic dimensions after the loop because any non-1
+      // static dimension takes precedence.
+      if (shapedType.isDynamicDim(idx)) {
+        dynamicShape = shapeOfOp;
+        numDynamic++;
+        continue;
+      }
+
+      if (shapedType.getDimSize(idx) == 1) continue;
+
+      // Return as soon as we see a non-1 static dim.
+      rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(
+          op, shapedType.getDimSize(idx));
+      return success();
+    }
+    if (numDynamic > 1) return failure();
+
+    // Replace with the single dynamic dimension or 1.
+    if (dynamicShape) {
+      rewriter.replaceOpWithNewOp<tensor::DimOp>(op, dynamicShape.getArg(),
+                                                 index);
+    } else {
+      rewriter.replaceOpWithNewOp<arith::ConstantIndexOp>(op, 1);
+    }
+    return success();
+  }
+};
+
+struct ShapeSimplificationPass
+    : public impl::ShapeSimplificationPassBase<ShapeSimplificationPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<mlir::arith::ArithDialect>();
+    registry.insert<mhlo::MhloDialect>();
+    registry.insert<mlir::func::FuncDialect>();
+    registry.insert<shape::ShapeDialect>();
+    registry.insert<tensor::TensorDialect>();
+  }
+
+  void runOnOperation() override {
+    MLIRContext *context = &getContext();
+    RewritePatternSet patterns(&getContext());
+
+    for (auto op : context->getRegisteredOperations()) {
+      if (isa<shape::ShapeDialect, mhlo::MhloDialect>(op.getDialect()))
+        op.getCanonicalizationPatterns(patterns, context);
+    }
+
+    patterns.add<BroadcastRemoveSubsumedOperandsPattern,
+                 ExtractFromBroadcastedTensorCanonicalizationPattern>(context);
+
+    auto func = getOperation();
+    if (failed(applyPatternsGreedily(func, std::move(patterns))))
+      return signalPassFailure();
+  }
+};
+
+}  // namespace
+
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/shape_to_descriptors_pass.cc

@@ -0,0 +1,88 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file combines patterns for lowering shape dialect to standard ops,
+// structured control flow and descriptors.
+
+#include <memory>
+#include <utility>
+
+#include "mlir/Conversion/ShapeToStandard/ShapeToStandard.h"  // from @llvm-project
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Math/IR/Math.h"  // from @llvm-project
+#include "mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
+#include "mlir/Dialect/SCF/IR/SCF.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Dialect/Tensor/IR/Tensor.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/PatternMatch.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+namespace {
+
+#define GEN_PASS_DEF_SHAPETODESCRIPTORSPASS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+struct ShapeToDescriptorsPass
+    : public impl::ShapeToDescriptorsPassBase<ShapeToDescriptorsPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<scf::SCFDialect>();
+  }
+
+ public:
+  void runOnOperation() override {
+    MLIRContext &ctx = getContext();
+
+    // Setup target legality.
+    ConversionTarget target(ctx);
+    target.addIllegalDialect<shape::ShapeDialect>();
+    target.addLegalDialect<arith::ArithDialect>();
+    target.addLegalDialect<scf::SCFDialect>();
+    target.addLegalDialect<memref::MemRefDialect>();
+    target.addLegalDialect<func::FuncDialect>();
+    target.addLegalDialect<math::MathDialect>();
+    target.addLegalDialect<tensor::TensorDialect>();
+    // Don't mark the primary Cstr/Assuming ops as illegal, so they can be
+    // lowered at a later time to assertions.
+    target.addLegalOp<shape::AssumingOp, shape::AssumingYieldOp,
+                      shape::AssumingAllOp, shape::CstrRequireOp>();
+
+    // Setup conversion patterns.
+    RewritePatternSet patterns(&getContext());
+    populateShapeRewritePatterns(patterns);
+    populateShapeToStandardConversionPatterns(patterns);
+
+    // Apply conversion.
+    auto module = getOperation();
+    if (failed(applyPartialConversion(module, target, std::move(patterns))))
+      signalPassFailure();
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<ModuleOp> > CreateShapeToDescriptorsPass() {
+  return std::make_unique<ShapeToDescriptorsPass>();
+}
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/tensorflow_abi_knowledge_propagation.cc

@@ -0,0 +1,221 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file contains the analysis and transformation to rewrite kernel
+// functions such that information about alignment, aliasing and zero offsets
+// steming from the tf_framework uses is propagated.
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+#include "llvm/ADT/Bitfields.h"
+#include "llvm/ADT/DenseMap.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"  // from @llvm-project
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"  // from @llvm-project
+#include "mlir/Dialect/MemRef/IR/MemRef.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+namespace {
+
+#define GEN_PASS_DEF_PROPAGATETFABIKNOWLEDGETOKERNELS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+struct PropagateTfAbiKnowledgeToKernelsPass
+    : public impl::PropagateTfAbiKnowledgeToKernelsBase<
+          PropagateTfAbiKnowledgeToKernelsPass> {
+  void runOnOperation() override {
+    func::FuncOp function = getOperation();
+    llvm::SmallVector<Value, 4> worklist;
+    // We currently only handle entry functions and do not propagate across
+    // functions.
+    if (function->getAttrOfType<mlir::UnitAttr>(
+            tf_framework::TFFrameworkDialect::kTFEntryAttrName)) {
+      // For all operands of this function, we know they are aligned. Also, by
+      // construction of kernel generator, we know that there is no offset and
+      // the inner stride is one.
+      // TODO(herhut): Insert asserts in debug mode to check this.
+      for (auto argument : function.getArguments()) {
+        if (mlir::isa<BaseMemRefType>(argument.getType())) {
+          worklist.push_back(argument);
+          allocated_by_tf_runtime.insert(argument);
+          offset_is_zero.insert(argument);
+          inner_stride_is_constant.insert({argument, 1});
+        }
+      }
+    }
+
+    // For locally allocated values, we know they are aligned and have offset
+    // zero. Further, they also do not alias with other memrefs, except in
+    // benign ways. This is by construction and ensured by the reuse analysis.
+    function.walk([&](tf_framework::TFAllocOp op) {
+      Value allocated = op.getResult();
+      worklist.push_back(allocated);
+      no_alias.insert(allocated);
+      allocated_by_tf_runtime.insert(allocated);
+      offset_is_zero.insert(allocated);
+      inner_stride_is_constant.insert({allocated, 1});
+    });
+
+    // Next, take what we have and propagate it through known operations.
+    propagateThroughUses(worklist);
+
+    // Now look at launches and make use of the knowledge we have.
+    function.walk([&](gpu::LaunchFuncOp launch) {
+      auto module = launch->getParentOfType<ModuleOp>();
+      auto kernel = module.lookupSymbol<LLVM::LLVMFuncOp>(launch.getKernel());
+
+      if (!kernel || kernel.isExternal()) return;
+
+      // Count the position of kernel operands independently, as they do not
+      // coincide with laucnh operands as memref parameters get expanded when
+      // lowered to llvm.
+      int kernel_p = 0;
+      OpBuilder b = OpBuilder::atBlockBegin(&kernel.getBody().front());
+      llvm::SmallDenseMap<int64_t, Value> constants;
+      auto loc = kernel.getLoc();
+      for (auto operand : launch.getKernelOperands()) {
+        auto memref = mlir::dyn_cast<MemRefType>(operand.getType());
+        if (!memref) {
+          // Scalar argument, advance kernel position by one.
+          kernel_p++;
+          continue;
+        }
+        if (allocated_by_tf_runtime.contains(operand)) {
+          // This was allocated by the tf runtime, so the two pointers in the
+          // descriptor coincide. Rewrite the kernel accordingly.
+          Value alloc_ptr = kernel.getArgument(kernel_p);
+          Value align_ptr = kernel.getArgument(kernel_p + 1);
+          alloc_ptr.replaceAllUsesWith(align_ptr);
+          kernel.setArgAttr(
+              kernel_p + 1, LLVM::LLVMDialect::getAlignAttrName(),
+              b.getIndexAttr(
+                  tf_framework::TFFrameworkDialect::kAllocationAlignment));
+        }
+        if (offset_is_zero.contains(operand)) {
+          Value offset = kernel.getArgument(kernel_p + 2);
+          Value &zero = constants[0];
+          if (!zero) {
+            zero = LLVM::ConstantOp::create(b, loc, offset.getType(),
+                                            b.getIndexAttr(0));
+          }
+          offset.replaceAllUsesWith(zero);
+        }
+        auto const_stride = inner_stride_is_constant.find(operand);
+        if (const_stride != inner_stride_is_constant.end()) {
+          // The stride is the last argument belonging to this memref.
+          Value inner_stride =
+              kernel.getArgument(kernel_p + 2 + memref.getRank() * 2);
+          Value &stride_val = constants[const_stride->second];
+          if (!stride_val) {
+            stride_val =
+                LLVM::ConstantOp::create(b, loc, inner_stride.getType(),
+                                         b.getIndexAttr(const_stride->second));
+          }
+          inner_stride.replaceAllUsesWith(stride_val);
+        }
+        if (no_alias.contains(operand)) {
+          // TODO(herhut): We also need to check whether any of the other args
+          //     are aliases. This is currently never the case by construction
+          //     but we could use the alias analysis from buffer placement here
+          //     to make sure.
+          // Add the no_alias attribute to the corresponding pointer.
+          kernel.setArgAttr(kernel_p + 1,
+                            LLVM::LLVMDialect::getNoAliasAttrName(),
+                            b.getUnitAttr());
+        }
+        // Advance base, aligned, offset, strides and sizes many arguments.
+        kernel_p += memref.getRank() * 2 + 3;
+      }
+    });
+  }
+
+ private:
+  void propagateThroughUses(SmallVectorImpl<Value> &worklist) {
+    while (!worklist.empty()) {
+      Value candidate = worklist.pop_back_val();
+      for (auto user : candidate.getUsers()) {
+        if (isa<memref::CastOp, memref::ReshapeOp>(user)) {
+          // Reshape and Cast propagate alignment, offset and innermost stride.
+          // TODO(herhut): This should be a trait.
+          Value result = user->getResult(0);
+          if (allocated_by_tf_runtime.contains(candidate)) {
+            allocated_by_tf_runtime.insert(result);
+          }
+          auto const_stride = inner_stride_is_constant.find(candidate);
+          if (const_stride != inner_stride_is_constant.end()) {
+            inner_stride_is_constant.insert({result, const_stride->second});
+          }
+          if (offset_is_zero.contains(candidate)) {
+            offset_is_zero.insert(result);
+          }
+          worklist.push_back(result);
+        }
+        if (auto cast = dyn_cast<memref::ReinterpretCastOp>(user)) {
+          // Check that we have offset 0.
+          Value result = cast.getResult();
+          if (!cast.isDynamicOffset(0) && cast.getStaticOffset(0) == 0) {
+            offset_is_zero.insert(result);
+          }
+          if (allocated_by_tf_runtime.contains(candidate)) {
+            allocated_by_tf_runtime.insert(result);
+          }
+          size_t last_stride = cast.getResultRank() - 1;
+          // TODO(herhut): Remove this once canonicalization handles this.
+          if (cast.isDynamicStride(last_stride)) {
+            auto dyn_stride = cast.getDynamicStride(last_stride)
+                                  .getDefiningOp<arith::ConstantIndexOp>();
+            if (dyn_stride) {
+              inner_stride_is_constant.insert({result, dyn_stride.value()});
+            }
+          } else {
+            inner_stride_is_constant.insert(
+                {result, cast.getStaticStride(last_stride)});
+          }
+          worklist.push_back(result);
+        }
+      }
+    }
+  }
+
+  // Set of values that were allocated by the tf runtime and hence are aligned.
+  llvm::SmallPtrSet<Value, 8> allocated_by_tf_runtime;
+  // Set of values that are known to not have an offset of 0.
+  llvm::SmallPtrSet<Value, 8> offset_is_zero;
+  // Set of values that are known to have a constant stride.
+  llvm::SmallDenseMap<Value, int64_t, 8> inner_stride_is_constant;
+  // Set of values we know do not alias other values.
+  llvm::SmallPtrSet<Value, 8> no_alias;
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<func::FuncOp>>
+CreatePropagateTfAbiKnowledgeToKernels() {
+  return std::make_unique<PropagateTfAbiKnowledgeToKernelsPass>();
+}
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/tf_framework_legalize_to_llvm.cc

@@ -0,0 +1,600 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cassert>
+#include <cstdint>
+#include <functional>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include "mlir/Conversion/LLVMCommon/Pattern.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewriters.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/utils.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace tf_framework {
+namespace {
+
+using transforms::CreateOrFindGlobalStringConstant;
+using transforms::GetGlobalName;
+using transforms::GetOrInsertLLVMFunction;
+
+static constexpr StringRef kCInterfaceAlloc = "_mlir_ciface_tf_alloc";
+static constexpr StringRef kCInterfaceDealloc = "_mlir_ciface_tf_dealloc";
+static constexpr StringRef kCInterfaceReportError =
+    "_mlir_ciface_tf_report_error";
+static constexpr StringRef kCInterfaceJITCompile =
+    "_mlir_ciface_tf_jit_compile";
+static constexpr StringRef kCInterfaceJITExecute =
+    "_mlir_ciface_tf_jit_execute";
+static constexpr StringRef kJITCodeGlobalBaseName = "jit_module_code";
+static constexpr StringRef kErrorMessageGlobalBaseName = "error_message";
+
+/// Base class for patterns converting TF Framework ops to function calls.
+template <typename OpTy>
+class ConvertToLLVMCallOpPattern : public ConvertOpToLLVMPattern<OpTy> {
+ public:
+  using ConvertOpToLLVMPattern<OpTy>::ConvertOpToLLVMPattern;
+
+ protected:
+  virtual StringRef GetFuncName() const = 0;
+  virtual Type GetFuncType() const = 0;
+
+  std::pair<Value, Value> ConvertArrayAttrToStackAllocatedArray(
+      Location loc, Type size_ty, Type element_ty,
+      std::optional<ArrayAttr> attr, ConversionPatternRewriter *rewriter,
+      std::function<Value(Attribute)> create_element) const {
+    Type ptr_ty = LLVM::LLVMPointerType::get(element_ty.getContext());
+
+    // If the attribute is missing or empty, set the element count to 0 and
+    // return NULL.
+    if (!attr.has_value() || attr.value().empty()) {
+      Value zero = LLVM::ConstantOp::create(
+          *rewriter, loc, size_ty, rewriter->getIntegerAttr(size_ty, 0));
+      Value null_ptr = LLVM::ZeroOp::create(*rewriter, loc, ptr_ty);
+      return std::make_pair(zero, null_ptr);
+    }
+
+    // Allocate array to store the elements.
+    auto &array_attr = attr.value();
+    Value array_size = LLVM::ConstantOp::create(
+        *rewriter, loc, size_ty,
+        rewriter->getIntegerAttr(size_ty, array_attr.size()));
+    Value array_ptr = LLVM::AllocaOp::create(*rewriter, loc, ptr_ty, element_ty,
+                                             array_size, /*alignment=*/0);
+    for (const auto &e : llvm::enumerate(array_attr)) {
+      Value index = LLVM::ConstantOp::create(
+          *rewriter, loc, size_ty,
+          rewriter->getIntegerAttr(size_ty, e.index()));
+      Value element_ptr = LLVM::GEPOp::create(*rewriter, loc, ptr_ty,
+                                              element_ty, array_ptr, index);
+      Value element = create_element(e.value());
+      LLVM::StoreOp::create(*rewriter, loc, element, element_ptr);
+    }
+    return std::make_pair(array_size, array_ptr);
+  }
+
+  std::pair<Value, Value> ConvertIntegerArrayAttrToStackAllocatedArray(
+      Location loc, Type size_ty, Type element_ty,
+      std::optional<ArrayAttr> attr,
+      ConversionPatternRewriter *rewriter) const {
+    assert(mlir::isa<IntegerType>(size_ty) && "expect integer size type");
+    assert(mlir::isa<IntegerType>(element_ty) && "expect integer element type");
+    return ConvertArrayAttrToStackAllocatedArray(
+        loc, size_ty, element_ty, attr, rewriter, [&](Attribute attr) {
+          return LLVM::ConstantOp::create(
+              *rewriter, loc, element_ty,
+              rewriter->getIntegerAttr(element_ty,
+                                       mlir::cast<IntegerAttr>(attr).getInt()));
+        });
+  }
+};
+
+class TFAllocOpConverter : public ConvertToLLVMCallOpPattern<TFAllocOp> {
+ public:
+  using ConvertToLLVMCallOpPattern<TFAllocOp>::ConvertToLLVMCallOpPattern;
+
+  LogicalResult matchAndRewrite(
+      TFAllocOp tf_alloc_op, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    mlir::Operation *op = tf_alloc_op.getOperation();
+    Location loc = op->getLoc();
+    MemRefType memref_type = tf_alloc_op.getType();
+
+    // Get memref descriptor sizes.
+    SmallVector<Value, 4> sizes;
+    SmallVector<Value, 4> strides;
+    Value sizeBytes;
+    getMemRefDescriptorSizes(loc, memref_type,
+                             llvm::to_vector<4>(adaptor.getDynSizes()),
+                             rewriter, sizes, strides, sizeBytes);
+    // Get number of elements.
+    Value num_elements =
+        getNumElements(loc, memref_type, adaptor.getDynSizes(), rewriter);
+    // Get element size.
+    Value element_size =
+        getSizeInBytes(loc, memref_type.getElementType(), rewriter);
+
+    // Convert `output_index` or set it to -1 if the attribute is missing.
+    Type llvmInt32Type = IntegerType::get(rewriter.getContext(), 32);
+    Value output_index = LLVM::ConstantOp::create(
+        rewriter, loc, llvmInt32Type,
+        rewriter.getI32IntegerAttr(tf_alloc_op.getOutputIndex().has_value()
+                                       ? tf_alloc_op.getOutputIndex().value()
+                                       : -1));
+
+    // Convert `candidate_input_indices`.
+    auto candidates_count_and_ptr =
+        ConvertIntegerArrayAttrToStackAllocatedArray(
+            loc, rewriter.getI32Type(), rewriter.getI32Type(),
+            tf_alloc_op.getInputIndices(), &rewriter);
+
+    // Insert function call.
+    FlatSymbolRefAttr tf_func_ref =
+        GetOrInsertLLVMFunction(GetFuncName(), GetFuncType(), op, &rewriter);
+    Value allocated_byte_ptr =
+        LLVM::CallOp::create(
+            rewriter, loc, getVoidPtrType(), tf_func_ref,
+            llvm::ArrayRef({adaptor.getCtx(), num_elements, element_size,
+                            output_index, candidates_count_and_ptr.first,
+                            candidates_count_and_ptr.second}))
+            .getResult();
+
+    MemRefDescriptor memRefDescriptor = CreateMemRefDescriptor(
+        loc, rewriter, memref_type, allocated_byte_ptr, sizes);
+
+    // Return the final value of the descriptor.
+    rewriter.replaceOp(op, {memRefDescriptor});
+    return success();
+  }
+
+ protected:
+  StringRef GetFuncName() const override { return kCInterfaceAlloc; }
+
+  Type GetFuncType() const override {
+    Type llvm_i32_type = IntegerType::get(getDialect().getContext(), 32);
+    Type llvm_ptr_type = LLVM::LLVMPointerType::get(getDialect().getContext());
+    return LLVM::LLVMFunctionType::get(
+        llvm_ptr_type,
+        llvm::ArrayRef({/*void* op_kernel_ctx*/ llvm_ptr_type,
+                        /*size_t num_elements*/ getIndexType(),
+                        /*size_t element_size*/ getIndexType(),
+                        /*int32_t output_index*/ llvm_i32_type,
+                        /*int32_t num_candidates*/ llvm_i32_type,
+                        /*int32_t* candidate_input_indices*/ llvm_ptr_type}));
+  }
+
+ private:
+  // TODO(pifon): Remove strides computation.
+  MemRefDescriptor CreateMemRefDescriptor(Location loc,
+                                          ConversionPatternRewriter &rewriter,
+                                          MemRefType memref_type,
+                                          Value allocated_byte_ptr,
+                                          ArrayRef<Value> sizes) const {
+    auto memref_desc = MemRefDescriptor::poison(
+        rewriter, loc, typeConverter->convertType(memref_type));
+
+    // TF AllocateRaw returns aligned pointer => AllocatedPtr == AlignedPtr.
+    memref_desc.setAllocatedPtr(rewriter, loc, allocated_byte_ptr);
+    memref_desc.setAlignedPtr(rewriter, loc, allocated_byte_ptr);
+    memref_desc.setConstantOffset(rewriter, loc, 0);
+
+    if (memref_type.getRank() == 0) {
+      return memref_desc;
+    }
+
+    // Compute strides and populate descriptor `size` and `stride` fields.
+    Value stride_carried =
+        createIndexAttrConstant(rewriter, loc, getIndexType(), 1);
+    for (int pos = sizes.size() - 1; pos >= 0; --pos) {
+      Value size = sizes[pos];
+      memref_desc.setSize(rewriter, loc, pos, size);
+      memref_desc.setStride(rewriter, loc, pos, stride_carried);
+      // Update stride
+      if (pos > 0) {
+        stride_carried =
+            LLVM::MulOp::create(rewriter, loc, stride_carried, size);
+      }
+    }
+    return memref_desc;
+  }
+};
+
+class TFDeallocOpConverter : public ConvertToLLVMCallOpPattern<TFDeallocOp> {
+ public:
+  using ConvertToLLVMCallOpPattern<TFDeallocOp>::ConvertToLLVMCallOpPattern;
+
+  LogicalResult matchAndRewrite(
+      TFDeallocOp op, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    // TODO(herhut) Support unranked memrefs.
+    if (!mlir::isa<MemRefType>(op.getMemref().getType())) return failure();
+    MemRefDescriptor memref(adaptor.getMemref());
+
+    Value allocated_bytes_ptr = memref.allocatedPtr(rewriter, op.getLoc());
+
+    // Insert function call.
+    FlatSymbolRefAttr tf_func_ref =
+        GetOrInsertLLVMFunction(GetFuncName(), GetFuncType(), op, &rewriter);
+    rewriter.replaceOpWithNewOp<LLVM::CallOp>(
+        op, mlir::TypeRange(), tf_func_ref,
+        llvm::ArrayRef({adaptor.getCtx(), allocated_bytes_ptr}));
+    return success();
+  }
+
+ protected:
+  StringRef GetFuncName() const override { return kCInterfaceDealloc; }
+  Type GetFuncType() const override {
+    return LLVM::LLVMFunctionType::get(getVoidType(),
+                                       {getVoidPtrType(), getVoidPtrType()});
+  }
+};
+
+class JITCompileFromStrOpConverter
+    : public ConvertToLLVMCallOpPattern<JITCompileFromStrOp> {
+  using ConvertToLLVMCallOpPattern<
+      JITCompileFromStrOp>::ConvertToLLVMCallOpPattern;
+
+  LogicalResult matchAndRewrite(
+      JITCompileFromStrOp op, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    if (adaptor.getCtx() == nullptr) return failure();
+    auto loc = op.getLoc();
+    std::string zero_terminated_code = op.getCode().str() + '\00';
+    Value jit_module_code = CreateOrFindGlobalStringConstant(
+        loc, GetGlobalName(kJITCodeGlobalBaseName, zero_terminated_code),
+        zero_terminated_code, &rewriter);
+    std::pair<Value, Value> tile_sizes =
+        ConvertIntegerArrayAttrToStackAllocatedArray(
+            loc, rewriter.getI64Type(), rewriter.getI64Type(),
+            op.getTileSizes(), &rewriter);
+    std::pair<Value, Value> unroll_factors =
+        ConvertIntegerArrayAttrToStackAllocatedArray(
+            loc, rewriter.getI64Type(), rewriter.getI64Type(),
+            op.getUnrollFactors(), &rewriter);
+    Value enable_ftz = LLVM::ConstantOp::create(
+        rewriter, loc, rewriter.getI1Type(), op.getEnableFtzAttr());
+    Value index_64bit = LLVM::ConstantOp::create(
+        rewriter, loc, rewriter.getI1Type(), op.getIndex64BitAttr());
+    Value cpu_codegen = LLVM::ConstantOp::create(
+        rewriter, loc, rewriter.getI1Type(), op.getCpuCodegenAttr());
+    FlatSymbolRefAttr tf_func_ref =
+        GetOrInsertLLVMFunction(GetFuncName(), GetFuncType(), op, &rewriter);
+    rewriter.replaceOpWithNewOp<LLVM::CallOp>(
+        op, getVoidPtrType(), tf_func_ref,
+        llvm::ArrayRef({adaptor.getCtx(), jit_module_code, tile_sizes.first,
+                        tile_sizes.second, unroll_factors.first,
+                        unroll_factors.second, enable_ftz, index_64bit,
+                        cpu_codegen}));
+    return success();
+  }
+
+ protected:
+  StringRef GetFuncName() const override { return kCInterfaceJITCompile; }
+
+  Type GetFuncType() const override {
+    auto ptr_ty = LLVM::LLVMPointerType::get(getContext());
+    auto i64_ty = IntegerType::get(getContext(), 64);
+    auto i1_ty = IntegerType::get(getContext(), 1);
+    return LLVM::LLVMFunctionType::get(
+        getVoidPtrType(), {/*void* op_kernel_ctx*/ getVoidPtrType(),
+                           /*char* code*/ ptr_ty,
+                           /*int64_t num_tile_sizes*/ i64_ty,
+                           /*int64_t* tile_sizes_ptr*/ ptr_ty,
+                           /*int64_t num_unroll_factors*/ i64_ty,
+                           /*int64_t* unroll_factors_ptr*/ ptr_ty,
+                           /*bool enable_ftz*/ i1_ty,
+                           /*bool index_64bit*/ i1_ty,
+                           /*bool cpu_codegen*/ i1_ty});
+  }
+};
+
+class JITExecuteOpConverter : public ConvertToLLVMCallOpPattern<JITExecuteOp> {
+ public:
+  using ConvertToLLVMCallOpPattern<JITExecuteOp>::ConvertToLLVMCallOpPattern;
+
+  LogicalResult matchAndRewrite(
+      JITExecuteOp op, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    // The TF context must be known for a successful lowering.
+    if (adaptor.getCtx() == nullptr || op.getInputs().empty()) {
+      return failure();
+    }
+
+    // Allocate result on stack.
+    auto loc = op.getLoc();
+    Type result_ty =
+        getTypeConverter()->convertType(op->getResultTypes().front());
+    Type ptr_ty = LLVM::LLVMPointerType::get(getContext());
+    Type i64_ty = rewriter.getI64Type();
+    Value one = LLVM::ConstantOp::create(rewriter, loc, i64_ty,
+                                         rewriter.getI64IntegerAttr(1));
+    auto result_ptr =
+        LLVM::AllocaOp::create(rewriter, loc, ptr_ty, result_ty, one);
+
+    // Pass the buffer arguments as a stack-allocated array.
+    Type args_elem_ty = adaptor.getInputs().front().getType();
+    Value num_args = LLVM::ConstantOp::create(
+        rewriter, loc, i64_ty,
+        rewriter.getI64IntegerAttr(
+            static_cast<int64_t>(adaptor.getInputs().size())));
+    Value args_ptr =
+        LLVM::AllocaOp::create(rewriter, loc, ptr_ty, args_elem_ty, num_args,
+                               /*alignment=*/0);
+    for (const auto &it : llvm::enumerate(adaptor.getInputs())) {
+      Value index = LLVM::ConstantOp::create(
+          rewriter, loc, i64_ty, rewriter.getI64IntegerAttr(it.index()));
+      Value element_ptr = LLVM::GEPOp::create(rewriter, loc, ptr_ty,
+                                              args_elem_ty, args_ptr, index);
+      LLVM::StoreOp::create(rewriter, loc, it.value(), element_ptr);
+    }
+
+    // Materialize runtime call.
+    FlatSymbolRefAttr tf_func_ref =
+        GetOrInsertLLVMFunction(GetFuncName(), GetFuncType(), op, &rewriter);
+    LLVM::CallOp::create(rewriter, loc, mlir::TypeRange(), tf_func_ref,
+                         ValueRange{adaptor.getCtx(), adaptor.getCallable(),
+                                    result_ptr, num_args, args_ptr});
+
+    // Copy result (including the descriptor) to a stack-allocated buffer and
+    // free the old descriptor.
+    llvm::SmallVector<Value, 1> final_result = {
+        LLVM::LoadOp::create(rewriter, loc, result_ty, result_ptr)};
+    if (failed(copyUnrankedDescriptors(rewriter, loc, op->getResultTypes(),
+                                       final_result,
+                                       /*toDynamic=*/false))) {
+      return failure();
+    }
+
+    rewriter.replaceOp(op, final_result.front());
+    return success();
+  }
+
+ protected:
+  StringRef GetFuncName() const override { return kCInterfaceJITExecute; }
+
+  Type GetFuncType() const override {
+    auto i64_ty = IntegerType::get(getContext(), 64);
+    auto ptr_ty = LLVM::LLVMPointerType::get(getContext());
+    return LLVM::LLVMFunctionType::get(getVoidType(),
+                                       {/*void* op_kernel_ctx*/ ptr_ty,
+                                        /*void* callable*/ ptr_ty,
+                                        /*void* result*/ ptr_ty,
+                                        /*int64_t num_args*/ i64_ty,
+                                        /*void* args_ptr*/ ptr_ty});
+  }
+};
+
+class ReportErrorOpConverter
+    : public ConvertToLLVMCallOpPattern<ReportErrorOp> {
+ public:
+  using ConvertToLLVMCallOpPattern<ReportErrorOp>::ConvertToLLVMCallOpPattern;
+
+  LogicalResult matchAndRewrite(
+      ReportErrorOp op, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    auto module = op->getParentOfType<ModuleOp>();
+    Value message_constant =
+        GenerateErrorMessageConstant(loc, module, adaptor.getMsg(), rewriter);
+
+    // Insert function call.
+    FlatSymbolRefAttr tf_func_ref =
+        GetOrInsertLLVMFunction(GetFuncName(), GetFuncType(), op, &rewriter);
+    Value error_code = LLVM::ConstantOp::create(
+        rewriter, loc, typeConverter->convertType(rewriter.getI32Type()),
+        adaptor.getErrorCodeAttr());
+    rewriter.replaceOpWithNewOp<LLVM::CallOp>(
+        op, mlir::TypeRange(), tf_func_ref,
+        llvm::ArrayRef({adaptor.getCtx(), error_code, message_constant}));
+    return success();
+  }
+
+ protected:
+  StringRef GetFuncName() const override { return kCInterfaceReportError; }
+  Type GetFuncType() const override {
+    MLIRContext *ctx = &getTypeConverter()->getContext();
+    auto ptr_type = LLVM::LLVMPointerType::get(ctx);
+    auto i32_type = IntegerType::get(ctx, 32);
+    return LLVM::LLVMFunctionType::get(getVoidType(),
+                                       {getVoidPtrType(), i32_type, ptr_type});
+  }
+
+ private:
+  // Generates an LLVM IR dialect global that contains the name of the given
+  // kernel function as a C string, and returns a pointer to its beginning.
+  Value GenerateErrorMessageConstant(Location loc, Operation *module,
+                                     StringRef message,
+                                     OpBuilder &builder) const {
+    std::string err_str;
+    llvm::raw_string_ostream err_stream(err_str);
+    err_stream << message;
+    if (!mlir::isa<UnknownLoc>(loc)) {
+      err_stream << " at ";
+      loc.print(err_stream);
+    }
+    err_stream << '\00';
+    StringRef generated_error(err_stream.str());
+    return CreateOrFindGlobalStringConstant(
+        loc, GetGlobalName(kErrorMessageGlobalBaseName, generated_error),
+        generated_error, &builder);
+  }
+};
+
+class NullContextOpConverter : public ConvertOpToLLVMPattern<NullContextOp> {
+ public:
+  using ConvertOpToLLVMPattern<NullContextOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult matchAndRewrite(
+      NullContextOp op, OpAdaptor /*adaptor*/,
+      ConversionPatternRewriter &rewriter) const override {
+    rewriter.replaceOpWithNewOp<LLVM::ZeroOp>(op, getVoidPtrType());
+    return success();
+  }
+};
+
+class NullMemRefOpConverter : public ConvertOpToLLVMPattern<NullMemRefOp> {
+ public:
+  using ConvertOpToLLVMPattern<NullMemRefOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult matchAndRewrite(
+      NullMemRefOp null_memref_op, OpAdaptor /*adaptor*/,
+      ConversionPatternRewriter &rewriter) const override {
+    Location loc = null_memref_op->getLoc();
+    const LLVMTypeConverter &type_converter = *getTypeConverter();
+    MLIRContext *ctx = null_memref_op.getContext();
+    mlir::Operation *op = null_memref_op.getOperation();
+
+    auto shaped_result_type =
+        mlir::cast<BaseMemRefType>(null_memref_op.getType());
+    auto mem_space = mlir::dyn_cast_or_null<IntegerAttr>(
+        shaped_result_type.getMemorySpace());
+    unsigned address_space =
+        static_cast<unsigned>(mem_space ? mem_space.getInt() : 0);
+    LLVM::LLVMPointerType llvm_ptr_type =
+        LLVM::LLVMPointerType::get(ctx, address_space);
+
+    Value zero = createIndexAttrConstant(rewriter, loc, getIndexType(), 0);
+    if (auto result_type =
+            mlir::dyn_cast<MemRefType>(null_memref_op.getType())) {
+      // Set all dynamic sizes to 1 and compute fake strides.
+      SmallVector<Value, 4> dyn_sizes(
+          result_type.getNumDynamicDims(),
+          createIndexAttrConstant(rewriter, loc, getIndexType(), 1));
+      SmallVector<Value, 4> sizes, strides;
+      Value sizeBytes;
+      getMemRefDescriptorSizes(loc, result_type, dyn_sizes, rewriter, sizes,
+                               strides, sizeBytes);
+
+      // Prepare packed args [allocatedPtr, alignedPtr, offset, sizes, strides]
+      // to create a memref descriptor.
+      Value null = LLVM::ZeroOp::create(rewriter, loc, llvm_ptr_type);
+      SmallVector<Value, 12> packed_values{null, null, zero};
+      packed_values.append(sizes);
+      packed_values.append(strides);
+
+      rewriter.replaceOp(
+          op, MemRefDescriptor::pack(rewriter, loc, type_converter, result_type,
+                                     packed_values));
+      return success();
+    }
+
+    auto result_type = mlir::cast<UnrankedMemRefType>(null_memref_op.getType());
+    Type llvm_result_type = type_converter.convertType(result_type);
+
+    auto desc =
+        UnrankedMemRefDescriptor::poison(rewriter, loc, llvm_result_type);
+    desc.setRank(rewriter, loc, zero);
+
+    // Extract address space and element type.
+    auto targetType =
+        mlir::cast<UnrankedMemRefType>(null_memref_op.getResult().getType());
+    unsigned addressSpace =
+        *getTypeConverter()->getMemRefAddressSpace(targetType);
+
+    // Due to the current way of handling unranked memref results escaping, we
+    // have to actually construct a ranked underlying descriptor instead of just
+    // setting its pointer to NULL.
+    Value alloca_size = UnrankedMemRefDescriptor::computeSize(
+        rewriter, loc, *getTypeConverter(), desc, addressSpace);
+    Value underlying_desc_ptr =
+        LLVM::AllocaOp::create(rewriter, loc, getVoidPtrType(),
+                               IntegerType::get(getContext(), 8), alloca_size);
+
+    // Populate underlying ranked descriptor.
+    Value null = LLVM::ZeroOp::create(rewriter, loc, llvm_ptr_type);
+    UnrankedMemRefDescriptor::setAllocatedPtr(
+        rewriter, loc, underlying_desc_ptr, llvm_ptr_type, null);
+    UnrankedMemRefDescriptor::setAlignedPtr(rewriter, loc, *getTypeConverter(),
+                                            underlying_desc_ptr, llvm_ptr_type,
+                                            null);
+    UnrankedMemRefDescriptor::setOffset(rewriter, loc, *getTypeConverter(),
+                                        underlying_desc_ptr, llvm_ptr_type,
+                                        zero);
+
+    desc.setMemRefDescPtr(rewriter, loc, underlying_desc_ptr);
+    rewriter.replaceOp(op, {desc});
+    return success();
+  }
+};
+
+class IsValidMemRefOpConverter
+    : public ConvertOpToLLVMPattern<IsValidMemRefOp> {
+ public:
+  using ConvertOpToLLVMPattern<IsValidMemRefOp>::ConvertOpToLLVMPattern;
+
+  LogicalResult matchAndRewrite(
+      IsValidMemRefOp op, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override {
+    Location loc = op.getLoc();
+    MemRefDescriptor desc(adaptor.getArg());
+
+    // Compare every size in the descriptor to 0 to check num_elements == 0.
+    int64_t rank = mlir::cast<MemRefType>(op.getArg().getType()).getRank();
+    Value is_empty_shape = LLVM::ConstantOp::create(
+        rewriter, loc, rewriter.getI1Type(), rewriter.getBoolAttr(false));
+    Value zero = createIndexAttrConstant(rewriter, loc, getIndexType(), 0);
+    for (int i = 0; i < rank; ++i) {
+      Value size = desc.size(rewriter, loc, i);
+      Value is_zero_size =
+          LLVM::ICmpOp::create(rewriter, loc, rewriter.getI1Type(),
+                               LLVM::ICmpPredicate::eq, size, zero);
+      is_empty_shape =
+          LLVM::OrOp::create(rewriter, loc, is_empty_shape, is_zero_size);
+    }
+
+    Value ptr = desc.allocatedPtr(rewriter, loc);
+    Value null = LLVM::ZeroOp::create(rewriter, loc, getVoidPtrType());
+    Value is_not_nullptr =
+        LLVM::ICmpOp::create(rewriter, loc, rewriter.getI1Type(),
+                             LLVM::ICmpPredicate::ne, ptr, null);
+
+    // Valid memref = ptr != NULL || num_elements == 0;
+    rewriter.replaceOpWithNewOp<LLVM::OrOp>(op, is_not_nullptr, is_empty_shape);
+    return success();
+  }
+};
+
+}  // namespace
+
+void PopulateTFFrameworkToLLVMConversionPatterns(LLVMTypeConverter *converter,
+                                                 RewritePatternSet *patterns) {
+  // clang-format off
+  patterns->add<
+      IsValidMemRefOpConverter,
+      JITCompileFromStrOpConverter,
+      JITExecuteOpConverter,
+      NullContextOpConverter,
+      NullMemRefOpConverter,
+      ReportErrorOpConverter,
+      TFAllocOpConverter,
+      TFDeallocOpConverter>(*converter);
+  // clang-format on
+}
+
+}  // namespace tf_framework
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/tf_kernel_to_llvm_pass.cc

@@ -0,0 +1,316 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <cassert>
+#include <memory>
+#include <utility>
+
+#include "llvm/ADT/STLExtras.h"
+#include "mlir/Conversion/ArithToLLVM/ArithToLLVM.h"  // from @llvm-project
+#include "mlir/Conversion/ComplexToLLVM/ComplexToLLVM.h"  // from @llvm-project
+#include "mlir/Conversion/ControlFlowToLLVM/ControlFlowToLLVM.h"  // from @llvm-project
+#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h"  // from @llvm-project
+#include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVMPass.h"  // from @llvm-project
+#include "mlir/Conversion/GPUCommon/GPUCommonPass.h"  // from @llvm-project
+#include "mlir/Conversion/LLVMCommon/Pattern.h"  // from @llvm-project
+#include "mlir/Conversion/MathToLLVM/MathToLLVM.h"  // from @llvm-project
+#include "mlir/Conversion/MathToLibm/MathToLibm.h"  // from @llvm-project
+#include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h"  // from @llvm-project
+#include "mlir/Conversion/ReconcileUnrealizedCasts/ReconcileUnrealizedCasts.h"  // from @llvm-project
+#include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"  // from @llvm-project
+#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
+#include "mlir/Dialect/Arith/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Dialect/Complex/IR/Complex.h"  // from @llvm-project
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/GPU/IR/GPUDialect.h"  // from @llvm-project
+#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"  // from @llvm-project
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"  // from @llvm-project
+#include "mlir/Dialect/LLVMIR/LLVMTypes.h"  // from @llvm-project
+#include "mlir/Dialect/Math/IR/Math.h"  // from @llvm-project
+#include "mlir/Dialect/MemRef/Transforms/Passes.h"  // from @llvm-project
+#include "mlir/Dialect/MemRef/Transforms/Transforms.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Transforms/DialectConversion.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tools/kernel_gen/ir/tf_framework_ops.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h"
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/rewriters.h"
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+namespace {
+
+constexpr StringRef kTfWrapperLibaryLaunchHelperName =
+    "_mlir_ciface_tf_launch_kernel";
+
+#define GEN_PASS_DEF_TFKERNELTOLLVMPASS
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc"
+
+/// A rewrite patter to convert gpu.launch_func operations into a runtime call
+/// for the TensorFlow runtime.
+class ConvertLaunchFuncOpToTfRuntimeCallPattern
+    : public ConvertOpToLLVMPattern<gpu::LaunchFuncOp> {
+ public:
+  ConvertLaunchFuncOpToTfRuntimeCallPattern(LLVMTypeConverter &type_converter,
+                                            StringRef gpu_binary_annotation)
+      : ConvertOpToLLVMPattern<gpu::LaunchFuncOp>(type_converter),
+        gpu_binary_annotation_(gpu_binary_annotation) {}
+
+ private:
+  Value generateParamsArray(gpu::LaunchFuncOp launch_op, OpAdaptor adaptor,
+                            OpBuilder &builder) const;
+
+  LogicalResult matchAndRewrite(
+      gpu::LaunchFuncOp launch_op, OpAdaptor adaptor,
+      ConversionPatternRewriter &rewriter) const override;
+
+  MLIRContext *context_ = &this->getTypeConverter()->getContext();
+
+  Type llvm_void_type_ = LLVM::LLVMVoidType::get(context_);
+  Type llvm_pointer_type_ = LLVM::LLVMPointerType::get(context_);
+  Type llvm_int8_type_ = IntegerType::get(context_, 8);
+  Type llvm_int32_type_ = IntegerType::get(context_, 32);
+  Type llvm_int64_type_ = IntegerType::get(context_, 64);
+  Type llvm_intptr_type_ = IntegerType::get(
+      context_, this->getTypeConverter()->getPointerBitwidth(0));
+
+  llvm::SmallString<32> gpu_binary_annotation_;
+};
+
+// Creates a struct containing all kernel parameters on the stack and returns
+// an array of type-erased pointers to the fields of the struct. The array can
+// then be passed to the CUDA / ROCm (HIP) kernel launch calls.
+// The generated code is essentially as follows:
+//
+// %struct = alloca(sizeof(struct { Parameters... }))
+// %array = alloca(NumParameters * sizeof(void *))
+// for (i : [0, NumParameters))
+//   %fieldPtr = llvm.getelementptr %struct[0, i]
+//   llvm.store parameters[i], %fieldPtr
+//   %elementPtr = llvm.getelementptr %array[i]
+//   llvm.store %fieldPtr, %elementPtr
+// return %array
+Value ConvertLaunchFuncOpToTfRuntimeCallPattern::generateParamsArray(
+    gpu::LaunchFuncOp launch_op, OpAdaptor adaptor, OpBuilder &builder) const {
+  auto loc = launch_op.getLoc();
+  auto num_kernel_operands = launch_op.getNumKernelOperands();
+  auto arguments = getTypeConverter()->promoteOperands(
+      loc, launch_op.getOperands().take_back(num_kernel_operands),
+      adaptor.getKernelOperands().take_back(num_kernel_operands), builder);
+  auto num_arguments = arguments.size();
+  SmallVector<Type, 4> argument_types;
+  argument_types.reserve(num_arguments);
+  for (auto argument : arguments) argument_types.push_back(argument.getType());
+  auto struct_type = LLVM::LLVMStructType::getNewIdentified(
+      context_, StringRef(), argument_types);
+  auto one = LLVM::ConstantOp::create(builder, loc, llvm_int32_type_,
+                                      builder.getI32IntegerAttr(1));
+  auto struct_ptr = LLVM::AllocaOp::create(builder, loc, llvm_pointer_type_,
+                                           struct_type, one, /*alignment=*/0);
+  auto array_size = LLVM::ConstantOp::create(
+      builder, loc, llvm_int32_type_, builder.getI32IntegerAttr(num_arguments));
+  auto array_ptr =
+      LLVM::AllocaOp::create(builder, loc, llvm_pointer_type_,
+                             llvm_pointer_type_, array_size, /*alignment=*/0);
+  auto zero = LLVM::ConstantOp::create(builder, loc, llvm_int32_type_,
+                                       builder.getI32IntegerAttr(0));
+  for (auto en : llvm::enumerate(arguments)) {
+    auto index = LLVM::ConstantOp::create(
+        builder, loc, llvm_int32_type_, builder.getI32IntegerAttr(en.index()));
+    auto field_ptr = LLVM::GEPOp::create(
+        builder, loc, llvm_pointer_type_, struct_type, struct_ptr,
+        ArrayRef<Value>{zero, index.getResult()});
+    LLVM::StoreOp::create(builder, loc, en.value(), field_ptr);
+    auto element_ptr =
+        LLVM::GEPOp::create(builder, loc, llvm_pointer_type_,
+                            llvm_pointer_type_, array_ptr, index.getResult());
+    LLVM::StoreOp::create(builder, loc, field_ptr, element_ptr);
+  }
+  return array_ptr;
+}
+
+// Emits LLVM IR to launch a kernel function. Expects the module that contains
+// the compiled kernel function as a cubin in the 'nvvm.cubin' attribute, or a
+// hsaco in the 'rocdl.hsaco' attribute of the kernel function in the IR.
+//
+// %0 = call %binarygetter
+// %1 = <pointer to kernel function name>
+// %2 = <see generateParamsArray>
+// call %tfLaunchKernel(%ctx, %0, %1, <launch_op operands 0..5>, %2)
+LogicalResult ConvertLaunchFuncOpToTfRuntimeCallPattern::matchAndRewrite(
+    gpu::LaunchFuncOp launch_op, OpAdaptor adaptor,
+    ConversionPatternRewriter &rewriter) const {
+  if (!launch_op.getAsyncDependencies().empty() || launch_op.getAsyncToken()) {
+    return rewriter.notifyMatchFailure(
+        launch_op, "Cannot convert with async dependency or result.");
+  }
+
+  Location loc = launch_op.getLoc();
+
+  // Create an LLVM global with CUBIN extracted from the kernel annotation and
+  // obtain a pointer to the first byte in it.
+  auto kernel_module = SymbolTable::lookupNearestSymbolFrom<gpu::GPUModuleOp>(
+      launch_op, launch_op.getKernelModuleName());
+  assert(kernel_module && "expected a kernel module");
+
+  auto binary_attr =
+      kernel_module->getAttrOfType<StringAttr>(gpu_binary_annotation_);
+  if (!binary_attr) {
+    kernel_module.emitOpError()
+        << "missing " << gpu_binary_annotation_ << " attribute";
+    return failure();
+  }
+
+  // Create a global for the module blob.
+  SmallString<128> name_buffer(kernel_module.getName());
+  name_buffer.append("_blob");
+  Value module_blob = LLVM::createGlobalString(loc, rewriter, name_buffer.str(),
+                                               binary_attr.getValue(),
+                                               LLVM::Linkage::Internal);
+
+  // Make sure the trailing zero is included in the constant.
+  auto kernel_name = launch_op.getKernelName().getValue();
+  SmallString<128> kernel_name_buffer(kernel_name);
+  kernel_name_buffer.push_back('\0');
+
+  // Create a global for the kernel name.
+  SmallString<128> kernel_name_global_name_buffer;
+  auto kernel_name_global_name =
+      (kernel_module.getName() + "_" + kernel_name + "_kernel_name")
+          .toStringRef(kernel_name_global_name_buffer);
+  auto kernel_name_global = LLVM::createGlobalString(
+      loc, rewriter, kernel_name_global_name, kernel_name_buffer,
+      LLVM::Linkage::Internal);
+
+  // The TensorFlow OpKernelContext is the first argument of the surrounding
+  // LLVMFunc.
+  Value context_arg =
+      launch_op->getParentOfType<LLVM::LLVMFuncOp>().getArgument(0);
+  auto kernel_params = generateParamsArray(launch_op, adaptor, rewriter);
+
+  auto libraryLaunchNameAttr =
+      mlir::StringAttr::get(loc.getContext(), kTfWrapperLibaryLaunchHelperName);
+  auto function = SymbolTable::lookupNearestSymbolFrom<LLVM::LLVMFuncOp>(
+      launch_op, libraryLaunchNameAttr);
+  if (!function) {
+    PatternRewriter::InsertionGuard guard(rewriter);
+    auto function_type = LLVM::LLVMFunctionType::get(
+        llvm_void_type_, {
+                             llvm_pointer_type_, /* void* context */
+                             llvm_pointer_type_, /* void* module_blob */
+                             llvm_pointer_type_, /* void* function_name */
+                             llvm_intptr_type_,  /* intptr_t grid_x_dim */
+                             llvm_intptr_type_,  /* intptr_t grid_y_dim */
+                             llvm_intptr_type_,  /* intptr_t grid_z_dim */
+                             llvm_intptr_type_,  /* intptr_t block_x_dim */
+                             llvm_intptr_type_,  /* intptr_t block_y_dim */
+                             llvm_intptr_type_,  /* intptr_t block_z_dim */
+                             llvm_pointer_type_, /* void **kernel_params */
+                         });
+    rewriter.setInsertionPointToStart(
+        launch_op->getParentOfType<ModuleOp>().getBody());
+    function = LLVM::LLVMFuncOp::create(
+        rewriter, loc, kTfWrapperLibaryLaunchHelperName, function_type);
+  }
+  LLVM::CallOp::create(
+      rewriter, loc, TypeRange(), mlir::SymbolRefAttr::get(function),
+
+      ArrayRef<Value>{context_arg, module_blob, kernel_name_global,
+                      adaptor.getGridSizeX(), adaptor.getGridSizeY(),
+                      adaptor.getGridSizeZ(), adaptor.getBlockSizeX(),
+                      adaptor.getBlockSizeY(), adaptor.getBlockSizeZ(),
+                      kernel_params});
+
+  rewriter.eraseOp(launch_op);
+  return success();
+}
+
+class TFKernelToLLVMPass
+    : public impl::TFKernelToLLVMPassBase<TFKernelToLLVMPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<LLVM::LLVMDialect>();
+  }
+
+ public:
+  explicit TFKernelToLLVMPass(StringRef blob_annotation) {
+    if (!blob_annotation.empty()) {
+      blob_annotation_ = blob_annotation.str();
+    }
+  }
+
+  void runOnOperation() override {
+    ModuleOp m = getOperation();
+
+    // Populate type conversions.
+    MLIRContext *ctx = m.getContext();
+    LowerToLLVMOptions options(&getContext());
+    LLVMTypeConverter type_converter(ctx, options);
+    type_converter.addConversion([&](tf_framework::OpKernelContextType type) {
+      return LLVM::LLVMPointerType::get(type.getContext());
+    });
+    type_converter.addConversion([&](tf_framework::JITCallableType type) {
+      return LLVM::LLVMPointerType::get(type.getContext());
+    });
+
+    // Populate patterns.
+    RewritePatternSet patterns(&getContext());
+    arith::populateArithExpandOpsPatterns(patterns);
+    memref::populateExpandOpsPatterns(patterns);
+    arith::populateArithToLLVMConversionPatterns(type_converter, patterns);
+    populateFinalizeMemRefToLLVMConversionPatterns(type_converter, patterns);
+    populateMathToLLVMConversionPatterns(type_converter, patterns);
+    populateFuncToLLVMConversionPatterns(type_converter, patterns);
+    cf::populateControlFlowToLLVMConversionPatterns(type_converter, patterns);
+    populateComplexToLLVMConversionPatterns(type_converter, patterns);
+    populateVectorToLLVMConversionPatterns(type_converter, patterns);
+    populateMathToLibmConversionPatterns(patterns);
+    tf_framework::PopulateTFFrameworkToLLVMConversionPatterns(&type_converter,
+                                                              &patterns);
+    patterns.add<ConvertLaunchFuncOpToTfRuntimeCallPattern>(type_converter,
+                                                            blob_annotation_);
+    //  Set target.
+    ConversionTarget target(*ctx);
+    target.addLegalDialect<LLVM::LLVMDialect>();
+    target.addIllegalDialect<
+        arith::ArithDialect, func::FuncDialect, complex::ComplexDialect,
+        gpu::GPUDialect, tf_framework::TFFrameworkDialect, math::MathDialect>();
+    // Mark modules as legal.
+    target.addLegalOp<ModuleOp, gpu::GPUModuleOp>();
+    // Do not look into gpu modules, only consider host-side.
+    target.markOpRecursivelyLegal<gpu::GPUModuleOp>();
+    // Unrealized conversion casts are cleaned up by a separate pass.
+    target.addLegalOp<UnrealizedConversionCastOp>();
+
+    if (failed(applyFullConversion(m, target, std::move(patterns)))) {
+      signalPassFailure();
+    }
+
+    // Finally, strip the GPU modules, as they are no longer needed.
+    for (auto op : llvm::make_early_inc_range(m.getOps<gpu::GPUModuleOp>())) {
+      op.erase();
+    }
+  }
+};
+
+}  // namespace
+
+std::unique_ptr<OperationPass<ModuleOp> > CreateTFKernelToLLVMPass(
+    StringRef blob_annotation) {
+  return std::make_unique<TFKernelToLLVMPass>(blob_annotation);
+}
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/utils.cc

@@ -0,0 +1,71 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/utils.h"
+
+#include <string>
+
+#include "llvm/Support/FormatVariadic.h"
+#include "mlir/Dialect/LLVMIR/LLVMAttrs.h"  // from @llvm-project
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+
+using LLVM::LLVMFuncOp;
+
+FlatSymbolRefAttr GetOrInsertLLVMFunction(StringRef func_name, Type func_type,
+                                          Operation* op, OpBuilder* b) {
+  auto module = op->getParentOfType<ModuleOp>();
+  auto tf_func = module.lookupSymbol<LLVMFuncOp>(func_name);
+  if (!tf_func) {
+    OpBuilder::InsertionGuard guard(*b);
+    b->setInsertionPointToStart(module.getBody());
+    tf_func = LLVMFuncOp::create(*b, b->getUnknownLoc(), func_name, func_type);
+  }
+  return SymbolRefAttr::get(b->getContext(), func_name);
+}
+
+std::string GetGlobalName(StringRef base, StringRef content) {
+  return llvm::formatv("{0}_{1}", base, llvm::hash_value(content));
+}
+
+Value CreateOrFindGlobalStringConstant(Location loc, StringRef global_name,
+                                       StringRef content, OpBuilder* b) {
+  auto module =
+      b->getInsertionBlock()->getParentOp()->getParentOfType<ModuleOp>();
+  Operation* global_constant = SymbolTable::lookupNearestSymbolFrom(
+      module, b->getStringAttr(global_name));
+  if (global_constant) {
+    auto global_op = cast<LLVM::GlobalOp>(global_constant);
+    StringRef symbol_name = global_op.getName();
+    Type symbol_type = global_op.getType();
+    Type ptr_type = LLVM::LLVMPointerType::get(b->getContext());
+    Value global_ptr =
+        LLVM::AddressOfOp::create(*b, loc, ptr_type, symbol_name);
+    Value c0 =
+        LLVM::ConstantOp::create(*b, loc, b->getI64Type(), b->getIndexAttr(0));
+    return LLVM::GEPOp::create(*b, loc, ptr_type, symbol_type, global_ptr,
+                               ValueRange{c0, c0});
+  }
+  return LLVM::createGlobalString(loc, *b, global_name, content,
+                                  LLVM::Linkage::Internal);
+}
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir

tensorflow/compiler/mlir/tools/kernel_gen/transforms/utils.h

@@ -0,0 +1,43 @@
+/* Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TRANSFORMS_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TRANSFORMS_UTILS_H_
+
+#include <string>
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+
+namespace mlir {
+namespace kernel_gen {
+namespace transforms {
+
+// Attempts to find function symbol in the module, adds it if not found.
+FlatSymbolRefAttr GetOrInsertLLVMFunction(StringRef func_name, Type func_type,
+                                          Operation* op, OpBuilder* b);
+
+// Attemts to find a global string constant in the module, adds it if not found.
+Value CreateOrFindGlobalStringConstant(Location loc, StringRef global_name,
+                                       StringRef content, OpBuilder* builder);
+
+// Generates a global name with the format "base_hash(content)".
+std::string GetGlobalName(StringRef base, StringRef content);
+
+}  // namespace transforms
+}  // namespace kernel_gen
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TRANSFORMS_UTILS_H_

tensorflow/compiler/mlir/tools/optimize/BUILD

@@ -0,0 +1,23 @@
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow:tensorflow.default.bzl", "get_compatible_with_portable")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "quantization_utils",
+    srcs = ["quantization_utils.cc"],
+    hdrs = ["quantization_utils.h"],
+    compatible_with = get_compatible_with_portable(),
+    deps = [
+        "//tensorflow/compiler/mlir/tools:safe_cast",
+        "//tensorflow/core/framework:tensor_shape",
+        "@com_google_absl//absl/log:check",
+        "@com_google_absl//absl/types:span",
+    ],
+)

tensorflow/compiler/mlir/tools/optimize/quantization_utils.cc

@@ -0,0 +1,129 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/compiler/mlir/tools/optimize/quantization_utils.h"
+
+#include <algorithm>
+#include <cmath>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <vector>
+
+#include "absl/log/check.h"
+#include "absl/types/span.h"
+#include "tensorflow/compiler/mlir/tools/safe_cast.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+
+namespace tflite_migration {
+namespace optimize {
+namespace utils {
+
+namespace {
+
+const int8_t kMinQuantizedValue8bit = -127;
+const int8_t kMaxQuantizedValue8bit = 127;
+
+// const int8_t kMinQuantizedValue4bit = -7;
+// const int8_t kMaxQuantizedValue4bit = 7;
+
+// The maximum number of dimensions supported in per-channel quantization.
+constexpr int kPerChannelMaxDim = 4;
+}  // namespace
+
+template <class BiasType>
+std::vector<BiasType> SymmetricBiasQuantize(const float* data,
+                                            uint64_t num_elements,
+                                            const std::vector<float>& scales) {
+  std::vector<BiasType> buffer(num_elements);
+  const BiasType kScale = std::numeric_limits<BiasType>::max();
+  float scaling_factor_inv_per_layer = (scales[0] == 0) ? 0 : 1.0 / scales[0];
+
+  for (int32_t idx = 0; idx < num_elements; idx++) {
+    float scaling_factor_inv =
+        scales.size() == 1 ? scaling_factor_inv_per_layer
+                           : ((scales[idx] == 0) ? 0 : 1.0 / scales[idx]);
+    const BiasType quantized_value =
+        tools::SafeCast<BiasType>(std::round(data[idx] * scaling_factor_inv));
+    buffer[idx] = std::min(kScale, std::max(-kScale, quantized_value));
+  }
+  return buffer;
+}
+
+template std::vector<std::int32_t> SymmetricBiasQuantize<std::int32_t>(
+    const float* data, uint64_t num_elements, const std::vector<float>& scales);
+
+template std::vector<std::int64_t> SymmetricBiasQuantize<std::int64_t>(
+    const float* data, uint64_t num_elements, const std::vector<float>& scales);
+
+std::vector<int16_t> SymmetricQuantizeFloatsToInt16(const float* data,
+                                                    uint64_t num_elements,
+                                                    float scaling_factor) {
+  // Compute the inverse of scale.
+  const float scaling_factor_inv =
+      (scaling_factor == 0) ? 0 : 1.0 / scaling_factor;
+  std::vector<int16_t> buffer(num_elements);
+  const int32_t kScale = std::numeric_limits<int16_t>::max();
+
+  for (size_t i = 0; i < num_elements; i++) {
+    const int32_t quantized_value =
+        static_cast<int32_t>(std::round(data[i] * scaling_factor_inv));
+    buffer[i] = std::min(kScale, std::max(-kScale, quantized_value));
+  }
+  return buffer;
+}
+
+void SymmetricPerChannelQuantizeValues(const float* const input,
+                                       const std::vector<float>& scales_inv,
+                                       const std::vector<int32_t>& dimension,
+                                       int32_t channel_dim_index,
+                                       std::vector<int8_t>* output_value) {
+  // Quantize the values.
+  int indices[kPerChannelMaxDim];
+  tensorflow::TensorShape unextended_shape;
+  CHECK_OK(tensorflow::TensorShapeUtils::MakeShape(absl::MakeSpan(dimension),
+                                                   &unextended_shape));
+  tensorflow::TensorShape shape;
+  for (int i = 0; i < kPerChannelMaxDim - unextended_shape.dims(); ++i) {
+    CHECK_OK(shape.AddDimWithStatus(1));
+  }
+  CHECK_OK(shape.AppendShapeWithStatus(unextended_shape));
+  channel_dim_index += kPerChannelMaxDim - unextended_shape.dims();
+
+  for (indices[0] = 0; indices[0] < shape.dim_size(0); indices[0]++) {
+    for (indices[1] = 0; indices[1] < shape.dim_size(1); indices[1]++) {
+      for (indices[2] = 0; indices[2] < shape.dim_size(2); indices[2]++) {
+        for (indices[3] = 0; indices[3] < shape.dim_size(3); indices[3]++) {
+          int channel_idx = indices[channel_dim_index];
+          int index = 0;
+          int current_stride = 1;
+          for (int i = kPerChannelMaxDim - 1; i >= 0; --i) {
+            index += indices[i] * current_stride;
+            current_stride *= shape.dim_size(i);
+          }
+          const float val = input[index];
+          const int32_t quantized_value =
+              static_cast<int32_t>(std::round(val * scales_inv[channel_idx]));
+          output_value->at(index) = std::min<int8_t>(
+              kMaxQuantizedValue8bit,
+              std::max<int8_t>(kMinQuantizedValue8bit, quantized_value));
+        }
+      }
+    }
+  }
+}
+
+}  // namespace utils
+}  // namespace optimize
+}  // namespace tflite_migration

tensorflow/compiler/mlir/tools/optimize/quantization_utils.h

@@ -0,0 +1,45 @@
+/* Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TOOLS_OPTIMIZE_QUANTIZATION_UTILS_H_
+#define TENSORFLOW_COMPILER_MLIR_TOOLS_OPTIMIZE_QUANTIZATION_UTILS_H_
+
+#include <cstdint>
+#include <vector>
+
+namespace tflite_migration {
+namespace optimize {
+namespace utils {
+
+template <typename BiasType>
+std::vector<BiasType> SymmetricBiasQuantize(const float* data,
+                                            uint64_t num_elements,
+                                            const std::vector<float>& scales);
+
+std::vector<int16_t> SymmetricQuantizeFloatsToInt16(const float* data,
+                                                    uint64_t num_elements,
+                                                    float scaling_factor);
+
+// Quantize the values given an array of scales.
+void SymmetricPerChannelQuantizeValues(const float* input,
+                                       const std::vector<float>& scales_inv,
+                                       const std::vector<int32_t>& dimension,
+                                       int32_t channel_dim_index,
+                                       std::vector<int8_t>* output_value);
+
+}  // namespace utils
+}  // namespace optimize
+}  // namespace tflite_migration
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOOLS_OPTIMIZE_QUANTIZATION_UTILS_H_

tensorflow/compiler/mlir/tools/safe_cast.h

@@ -0,0 +1,89 @@
+/* Copyright 2025 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_COMPILER_MLIR_TOOLS_SAFE_CAST_H_
+#define TENSORFLOW_COMPILER_MLIR_TOOLS_SAFE_CAST_H_
+
+#include <cmath>
+#include <limits>
+
+namespace tools {
+
+// Converts a floating-point number to an integer. For all inputs x where
+// static_cast<IntOut>(x) is legal according to the C++ standard, the result
+// is identical to that cast (i.e. the result is x with its fractional part
+// truncated whenever that is representable as IntOut).
+//
+// static_cast would cause undefined behavior for the following cases, which
+// have well-defined behavior for this function:
+//
+//  1. If x is NaN, the result is zero.
+//
+//  2. If the truncated form of x is above the representable range of IntOut,
+//     the result is std::numeric_limits<IntOut>::max().
+//
+//  3. If the truncated form of x is below the representable range of IntOut,
+//     the result is std::numeric_limits<IntOut>::min().
+//
+// Note that cases #2 and #3 cover infinities as well as finite numbers.
+//
+// The range of FloatIn must include the range of IntOut, otherwise
+// the results are undefined.
+// TODO(sfeuz): Replace by absl::SafeCast once available.
+template <class IntOut, class FloatIn>
+IntOut SafeCast(FloatIn x) {
+  static_assert(!std::numeric_limits<FloatIn>::is_integer,
+                "FloatIn is integer");
+  static_assert(std::numeric_limits<IntOut>::is_integer,
+                "IntOut is not integer");
+  static_assert(std::numeric_limits<IntOut>::radix == 2, "IntOut is base 2");
+
+  // Special case NaN, for which the logic below doesn't work.
+  if (std::isnan(x)) {
+    return 0;
+  }
+
+  // Negative values all clip to zero for unsigned results.
+  if (!std::numeric_limits<IntOut>::is_signed && x < 0) {
+    return 0;
+  }
+
+  // Handle infinities.
+  if (std::isinf(x)) {
+    return x < 0 ? std::numeric_limits<IntOut>::min()
+                 : std::numeric_limits<IntOut>::max();
+  }
+
+  // Set exp such that x == f * 2^exp for some f with |f| in [0.5, 1.0),
+  // unless x is zero in which case exp == 0. Note that this implies that the
+  // magnitude of x is strictly less than 2^exp.
+  int exp = 0;
+  std::frexp(x, &exp);
+
+  // Let N be the number of non-sign bits in the representation of IntOut. If
+  // the magnitude of x is strictly less than 2^N, the truncated version of x
+  // is representable as IntOut. The only representable integer for which this
+  // is not the case is kMin for signed types (i.e. -2^N), but that is covered
+  // by the fall-through below.
+  if (exp <= std::numeric_limits<IntOut>::digits) {
+    return x;
+  }
+
+  // Handle numbers with magnitude >= 2^N.
+  return x < 0 ? std::numeric_limits<IntOut>::min()
+               : std::numeric_limits<IntOut>::max();
+}
+}  // namespace tools
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOOLS_SAFE_CAST_H_

tensorflow/compiler/mlir/tools/tf_mlir_translate_cl.cc

@@ -0,0 +1,161 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/compiler/mlir/tools/tf_mlir_translate_cl.h"
+
+#include <string>
+
+#include "llvm/Support/CommandLine.h"
+
+// These command-line options are following LLVM conventions because we also
+// need to register the TF Graph(Def) to MLIR conversion with mlir-translate,
+// which expects command-line options of such style.
+
+using llvm::cl::opt;
+
+// Import options.
+// NOLINTNEXTLINE
+opt<std::string> input_arrays(
+    "tf-input-arrays", llvm::cl::desc("Input tensor names, separated by ','"),
+    llvm::cl::init(""));
+
+// NOLINTNEXTLINE
+opt<std::string> input_dtypes(
+    "tf-input-data-types",
+    llvm::cl::desc("(Optional) Input tensor data types, separated by ','. Use "
+                   "'' if a single data type is skipped. The data type from "
+                   "the import graph is used if it is skipped."),
+    llvm::cl::init(""));
+
+// NOLINTNEXTLINE
+opt<std::string> input_shapes(
+    "tf-input-shapes",
+    llvm::cl::desc(
+        "Input tensor shapes. Shapes for different tensors are separated by "
+        "':', and dimension sizes for the same tensor are separated by ','"),
+    llvm::cl::init(""));
+
+// NOLINTNEXTLINE
+opt<std::string> output_arrays(
+    "tf-output-arrays", llvm::cl::desc("Output tensor names, separated by ','"),
+    llvm::cl::init(""));
+
+// NOLINTNEXTLINE
+opt<std::string> control_output_arrays(
+    "tf-control-output-arrays",
+    llvm::cl::desc("Control output node names, separated by ','"),
+    llvm::cl::init(""));
+
+// NOLINTNEXTLINE
+opt<std::string> inference_type(
+    "tf-inference-type",
+    llvm::cl::desc(
+        "Sets the type of real-number arrays in the output file. Only allows "
+        "float and quantized types"),
+    llvm::cl::init(""));
+
+// NOLINTNEXTLINE
+opt<std::string> min_values(
+    "tf-input-min-values",
+    llvm::cl::desc(
+        "Sets the lower bound of the input data. Separated by ','; Each entry "
+        "in the list should match an entry in -tf-input-arrays. This is "
+        "used when -tf-inference-type is a quantized type."),
+    llvm::cl::Optional, llvm::cl::init(""));
+
+// NOLINTNEXTLINE
+opt<std::string> max_values(
+    "tf-input-max-values",
+    llvm::cl::desc(
+        "Sets the upper bound of the input data. Separated by ','; Each entry "
+        "in the list should match an entry in -tf-input-arrays. This is "
+        "used when -tf-inference-type is a quantized type."),
+    llvm::cl::Optional, llvm::cl::init(""));
+
+// NOLINTNEXTLINE
+opt<std::string> debug_info_file(
+    "tf-debug-info",
+    llvm::cl::desc("Path to the debug info file of the input graph def"),
+    llvm::cl::init(""));
+
+// NOLINTNEXTLINE
+opt<std::string> xla_compile_device_type(
+    "tf-xla-compile-device-type",
+    llvm::cl::desc("Sets the compilation device type of the input graph def"),
+    llvm::cl::init(""));
+
+// TODO(b/134792656): If pruning is moved into TF dialect as a pass
+// we should remove this.
+// NOLINTNEXTLINE
+opt<bool> prune_unused_nodes(
+    "tf-prune-unused-nodes",
+    llvm::cl::desc("Prune unused nodes in the input graphdef"),
+    llvm::cl::init(false));
+
+// NOLINTNEXTLINE
+opt<bool> convert_legacy_fed_inputs(
+    "tf-convert-legacy-fed-inputs",
+    llvm::cl::desc(
+        "Eliminate LegacyFedInput nodes by replacing them with Placeholder"),
+    llvm::cl::init(false));
+
+// NOLINTNEXTLINE
+opt<bool> graph_as_function("tf-graph-as-function",
+                            llvm::cl::desc("Treat main graph as a function"),
+                            llvm::cl::init(false));
+
+// NOLINTNEXTLINE
+opt<bool> upgrade_legacy("tf-upgrade-legacy",
+                         llvm::cl::desc("Upgrade legacy TF graph behavior"),
+                         llvm::cl::init(false));
+
+// NOLINTNEXTLINE
+opt<bool> enable_shape_inference(
+    "tf-enable-shape-inference-on-import",
+    llvm::cl::desc("Enable shape inference on import (temporary)"),
+    llvm::cl::init(false));
+
+// NOLINTNEXTLINE
+opt<bool> unconditionally_use_set_output_shapes(
+    "tf-enable-unconditionally-use-set-output-shapes-on-import",
+    llvm::cl::desc("Enable using the _output_shapes unconditionally on import "
+                   "(temporary)"),
+    llvm::cl::init(false));
+
+// NOLINTNEXTLINE
+opt<bool> enable_soft_placement(
+    "tf-enable-soft-placement-on-import",
+    llvm::cl::desc("Enable soft device placement on import."),
+    llvm::cl::init(false));
+
+// NOLINTNEXTLINE
+opt<bool> set_original_tf_func_name(
+    "tf-set-original-tf-func-name-on-import",
+    llvm::cl::desc("Set original TF function name on importi."),
+    llvm::cl::init(false));
+
+// Export options.
+// NOLINTNEXTLINE
+opt<bool> export_entry_func_to_flib(
+    "tf-export-entry-func-to-flib",
+    llvm::cl::desc(
+        "Export entry function to function library instead of graph"),
+    llvm::cl::init(false));
+// NOLINTNEXTLINE
+opt<bool> export_original_tf_func_name(
+    "tf-export-original-func-name",
+    llvm::cl::desc("Export functions using the name set in the attribute "
+                   "'tf._original_func_name' if it exists."),
+    llvm::cl::init(false));

tensorflow/compiler/mlir/tools/tf_mlir_translate_cl.h

@@ -0,0 +1,54 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TOOLS_TF_MLIR_TRANSLATE_CL_H_
+#define TENSORFLOW_COMPILER_MLIR_TOOLS_TF_MLIR_TRANSLATE_CL_H_
+
+// This file contains command-line options aimed to provide the parameters
+// required by the TensorFlow Graph(Def) to MLIR module conversion. It is only
+// intended to be included by binaries.
+
+#include <string>
+
+#include "llvm/Support/CommandLine.h"
+
+// Please see the implementation file for documentation of these options.
+
+// Import options.
+extern llvm::cl::opt<std::string> input_arrays;
+extern llvm::cl::opt<std::string> input_dtypes;
+extern llvm::cl::opt<std::string> input_shapes;
+extern llvm::cl::opt<std::string> output_arrays;
+extern llvm::cl::opt<std::string> control_output_arrays;
+extern llvm::cl::opt<std::string> inference_type;
+extern llvm::cl::opt<std::string> min_values;
+extern llvm::cl::opt<std::string> max_values;
+extern llvm::cl::opt<std::string> debug_info_file;
+extern llvm::cl::opt<std::string> xla_compile_device_type;
+extern llvm::cl::opt<bool> prune_unused_nodes;
+extern llvm::cl::opt<bool> convert_legacy_fed_inputs;
+extern llvm::cl::opt<bool> graph_as_function;
+extern llvm::cl::opt<bool> upgrade_legacy;
+// TODO(jpienaar): Temporary flag, flip default and remove.
+extern llvm::cl::opt<bool> enable_shape_inference;
+extern llvm::cl::opt<bool> unconditionally_use_set_output_shapes;
+extern llvm::cl::opt<bool> enable_soft_placement;
+extern llvm::cl::opt<bool> set_original_tf_func_name;
+
+// Export options.
+extern llvm::cl::opt<bool> export_entry_func_to_flib;
+extern llvm::cl::opt<bool> export_original_tf_func_name;
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TOOLS_TF_MLIR_TRANSLATE_CL_H_

tensorflow/compiler/mlir/tools/tf_mlir_translate_registration.cc

@@ -0,0 +1,51 @@
+/* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+// This file wraps TensorFlow Graph(Def) to MLIR module conversion into passes
+// to satisfy the API of MLIR pass registration. In order to do this, the
+// command-line option header is pulled in.
+
+#include <utility>
+
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/Tools/mlir-translate/Translation.h"  // from @llvm-project
+#include "tensorflow/compiler/mlir/tensorflow/translate/tools/file_tf_mlir_translate.h"
+#include "tensorflow/compiler/mlir/tools/tf_mlir_translate_cl.h"
+#include "tensorflow/core/framework/graph.pb.h"
+
+namespace mlir {
+using tsl::Status;
+using tsl::StatusOr;
+
+
+static OwningOpRef<mlir::ModuleOp> GraphdefToSplattedMlirTranslateFunction(
+    llvm::StringRef input, MLIRContext* context) {
+  tensorflow::GraphdefToMlirOptions options{
+      debug_info_file,        xla_compile_device_type,
+      prune_unused_nodes,     convert_legacy_fed_inputs,
+      graph_as_function,      upgrade_legacy,
+      enable_shape_inference, unconditionally_use_set_output_shapes};
+  auto module_or = tensorflow::GraphdefToSplattedMlirTranslateFunction(
+      input, input_arrays, input_dtypes, input_shapes, output_arrays,
+      control_output_arrays, options, context);
+  if (!module_or.status().ok()) return nullptr;
+  return std::move(module_or).value();
+}
+
+static TranslateToMLIRRegistration GraphdefToSplattedMlirTranslate(
+    "graphdef-to-splatted-mlir", "graphdef-to-splatted-mlir",
+    GraphdefToSplattedMlirTranslateFunction);
+
+}  // namespace mlir