|
| 1 | +# Licensed to the Apache Software Foundation (ASF) under one |
| 2 | +# or more contributor license agreements. See the NOTICE file |
| 3 | +# distributed with this work for additional information |
| 4 | +# regarding copyright ownership. The ASF licenses this file |
| 5 | +# to you under the Apache License, Version 2.0 (the |
| 6 | +# "License"); you may not use this file except in compliance |
| 7 | +# with the License. You may obtain a copy of the License at |
| 8 | +# |
| 9 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | +# |
| 11 | +# Unless required by applicable law or agreed to in writing, |
| 12 | +# software distributed under the License is distributed on an |
| 13 | +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| 14 | +# KIND, either express or implied. See the License for the |
| 15 | +# specific language governing permissions and limitations |
| 16 | +# under the License. |
| 17 | +""" |
| 18 | +Using Pipeline Executor in Relay |
| 19 | +================================= |
| 20 | +**Author**: `Hua Jiang <https://https://github.com/huajsj>`_ |
| 21 | +
|
| 22 | +This is a short tutorial on how to use "Pipeline Executor" with Relay. |
| 23 | +""" |
| 24 | +import tvm |
| 25 | +from tvm import te |
| 26 | +import numpy as np |
| 27 | +from tvm.contrib import graph_executor as runtime |
| 28 | +from tvm.relay.op.contrib.cutlass import partition_for_cutlass |
| 29 | +from tvm import relay |
| 30 | +from tvm.relay import testing |
| 31 | +import tvm.testing |
| 32 | +from tvm.contrib.cutlass import ( |
| 33 | + has_cutlass, |
| 34 | + num_cutlass_partitions, |
| 35 | + finalize_modules, |
| 36 | + finalize_modules_vm, |
| 37 | +) |
| 38 | + |
| 39 | +img_size = 8 |
| 40 | +####################################################################### |
| 41 | +# Create a simple network, this network can be a pre-trained model too. |
| 42 | +# --------------------------------------------------------------------- |
| 43 | +# Let's create a very simple network for demonstration. |
| 44 | +# It consists of convolution, batch normalization, dense, and ReLU activation. |
| 45 | +def get_network(): |
| 46 | + out_channels = 16 |
| 47 | + batch_size = 1 |
| 48 | + data = relay.var("data", relay.TensorType((batch_size, 3, img_size, img_size), "float16")) |
| 49 | + dense_weight = relay.var( |
| 50 | + "dweight", relay.TensorType((batch_size, 16 * img_size * img_size), "float16") |
| 51 | + ) |
| 52 | + weight = relay.var("weight") |
| 53 | + second_weight = relay.var("second_weight") |
| 54 | + bn_gamma = relay.var("bn_gamma") |
| 55 | + bn_beta = relay.var("bn_beta") |
| 56 | + bn_mmean = relay.var("bn_mean") |
| 57 | + bn_mvar = relay.var("bn_var") |
| 58 | + simple_net = relay.nn.conv2d( |
| 59 | + data=data, weight=weight, kernel_size=(3, 3), channels=out_channels, padding=(1, 1) |
| 60 | + ) |
| 61 | + simple_net = relay.nn.batch_norm(simple_net, bn_gamma, bn_beta, bn_mmean, bn_mvar)[0] |
| 62 | + simple_net = relay.nn.relu(simple_net) |
| 63 | + simple_net = relay.nn.batch_flatten(simple_net) |
| 64 | + simple_net = relay.nn.dense(simple_net, dense_weight) |
| 65 | + simple_net = relay.Function(relay.analysis.free_vars(simple_net), simple_net) |
| 66 | + data_shape = (batch_size, 3, img_size, img_size) |
| 67 | + net, params = testing.create_workload(simple_net) |
| 68 | + return net, params, data_shape |
| 69 | + |
| 70 | + |
| 71 | +net, params, data_shape = get_network() |
| 72 | +########################################### |
| 73 | +# Splitting the network into two subgraphs. |
| 74 | +# ----------------------------------------- |
| 75 | +# This function called 'graph_split' from a unit test is just an example. User can create a customized logic |
| 76 | +# to split the graph. |
| 77 | +import inspect |
| 78 | +import os |
| 79 | + |
| 80 | +tutorial_dir = os.path.dirname(inspect.getfile(lambda: None)) |
| 81 | +os.sys.path.append(os.path.join(tutorial_dir, "../../../tests/python/relay")) |
| 82 | +from test_pipeline_executor import graph_split |
| 83 | + |
| 84 | +########################################### |
| 85 | +# Splitting the network into two subgraphs. |
| 86 | +split_config = [{"op_name": "nn.relu", "op_index": 0}] |
| 87 | +subgraphs = graph_split(net["main"], split_config, params) |
| 88 | +########################################################### |
| 89 | +# The generated subgraphs should look something like below. |
| 90 | + |
| 91 | +""" |
| 92 | +#subgraphs[0]) |
| 93 | +
|
| 94 | + def @main(%data: Tensor[(1, 3, img_size, img_size), float16]) { |
| 95 | + %0 = nn.conv2d(%data, meta[relay.Constant][0] /* ty=Tensor[(16, 3, 3, 3), float16] */, padding=[1, 1, 1, 1], channels=16, kernel_size=[3, 3]) /* ty=Tensor[(1, 16, img_size, img_size), float16] */; |
| 96 | + %1 = nn.batch_norm(%0, meta[relay.Constant][1] /* ty=Tensor[(16), float16] */, meta[relay.Constant][2] /* ty=Tensor[(16), float16]*/, meta[relay.Constant][3] /* ty=Tensor[(16), float16] */, meta[relay.Constant][4] /* ty=Tensor[(16), float16] */) /* ty=(Tensor[(1,16, img_size, img_size), float16], Tensor[(16), float16], Tensor[(16), float16]) */; |
| 97 | + %2 = %1.0; |
| 98 | + nn.relu(%2) /* ty=Tensor[(1, 16, img_size, img_size), float16] */ |
| 99 | + } |
| 100 | +
|
| 101 | +#subgraphs[1] |
| 102 | +
|
| 103 | + def @main(%data_n_0: Tensor[(1, 16, 8, 8), float16] /* ty=Tensor[(1, 16, 8, 8), float16] */) { |
| 104 | + %0 = nn.batch_flatten(%data_n_0) /* ty=Tensor[(1, 1024), float16] */; |
| 105 | + nn.dense(%0, meta[relay.Constant][0] /* ty=Tensor[(1, 1024), float16] */, units=None) /* ty=Tensor[(1, 1), float16] */ |
| 106 | + } |
| 107 | +
|
| 108 | +""" |
| 109 | + |
| 110 | +# sphinx_gallery_start_ignore |
| 111 | +from tvm import testing |
| 112 | + |
| 113 | +testing.utils.install_request_hook(depth=3) |
| 114 | +# sphinx_gallery_end_ignore |
| 115 | + |
| 116 | +######################################### |
| 117 | +# Build the subgraph with cutlass target. |
| 118 | +# --------------------------------------- |
| 119 | + |
| 120 | +cutlass = tvm.target.Target( |
| 121 | + { |
| 122 | + "kind": "cutlass", |
| 123 | + "sm": int(tvm.target.Target("cuda").arch.split("_")[1]), |
| 124 | + "use_3xtf32": True, |
| 125 | + "split_k_slices": [1], |
| 126 | + "profile_all_alignments": False, |
| 127 | + "find_first_valid": True, |
| 128 | + "use_multiprocessing": True, |
| 129 | + "use_fast_math": False, |
| 130 | + "tmp_dir": "./tmp", |
| 131 | + }, |
| 132 | + host=tvm.target.Target("llvm"), |
| 133 | +) |
| 134 | + |
| 135 | + |
| 136 | +def cutlass_build(mod, target, params=None, target_host=None, mod_name="default"): |
| 137 | + target = [target, cutlass] |
| 138 | + lib = relay.build_module.build( |
| 139 | + mod, target=target, params=params, target_host=target_host, mod_name=mod_name |
| 140 | + ) |
| 141 | + return lib |
| 142 | + |
| 143 | + |
| 144 | +########################################################### |
| 145 | +# Run the two subgraphs in pipeline with pipeline executor. |
| 146 | +# --------------------------------------------------------- |
| 147 | +# Set 'USE_PIPELINE_EXECUTOR' as ON, and set USE_CUTLASS' as ON in cmake. |
| 148 | +from tvm.contrib import graph_executor, pipeline_executor, pipeline_executor_build |
| 149 | + |
| 150 | +######################################### |
| 151 | +# Create subgraph pipeline configuration. |
| 152 | +# Associate a subgraph module with a target. |
| 153 | +# Use CUTLASS BYOC to build the second subgraph module. |
| 154 | +mod0, mod1 = subgraphs[0], subgraphs[1] |
| 155 | +# Use cutlass as the codegen. |
| 156 | +mod1 = partition_for_cutlass(mod1) |
| 157 | +################################################# |
| 158 | +# Get the pipeline executor configuration object. |
| 159 | +pipe_config = pipeline_executor_build.PipelineConfig() |
| 160 | +########################################################################### |
| 161 | +# Set the compile target of the subgraph module. |
| 162 | +pipe_config[mod0].target = "llvm" |
| 163 | +pipe_config[mod0].dev = tvm.cpu(0) |
| 164 | +############################################################## |
| 165 | +# Set the compile target of the second subgraph module as cuda. |
| 166 | +pipe_config[mod1].target = "cuda" |
| 167 | +pipe_config[mod1].dev = tvm.device("cuda", 0) |
| 168 | +pipe_config[mod1].build_func = cutlass_build |
| 169 | +pipe_config[mod1].export_cc = "nvcc" |
| 170 | +# Create the pipeline by connecting the subgraph modules. |
| 171 | +# The global input will be forwarded to the input interface of the first module named mod0 |
| 172 | +pipe_config["input"]["data"].connect(pipe_config[mod0]["input"]["data"]) |
| 173 | +# The first output of mod0 will be forwarded to the input interface of mod1 |
| 174 | +pipe_config[mod0]["output"][0].connect(pipe_config[mod1]["input"]["data_n_0"]) |
| 175 | +# The first output of mod1 will be the first global output. |
| 176 | +pipe_config[mod1]["output"][0].connect(pipe_config["output"][0]) |
| 177 | +###################################### |
| 178 | +# The pipeline configuration as below. |
| 179 | +""" |
| 180 | +print(pipe_config) |
| 181 | + Inputs |
| 182 | + |data: mod0:data |
| 183 | +
|
| 184 | + output |
| 185 | + |output(0) : mod1.output(0) |
| 186 | +
|
| 187 | + connections |
| 188 | + |mod0.output(0)-> mod1.data_n_0 |
| 189 | +""" |
| 190 | + |
| 191 | +# sphinx_gallery_start_ignore |
| 192 | +from tvm import testing |
| 193 | + |
| 194 | +# testing.utils.install_request_hook(depth=3) |
| 195 | +# sphinx_gallery_end_ignore |
| 196 | +############################## |
| 197 | +# Build the pipeline executor. |
| 198 | +# ---------------------------- |
| 199 | +with tvm.transform.PassContext(opt_level=3): |
| 200 | + pipeline_mod_factory = pipeline_executor_build.build(pipe_config) |
| 201 | +############################################### |
| 202 | +# Export the parameter configuration to a file. |
| 203 | +directory_path = tvm.contrib.utils.tempdir().temp_dir |
| 204 | +os.makedirs(directory_path, exist_ok=True) |
| 205 | +config_file_name = pipeline_mod_factory.export_library(directory_path) |
| 206 | +################################################################ |
| 207 | +# Use the load function to create and initialize PipelineModule. |
| 208 | +# -------------------------------------------------------------- |
| 209 | +pipeline_module = pipeline_executor.PipelineModule.load_library(config_file_name) |
| 210 | + |
| 211 | +############################ |
| 212 | +# Run the pipeline executor. |
| 213 | +# -------------------------- |
| 214 | +# Allocate input data. |
| 215 | +data = np.random.uniform(-1, 1, size=data_shape).astype("float16") |
| 216 | +pipeline_module.set_input("data", tvm.nd.array(data)) |
| 217 | +########################################################################## |
| 218 | +# Run the two subgraph in the pipeline mode to get the output asynchronously |
| 219 | +# or synchronously. In the following example, it is synchronous. |
| 220 | +pipeline_module.run() |
| 221 | +outputs = pipeline_module.get_output() |
| 222 | +###################################### |
| 223 | +# Use graph_executor for verification. |
| 224 | +# ------------------------------------ |
| 225 | +# Run these two subgraphs in sequence with graph_executor to get the output. |
| 226 | +target = "llvm" |
| 227 | +dev0 = tvm.device(target, 0) |
| 228 | +lib0 = relay.build_module.build(mod0, target, params=params) |
| 229 | +module0 = runtime.GraphModule(lib0["default"](dev0)) |
| 230 | +cuda = tvm.target.Target("cuda", host=tvm.target.Target("llvm")) |
| 231 | +lib1 = relay.build_module.build(mod1, [cuda, cutlass], params=params) |
| 232 | +lib1 = finalize_modules(lib1, "compile.so", "./tmp") |
| 233 | + |
| 234 | +dev1 = tvm.device("cuda", 0) |
| 235 | + |
| 236 | +module1 = runtime.GraphModule(lib1["default"](dev1)) |
| 237 | + |
| 238 | +module0.set_input("data", data) |
| 239 | +module0.run() |
| 240 | +out_shape = (1, 16, img_size, img_size) |
| 241 | +out = module0.get_output(0, tvm.nd.empty(out_shape, "float16")) |
| 242 | +module1.set_input("data_n_0", out) |
| 243 | +module1.run() |
| 244 | +out_shape = (1, 1) |
| 245 | +out = module1.get_output(0, tvm.nd.empty(out_shape, "float16")) |
| 246 | +#################### |
| 247 | +# Verify the result. |
| 248 | +tvm.testing.assert_allclose(outputs[0].numpy(), out.numpy()) |
0 commit comments