From 884b55fe86627271d9483e510e4b5d0d3eccf53d Mon Sep 17 00:00:00 2001 From: Adrian Kuegel Date: Thu, 25 Jun 2020 06:08:43 -0700 Subject: [PATCH] Workaround for the missing libtensorflow_framework.so when running tf_to_cubin. PiperOrigin-RevId: 318257426 Change-Id: I83440763f2224463966c40138b5d72f3de741e51 --- tensorflow/core/kernels/cubin_headers/build_defs.bzl | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tensorflow/core/kernels/cubin_headers/build_defs.bzl b/tensorflow/core/kernels/cubin_headers/build_defs.bzl index cdf77e59433..d60aaadcbf6 100644 --- a/tensorflow/core/kernels/cubin_headers/build_defs.bzl +++ b/tensorflow/core/kernels/cubin_headers/build_defs.bzl @@ -27,7 +27,7 @@ def _gen_kernel_cubin_impl(ctx): filename = "%s.%s.cubin" % (name, arch) cubin = ctx.actions.declare_file(filename) ctx.actions.run( - inputs = [ctx.file.mlir_op], + inputs = [ctx.file.mlir_op, ctx.file._tfso], outputs = [cubin], executable = ctx.executable._tool, arguments = cmd_args + [ @@ -49,6 +49,11 @@ _gen_kernel_cubin_rule = rule( "same_shape": attr.string(), "unroll_factors": attr.string(), "gpu_archs": attr.string_list(mandatory = True), + "_tfso": attr.label( + default = Label("//tensorflow:libtensorflow_framework.so.2"), + cfg = "host", + allow_single_file = True, + ), "_tool": attr.label( executable = True, default = Label("//tensorflow/compiler/mlir/tools/kernel_gen:tf_to_cubin"),