experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Workaround for the missing libtensorflow_framework.so when running tf_to_cubin.

Browse Source 
PiperOrigin-RevId: 318257426
Change-Id: I83440763f2224463966c40138b5d72f3de741e51
This commit is contained in:
Adrian Kuegel 2020-06-25 06:08:43 -07:00 committed by TensorFlower Gardener
parent 3398514578
commit 884b55fe86
1 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
tensorflow/core/kernels/cubin_headers/build_defs.bzl
View File

@ -27,7 +27,7 @@ def _gen_kernel_cubin_impl(ctx):
filename = "%s.%s.cubin" % (name, arch) filename = "%s.%s.cubin" % (name, arch)
cubin = ctx.actions.declare_file(filename) cubin = ctx.actions.declare_file(filename)
ctx.actions.run( ctx.actions.run(
inputs = [ctx.file.mlir_op], inputs = [ctx.file.mlir_op, ctx.file._tfso],
outputs = [cubin], outputs = [cubin],
executable = ctx.executable._tool, executable = ctx.executable._tool,
arguments = cmd_args + [ arguments = cmd_args + [
@ -49,6 +49,11 @@ _gen_kernel_cubin_rule = rule(
"same_shape": attr.string(), "same_shape": attr.string(),
"unroll_factors": attr.string(), "unroll_factors": attr.string(),
"gpu_archs": attr.string_list(mandatory = True), "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( "_tool": attr.label(
executable = True, executable = True,
default = Label("//tensorflow/compiler/mlir/tools/kernel_gen:tf_to_cubin"), default = Label("//tensorflow/compiler/mlir/tools/kernel_gen:tf_to_cubin"),