Home | History | Annotate | Download | only in fused_conv 
      1 # Description:
      2 #   A Fused Conv Bias Activation operator wrapper.
      3 #   APIs are meant to change over time.
      4 
      5 package(
      6     default_visibility = ["//visibility:private"],
      7 )
      8 
      9 package_group(
     10     name = "friends",
     11     packages = [
     12         "//tensorflow/...",
     13     ],
     14 )
     15 
     16 licenses(["notice"])  # Apache 2.0
     17 
     18 exports_files(["LICENSE"])
     19 
     20 load(
     21     "//tensorflow:tensorflow.bzl",
     22     "tf_kernel_library",
     23     "tf_custom_op_library",
     24     "tf_gen_op_libs",
     25     "tf_gen_op_wrapper_py",
     26 )
     27 load("//tensorflow:tensorflow.bzl", "cuda_py_test")
     28 load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library")
     29 
     30 tf_custom_op_py_library(
     31     name = "fused_conv_py",
     32     srcs = ["__init__.py"] + glob(["python/ops/*.py"]),
     33     dso = [":python/ops/_fused_conv2d_bias_activation_op.so"],
     34     kernels = [
     35         ":fused_conv2d_bias_activation_op_kernels",
     36         ":fused_conv2d_bias_activation_op_op_lib",
     37     ],
     38     srcs_version = "PY2AND3",
     39     visibility = ["//visibility:public"],
     40     deps = [
     41         ":fused_conv2d_bias_activation_op",
     42         "//tensorflow/contrib/util:util_py",
     43         "//tensorflow/python:array_ops",
     44         "//tensorflow/python:client_testlib",
     45         "//tensorflow/python:constant_op",
     46         "//tensorflow/python:control_flow_ops",
     47         "//tensorflow/python:dtypes",
     48         "//tensorflow/python:errors",
     49         "//tensorflow/python:framework_ops",
     50         "//tensorflow/python:framework_test_lib",
     51         "//tensorflow/python:nn_ops",
     52         "//tensorflow/python:platform",
     53         "//tensorflow/python:random_ops",
     54         "//tensorflow/python:session",
     55         "//tensorflow/python:util",
     56         "//tensorflow/python:variables",
     57         "//third_party/py/numpy",
     58     ],
     59 )
     60 
     61 tf_kernel_library(
     62     name = "fused_conv2d_bias_activation_op_kernels",
     63     srcs = [
     64         "kernels/fused_conv2d_bias_activation_op.cc",
     65         "kernels/fused_conv2d_bias_activation_op.h",
     66         "kernels/fused_conv_ops_gpu.h",
     67     ],
     68     prefix = "fused_conv2d_bias_activation_op",
     69     visibility = ["//visibility:public"],
     70     deps = [
     71         "//tensorflow/core:autotuning_proto_cc",
     72         "//tensorflow/core:conv_autotuning_proto_cc",
     73         "//tensorflow/core:framework",
     74         "//tensorflow/core:lib",
     75         "//tensorflow/core:lib_proto_parsing",
     76         "//tensorflow/core:stream_executor",
     77         "//tensorflow/core/kernels:bounds_check",
     78         "//tensorflow/core/kernels:conv_2d_hdrs",
     79         "//tensorflow/core/kernels:conv_ops_gpu_hdrs",
     80         "//tensorflow/core/kernels:gpu_util_hdrs",
     81         "//tensorflow/core/kernels:ops_util_hdrs",
     82         "//third_party/eigen3",
     83         "@local_config_cuda//cuda:cudnn_header",
     84     ],
     85     alwayslink = 1,
     86 )
     87 
     88 tf_custom_op_library(
     89     name = "python/ops/_fused_conv2d_bias_activation_op.so",
     90     srcs = [
     91         "kernels/fused_conv2d_bias_activation_op.cc",
     92         "kernels/fused_conv2d_bias_activation_op.h",
     93         "kernels/fused_conv_ops_gpu.h",
     94         "ops/fused_conv2d_bias_activation_op.cc",
     95     ],
     96     deps = [
     97         "//tensorflow/core:autotuning_proto_cc",
     98         "//tensorflow/core:conv_autotuning_proto_cc",
     99         "//tensorflow/core:lib_proto_parsing",
    100         "//tensorflow/core/kernels:bounds_check_lib",
    101         "//tensorflow/core/kernels:conv_2d_hdrs",
    102         "//tensorflow/core/kernels:conv_ops_gpu_hdrs",
    103         "//tensorflow/core/kernels:gpu_util_hdrs",
    104         "//tensorflow/core/kernels:ops_util_hdrs",
    105         "@local_config_cuda//cuda:cudnn_header",
    106     ],
    107 )
    108 
    109 tf_gen_op_libs(
    110     op_lib_names = ["fused_conv2d_bias_activation_op"],
    111     deps = ["//tensorflow/core:lib_proto_parsing"],
    112 )
    113 
    114 tf_gen_op_wrapper_py(
    115     name = "fused_conv2d_bias_activation_op",
    116     deps = [":fused_conv2d_bias_activation_op_op_lib"],
    117 )
    118 
    119 py_library(
    120     name = "fused_conv2d_bias_activation_op_test_base",
    121     testonly = 1,
    122     srcs = ["python/ops/fused_conv2d_bias_activation_op_test_base.py"],
    123     visibility = ["//tensorflow/compiler/tf2xla:internal"],
    124     deps = [
    125         ":fused_conv_py",
    126         "//tensorflow/python:array_ops",
    127         "//tensorflow/python:client",
    128         "//tensorflow/python:client_testlib",
    129         "//tensorflow/python:errors",
    130         "//tensorflow/python:framework_for_generated_wrappers",
    131         "//tensorflow/python:framework_test_lib",
    132         "//tensorflow/python:nn",
    133         "//tensorflow/python:nn_ops",
    134         "//tensorflow/python:platform",
    135         "//tensorflow/python:random_ops",
    136         "//tensorflow/python:training",
    137         "//tensorflow/python:variables",
    138         "//third_party/py/numpy",
    139         "@absl_py//absl/testing:parameterized",
    140     ],
    141 )
    142 
    143 cuda_py_test(
    144     name = "fused_conv2d_bias_activation_op_test",
    145     size = "large",
    146     srcs = ["python/ops/fused_conv2d_bias_activation_op_test.py"],
    147     additional_deps = [
    148         ":fused_conv2d_bias_activation_op_test_base",
    149         "//tensorflow/python:client_testlib",
    150     ],
    151     tags = [
    152         "manual",  # TODO(b/117128481): re-enable after fixing OSS build
    153         "no_pip",
    154         "requires-gpu-sm70",
    155     ],
    156 )
    157 
    158 cuda_py_test(
    159     name = "fused_conv2d_bias_activation_benchmark",
    160     srcs = ["python/ops/fused_conv2d_bias_activation_benchmark.py"],
    161     additional_deps = [
    162         ":fused_conv_py",
    163         "//tensorflow/python:client",
    164         "//tensorflow/python:client_testlib",
    165         "//tensorflow/python:control_flow_ops",
    166         "//tensorflow/python:framework_for_generated_wrappers",
    167         "//tensorflow/python:nn_ops",
    168         "//tensorflow/python:platform",
    169         "//tensorflow/python:platform_benchmark",
    170         "//tensorflow/python:random_ops",
    171         "//tensorflow/python:variables",
    172         "//third_party/py/numpy",
    173         "//tensorflow/core:protos_all_py",
    174     ],
    175     main = "python/ops/fused_conv2d_bias_activation_benchmark.py",
    176     tags = [
    177         "manual",  # TODO(b/117128481): re-enable after fixing OSS build
    178         "nogpu",
    179         "requires-gpu-sm70",
    180     ],
    181 )
    182