From b7cd197fb98973ca59ca1b1884ffff0460d45042 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 14 Dec 2021 14:18:30 +0900 Subject: [PATCH 1/8] Add cutlass conv2d profiler commit 1c0bbb297da43dab75ad995afbcacd59e9fe4c87 Author: Masahiro Masuda Date: Sun Dec 12 18:29:03 2021 +0900 fix lint commit 463574ce087ca0444b23d4f47baf8066b8fbd3df Author: Masahiro Masuda Date: Sun Dec 12 17:28:38 2021 +0900 fixed conv2d check commit 588c5abe15abbf0339a7972d81b8235bc7460620 Author: Masahiro Masuda Date: Sun Dec 12 15:05:27 2021 +0900 update test commit a447b57ade7c99da4efd38e1bdfc213a64a80fd2 Author: Masahiro Masuda Date: Sun Dec 12 14:54:52 2021 +0900 speed up profiling by removing initialization commit 93cd039ba04dd80e887ec1a71f358cd86a1e5221 Author: Masahiro Masuda Date: Sun Dec 12 08:26:29 2021 +0900 fixed nhwc cudnn depthwise conv commit 6db71727f553ee2009c9d3feb2b019c24459f4d2 Author: Masahiro Masuda Date: Sat Dec 11 15:39:05 2021 +0900 add cache commit f7d17a116acd80c57dfa04aa86577fa30898908d Author: Masahiro Masuda Date: Sat Dec 11 15:05:38 2021 +0900 removed im2col profiling for conv2d commit b724f446d07030e45f30474a1e3c124f1541b496 Author: Masahiro Masuda Date: Fri Dec 10 22:57:54 2021 +0900 black commit fe4687b9f6d41eff66742cdde694680538a3d5e4 Author: Masahiro Masuda Date: Fri Dec 10 22:49:13 2021 +0900 fixed cmd arguement commit ab114f5c3e1a3086255caccd6c0f5df09d7c3755 Author: Masahiro Masuda Date: Fri Dec 10 22:22:19 2021 +0900 conv2d profiler working commit 49ee61f5583f73f0d9b2c18c1861c90fa028f64a Author: Masahiro Masuda Date: Fri Dec 10 20:26:15 2021 +0900 add conv2d profiler commit 49e2c8918a1a2632f60c700675f12f9d83adef24 Author: Masahiro Masuda Date: Sun Dec 12 08:03:36 2021 +0900 do not offload depthwise conv2d commit cd8367768229720020d81d597d751bfa95ec014e Author: Masahiro Masuda Date: Fri Dec 10 13:20:01 2021 +0900 lint fix commit 870823c6d54114896aa9db91eb72c132cdef762f Author: Masahiro Masuda Date: Fri Dec 10 12:54:38 2021 +0900 add comment on IC == 3 case commit 6b780db7f8059a9a58bfc257fbb9cf7cb60b72a2 Author: Masahiro Masuda Date: Fri Dec 10 12:48:33 2021 +0900 check align on N dim commit 308c4dac39f761ac157f032b9fe7028820980294 Author: Masahiro Masuda Date: Fri Dec 10 12:34:42 2021 +0900 fixed check functions for fused cases, run infer type before mergecomposite commit 8d6a1bfee26e9dfdeb1ab001aa89b43b9cd33d74 Author: Masahiro Masuda Date: Fri Dec 10 12:10:59 2021 +0900 test IC=3 convolution commit ffce47de724398222f672b220156b19c8f48a700 Author: Masahiro Masuda Date: Fri Dec 10 12:10:16 2021 +0900 use align1 kernel for unusual channel cases (IC = 3 etc) commit 6cdf205a3451b5fa7ea0cbff0228ae1da775573a Author: Masahiro Masuda Date: Fri Dec 10 12:06:56 2021 +0900 add dtype and layout check in parttern match commit 7743cc6dde442b45d66b4be320268ed24f352422 Author: Masahiro Masuda Date: Fri Dec 10 10:40:53 2021 +0900 add sm75 kernels to sm80 profilings commit efceccb994ec71bc4ba847c791fc4f696eb7f242 Author: Masahiro Masuda Date: Fri Dec 10 10:40:42 2021 +0900 skip legalize when batch size is dynamic commit 65fbc0a0813cb4e980e53240f34f3ba18754ab99 Author: Masahiro Masuda Date: Fri Dec 10 10:36:36 2021 +0900 bug fix in im2col encoding --- python/tvm/contrib/cutlass/build.py | 9 + python/tvm/contrib/cutlass/conv2d_profiler.py | 163 ++++++++++++++++++ python/tvm/contrib/cutlass/gen_conv2d.py | 81 +++++++-- python/tvm/contrib/cutlass/gen_gemm.py | 9 +- python/tvm/contrib/cutlass/gen_tensor_op.py | 14 +- python/tvm/relay/op/strategy/cuda.py | 5 +- tests/python/contrib/test_cutlass.py | 12 +- 7 files changed, 259 insertions(+), 34 deletions(-) create mode 100644 python/tvm/contrib/cutlass/conv2d_profiler.py diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py index c3a8fdc1ad8c..90b7f9320c7f 100644 --- a/python/tvm/contrib/cutlass/build.py +++ b/python/tvm/contrib/cutlass/build.py @@ -185,6 +185,9 @@ def handle_conv2d( d_shape, w_shape, out_shape, + padding, + strides, + dilation, out_dtype, profile_all, use_multiprocessing, @@ -198,6 +201,9 @@ def handle_conv2d( d_shape, w_shape, out_shape, + padding, + strides, + dilation, out_dtype, profile_all=profile_all, use_multiprocessing=use_multiprocessing, @@ -279,6 +285,9 @@ def tune_cutlass_kernels(mod, sm, profile_all=True, use_multiprocessing=False, t arg0_shape, arg1_shape, annotator.signature["ret_shape"], + annotator.op_attrs.padding, + annotator.op_attrs.strides, + annotator.op_attrs.dilation, out_dtype, profile_all, use_multiprocessing, diff --git a/python/tvm/contrib/cutlass/conv2d_profiler.py b/python/tvm/contrib/cutlass/conv2d_profiler.py new file mode 100644 index 000000000000..e4ae03a4e3c7 --- /dev/null +++ b/python/tvm/contrib/cutlass/conv2d_profiler.py @@ -0,0 +1,163 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you 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. +# pylint: disable=import-outside-toplevel, invalid-name +"""Instantiate a C++ source for profiling CUTLASS kernels.""" + + +class Conv2dProfilerEmitter(object): + """Emit a C++ source for profiling CUTLASS kernels.""" + + def __init__(self): + from jinja2 import Template + + self.template = Template( + """ +#include +#include "cutlass/cutlass.h" +#include "cutlass/conv/kernel/default_conv2d_fprop.h" +#include "cutlass/conv/device/implicit_gemm_convolution.h" +#include "cutlass/util/command_line.h" +#include "cutlass/util/host_tensor.h" +#include "cutlass/util/reference/host/tensor_fill.h" + +#define CUTLASS_CHECK(status) \ + { \ + cutlass::Status error = status; \ + if (error != cutlass::Status::kSuccess) { \ + std::cerr << "Got cutlass error: " << cutlassGetStatusString(error) << " at: " << __LINE__ \ + << std::endl; \ + exit(EXIT_FAILURE); \ + } \ + } + +{{OperatorDef}} +using ImplicitGemm = cutlass::conv::device::ImplicitGemmConvolution<{{OperatorName}}>; + +struct Options { + cutlass::Tensor4DCoord input_size; + cutlass::Tensor4DCoord filter_size; + cutlass::Tensor4DCoord padding; + cutlass::MatrixCoord conv_stride; + cutlass::MatrixCoord dilation; + + void parse(int argc, char const **args) { + cutlass::CommandLine cmd(argc, args); + cmd.get_cmd_line_argument("n", input_size.n()); + cmd.get_cmd_line_argument("h", input_size.h()); + cmd.get_cmd_line_argument("w", input_size.w()); + cmd.get_cmd_line_argument("c", input_size.c()); + cmd.get_cmd_line_argument("k", filter_size.n()); + cmd.get_cmd_line_argument("r", filter_size.h()); + cmd.get_cmd_line_argument("s", filter_size.w()); + int pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w; + cmd.get_cmd_line_argument("pad_h", pad_h); + cmd.get_cmd_line_argument("pad_w", pad_w); + cmd.get_cmd_line_argument("stride_h", stride_h); + cmd.get_cmd_line_argument("stride_w", stride_w); + cmd.get_cmd_line_argument("dilation_h", dilation_h); + cmd.get_cmd_line_argument("dilation_w", dilation_w); + filter_size.c() = input_size.c(); + padding = {pad_h, pad_h, pad_w, pad_w}; + conv_stride = {stride_h, stride_w}; + dilation = {dilation_h, dilation_w}; + } + + cutlass::Tensor4DCoord output_size() const { + auto dilated_h = (filter_size.h() - 1) * dilation.row() + 1; + auto dilated_w = (filter_size.w() - 1) * dilation.column() + 1; + auto h = (input_size.h() + padding.n() + padding.h() - dilated_h) / conv_stride.row() + 1; + auto w = (input_size.w() + padding.w() + padding.c() - dilated_w) / conv_stride.column() + 1; + return cutlass::Tensor4DCoord(input_size.n(), h, w, filter_size.n()); + } +}; + +double profile_convolution(Options const &options) { + using ElementOutput = typename ImplicitGemm::ElementC; + using ElementInputA = typename ImplicitGemm::ElementA; + using ElementInputB = typename ImplicitGemm::ElementB; + auto oshape = options.output_size(); + cutlass::HostTensor tensor_a(options.input_size); + cutlass::HostTensor tensor_b(options.filter_size); + cutlass::HostTensor tensor_c(oshape); + cutlass::HostTensor tensor_ref_c(oshape); + + cutlass::conv::Conv2dProblemSize problem_size( + options.input_size, + options.filter_size, + options.padding, + options.conv_stride, + options.dilation, + options.output_size(), + cutlass::conv::Mode::kCrossCorrelation, + 1 + ); + + using ElementComputeEpilogue = typename ImplicitGemm::ElementCompute; + typename ImplicitGemm::Arguments arguments{ + problem_size, + tensor_a.device_ref(), + tensor_b.device_ref(), + tensor_c.device_ref(), + tensor_c.device_ref(), + {ElementComputeEpilogue(1), ElementComputeEpilogue(0)}, + }; + + ImplicitGemm implicit_gemm_op; + size_t workspace_size = implicit_gemm_op.get_workspace_size(arguments); + cutlass::device_memory::allocation workspace(workspace_size); + auto status = implicit_gemm_op.can_implement(arguments); + CUTLASS_CHECK(status); + + status = implicit_gemm_op.initialize(arguments, workspace.get()); + CUTLASS_CHECK(status); + status = implicit_gemm_op(); + CUTLASS_CHECK(status); + + cudaEvent_t events[2]; + for (auto & event : events) { + cudaEventCreate(&event); + } + cudaEventRecord(events[0]); + + for (int iteration = 0; iteration < 100; ++iteration) { + auto status = implicit_gemm_op(); + CUTLASS_CHECK(status); + } + + cudaEventRecord(events[1]); + cudaEventSynchronize(events[1]); + float runtime_ms = 0; + cudaEventElapsedTime(&runtime_ms, events[0], events[1]); + + for (auto event : events) { + (void)cudaEventDestroy(event); + } + return double(runtime_ms) / 100.0; +} + +int main(int argc, char const **args) { + Options options; + options.parse(argc, args); + std::cout << profile_convolution(options) << std::endl; + return 0; +} +""" + ) + + def emit(self, op_def, op_name): + src = self.template.render(OperatorDef=op_def, OperatorName=op_name) + return src diff --git a/python/tvm/contrib/cutlass/gen_conv2d.py b/python/tvm/contrib/cutlass/gen_conv2d.py index 5a616c9b6e02..a5beb357b309 100644 --- a/python/tvm/contrib/cutlass/gen_conv2d.py +++ b/python/tvm/contrib/cutlass/gen_conv2d.py @@ -16,8 +16,14 @@ # under the License. # pylint: disable=invalid-name """Conv2d kernel generator and profiler for CUTLASS.""" +import re from .conv2d_operation import Conv2dOperation, EmitConv2dInstance from .gen_gemm import CutlassGemmProfiler +from .conv2d_profiler import Conv2dProfilerEmitter +from .gen_tensor_op import ( + ProfilerEngine, + GENERATOR_FUNC_TABLE, +) from .library import ( EpilogueFunctor, SwizzlingFunctor, @@ -39,6 +45,7 @@ def create_conv2d_operator( ret = [] kernel_emitter = EmitConv2dInstance() + profiler_emitter = Conv2dProfilerEmitter() element_a, element_b, element_c, element_epilogue = data_type iterator_algorithms = [IteratorAlgorithm.Optimized] @@ -75,6 +82,7 @@ def create_conv2d_operator( # TODO(masahi): Add profiler source here op_entry["opdef"] = kernel_emitter.emit(op) op_entry["op"] = op + op_entry["src"] = profiler_emitter.emit(op_entry["opdef"], op.procedural_name()) op_entry["name"] = op.procedural_name() op_entry["runtime"] = 9999999 @@ -113,6 +121,9 @@ class CutlassConv2DProfiler: def __init__(self, sm, cutlass_path, binary_path): self.gemm_profiler = CutlassGemmProfiler(sm, cutlass_path, binary_path) self.sm = sm + assert sm in GENERATOR_FUNC_TABLE, "sm%d not supported yet." % sm + self.engine = ProfilerEngine(sm, cutlass_path, binary_path) + self.cache = {} def get_default(self, out_dtype): gemm_profile_result = self.gemm_profiler.get_default(out_dtype) @@ -121,27 +132,71 @@ def get_default(self, out_dtype): data_type = gemm_profile_result["data_type"] return create_conv2d_operator([tile_description], data_type, [alignment])[0] + def check_align(self, op_name, C, K): + """Filter out kernels that cannot be supported.""" + aligns = re.findall(r"align[1|2|4|8]", op_name) + assert len(aligns) == 1 + align = int(aligns[0][-1]) + return all([dim % align == 0 for dim in [C, K]]) + def profile( - self, d_shape, w_shape, out_shape, out_dtype, profile_all=True, use_multiprocessing=False + self, + d_shape, + w_shape, + out_shape, + padding, + stride, + dilation, + out_dtype, + profile_all=True, + use_multiprocessing=False, ): """Profile and select the best kernel from candidate kernels. If profile_all is False, return immediately after the first applicable kernel is found. If use_multiprocessing is True, compile all profiler executables in parallel. """ - B, _, _, IC = d_shape + N, H, W, IC = d_shape OC, R, S, _ = w_shape - _, P, Q, _ = out_shape + workload = ( + N, + H, + W, + IC, + OC, + R, + S, + padding[0], + padding[1], + stride[0], + stride[1], + dilation[0], + dilation[1], + ) - M = B * P * Q - N = OC - K = R * S * IC + if workload in self.cache: + return self.cache[workload] - gemm_profile_result = self.gemm_profiler.profile( - M, N, K, out_dtype, profile_all=profile_all, use_multiprocessing=use_multiprocessing - ) + ops = GENERATOR_FUNC_TABLE[self.sm](out_dtype, op_creator=create_conv2d_operator) + ops = list(filter(lambda op: self.check_align(op["name"], IC, OC), ops)) - tile_description = gemm_profile_result["tile_description"] - alignment = gemm_profile_result["alignment"] - data_type = gemm_profile_result["data_type"] + for op in ops: + op["runtime"] = -1 - return create_conv2d_operator([tile_description], data_type, [alignment])[0] + if profile_all: + self.engine.compile_all(ops, use_multiprocessing) + + args = ( + "--n=%d --h=%d --w=%d --c=%d --k=%d --r=%d --s=%d --pad_h=%d --pad_w=%d " + "--stride_h=%d --stride_w=%d --dilation_h=%d --dilation_w=%d" + ) % workload + + for op in ops: + out = self.engine.evaluate(op, args.split(" ")) + op["runtime"] = out + if out > 0 and profile_all is False: + break + + valid_ops = filter(lambda op: op["runtime"] > 0, ops) + output = sorted(valid_ops, key=lambda i: i["runtime"]) + self.cache[workload] = output[0] + return output[0] diff --git a/python/tvm/contrib/cutlass/gen_gemm.py b/python/tvm/contrib/cutlass/gen_gemm.py index 58d690f8191b..7287d2471817 100644 --- a/python/tvm/contrib/cutlass/gen_gemm.py +++ b/python/tvm/contrib/cutlass/gen_gemm.py @@ -22,8 +22,7 @@ from .gemm_profiler import GemmProfilerEmitter from .gen_tensor_op import ( ProfilerEngine, - generate_sm75_tensor_op_1688, - generate_sm80_tensor_op_16816, + GENERATOR_FUNC_TABLE, ) from .library import ( EpilogueFunctor, @@ -132,12 +131,6 @@ def create_gemm_operator( return ret -GENERATOR_FUNC_TABLE = { - 75: generate_sm75_tensor_op_1688, - 80: generate_sm80_tensor_op_16816, -} - - # TODO(masahi): A sensible way to pick reasonable default kernels DEFAULT_KERNELS = { 75: { diff --git a/python/tvm/contrib/cutlass/gen_tensor_op.py b/python/tvm/contrib/cutlass/gen_tensor_op.py index cc228737cefc..01d340912f16 100644 --- a/python/tvm/contrib/cutlass/gen_tensor_op.py +++ b/python/tvm/contrib/cutlass/gen_tensor_op.py @@ -159,6 +159,12 @@ def get_tile_descriptions(math_inst): return sm75_kernels + sm80_kernels +GENERATOR_FUNC_TABLE = { + 75: generate_sm75_tensor_op_1688, + 80: generate_sm80_tensor_op_16816, +} + + class ProfilerEngine: """Compile and run a given profiler executable.""" @@ -204,12 +210,8 @@ def evaluate(self, op, args): if not os.path.exists(opath): self._compile(op) cmd = [opath] - if args is not None: - cmd.append(str(args[0])) - cmd.append(str(args[1])) - cmd.append(str(args[2])) - if len(args) > 3: - cmd.append(str(args[3])) + for arg in args: + cmd.append(str(arg)) try: sp = subprocess.run(cmd, capture_output=True, check=True) rt = float(sp.stdout) diff --git a/python/tvm/relay/op/strategy/cuda.py b/python/tvm/relay/op/strategy/cuda.py index 7bc04b45774c..307720fe5883 100644 --- a/python/tvm/relay/op/strategy/cuda.py +++ b/python/tvm/relay/op/strategy/cuda.py @@ -324,7 +324,10 @@ def conv2d_strategy_cuda(attrs, inputs, out_type, target): plevel=25, ) - elif is_depthwise_conv2d(data.shape, layout, kernel.shape, kernel_layout, groups): + elif ( + is_depthwise_conv2d(data.shape, layout, kernel.shape, kernel_layout, groups) + and "cudnn" not in target.libs + ): if layout == "NCHW": assert kernel_layout == "OIHW" strategy.add_implementation( diff --git a/tests/python/contrib/test_cutlass.py b/tests/python/contrib/test_cutlass.py index 585b42a21425..672dda6e4bac 100644 --- a/tests/python/contrib/test_cutlass.py +++ b/tests/python/contrib/test_cutlass.py @@ -85,7 +85,7 @@ def get_dense_bias(M, N, K, out_dtype="float16"): def get_dense_bias_relu(M, N, K, out_dtype="float16"): - return relay.nn.relu(get_dense_bias(M, N, K, out_dtype="float16")) + return relay.nn.relu(get_dense_bias(M, N, K, out_dtype=out_dtype)) def get_dense_bias_gelu(M, N, K, out_dtype="float16"): @@ -110,7 +110,7 @@ def get_batch_matmul(batch, M, N, K, out_dtype="float16"): return get_batch_matmul_with_shape((batch, M, K), (batch, N, K), out_dtype="float16") -def get_conv2d_nchw(d_shape, w_shape): +def get_conv2d_nchw(d_shape, w_shape, padding, out_dtype="float16"): data = relay.var("data", shape=d_shape, dtype="float16") weight = relay.var("weight", shape=w_shape, dtype="float16") out_channel = w_shape[0] @@ -118,10 +118,10 @@ def get_conv2d_nchw(d_shape, w_shape): relay.nn.conv2d( data=data, weight=weight, - kernel_size=(3, 3), + kernel_size=w_shape[2:], channels=out_channel, - padding=(1, 1), - out_dtype="float16", + padding=padding, + out_dtype=out_dtype, ) ) @@ -129,7 +129,7 @@ def get_conv2d_nchw(d_shape, w_shape): def profile_and_build(mod, params, sm, tmp_dir="./tmp", lib_path="compile.so"): mod = partition_for_cutlass(mod) mod, num_cutlass_partition = tune_cutlass_kernels( - mod, sm, profile_all=False, use_multiprocessing=False, tmp_dir=tmp_dir + mod, sm, profile_all=True, use_multiprocessing=True, tmp_dir=tmp_dir ) with tvm.transform.PassContext(opt_level=3): lib = relay.build(mod, target="cuda", params=params) From 8e47ec89e46ebe63c76ca1f42a2e852008638681 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 14 Dec 2021 14:32:38 +0900 Subject: [PATCH 2/8] minor fix --- python/tvm/contrib/cutlass/build.py | 1 - python/tvm/contrib/cutlass/gen_conv2d.py | 2 -- tests/python/contrib/test_cutlass.py | 10 ++++++---- 3 files changed, 6 insertions(+), 7 deletions(-) diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py index 90b7f9320c7f..cf9888d17c3b 100644 --- a/python/tvm/contrib/cutlass/build.py +++ b/python/tvm/contrib/cutlass/build.py @@ -200,7 +200,6 @@ def handle_conv2d( out = cutlass_profiler.profile( d_shape, w_shape, - out_shape, padding, strides, dilation, diff --git a/python/tvm/contrib/cutlass/gen_conv2d.py b/python/tvm/contrib/cutlass/gen_conv2d.py index a5beb357b309..0696003951a9 100644 --- a/python/tvm/contrib/cutlass/gen_conv2d.py +++ b/python/tvm/contrib/cutlass/gen_conv2d.py @@ -79,7 +79,6 @@ def create_conv2d_operator( swizzling_functor_, ) - # TODO(masahi): Add profiler source here op_entry["opdef"] = kernel_emitter.emit(op) op_entry["op"] = op op_entry["src"] = profiler_emitter.emit(op_entry["opdef"], op.procedural_name()) @@ -143,7 +142,6 @@ def profile( self, d_shape, w_shape, - out_shape, padding, stride, dilation, diff --git a/tests/python/contrib/test_cutlass.py b/tests/python/contrib/test_cutlass.py index 672dda6e4bac..fee84d252081 100644 --- a/tests/python/contrib/test_cutlass.py +++ b/tests/python/contrib/test_cutlass.py @@ -129,7 +129,7 @@ def get_conv2d_nchw(d_shape, w_shape, padding, out_dtype="float16"): def profile_and_build(mod, params, sm, tmp_dir="./tmp", lib_path="compile.so"): mod = partition_for_cutlass(mod) mod, num_cutlass_partition = tune_cutlass_kernels( - mod, sm, profile_all=True, use_multiprocessing=True, tmp_dir=tmp_dir + mod, sm, profile_all=False, use_multiprocessing=False, tmp_dir=tmp_dir ) with tvm.transform.PassContext(opt_level=3): lib = relay.build(mod, target="cuda", params=params) @@ -376,7 +376,8 @@ def test_conv2d(): for IC in [3, 16]: d_shape = (16, IC, 32, 32) w_shape = (32, IC, 3, 3) - mod_nchw = get_conv2d_nchw(d_shape, w_shape) + padding = (1, 1) + mod_nchw = get_conv2d_nchw(d_shape, w_shape, padding) verify_conv2d( mod_nchw, @@ -392,10 +393,11 @@ def test_conv2d(): d_shape = (16, 16, 32, 32) w_shape = (32, 16, 3, 3) + padding = (1, 1) dyn_batch_shape = (relay.Any(),) + d_shape[1:] - mod_nchw = get_conv2d_nchw(d_shape, w_shape) - mod_dyn = get_conv2d_nchw(dyn_batch_shape, w_shape) + mod_nchw = get_conv2d_nchw(d_shape, w_shape, padding) + mod_dyn = get_conv2d_nchw(dyn_batch_shape, w_shape, padding) verify_conv2d( mod_dyn, mod_nchw, d_shape, w_shape, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False From 3e634a44d64ff17e03c231a893b47a5f51c439e6 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Tue, 14 Dec 2021 15:21:49 +0900 Subject: [PATCH 3/8] lint fix --- python/tvm/contrib/cutlass/build.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/python/tvm/contrib/cutlass/build.py b/python/tvm/contrib/cutlass/build.py index cf9888d17c3b..32a2ebaa2842 100644 --- a/python/tvm/contrib/cutlass/build.py +++ b/python/tvm/contrib/cutlass/build.py @@ -184,7 +184,6 @@ def handle_conv2d( op_type, d_shape, w_shape, - out_shape, padding, strides, dilation, @@ -283,7 +282,6 @@ def tune_cutlass_kernels(mod, sm, profile_all=True, use_multiprocessing=False, t op_type, arg0_shape, arg1_shape, - annotator.signature["ret_shape"], annotator.op_attrs.padding, annotator.op_attrs.strides, annotator.op_attrs.dilation, From 56ef28d010ddc686657912a5a587cd189a8c30a0 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 15 Dec 2021 03:59:26 +0900 Subject: [PATCH 4/8] allow autotvm NCHW depthwise conv2d schedule even if -libs=cudnn --- python/tvm/relay/op/strategy/cuda.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/python/tvm/relay/op/strategy/cuda.py b/python/tvm/relay/op/strategy/cuda.py index 307720fe5883..607b2a59045a 100644 --- a/python/tvm/relay/op/strategy/cuda.py +++ b/python/tvm/relay/op/strategy/cuda.py @@ -324,10 +324,9 @@ def conv2d_strategy_cuda(attrs, inputs, out_type, target): plevel=25, ) - elif ( - is_depthwise_conv2d(data.shape, layout, kernel.shape, kernel_layout, groups) - and "cudnn" not in target.libs - ): + elif is_depthwise_conv2d(data.shape, layout, kernel.shape, kernel_layout, groups) and ( + layout == "NCHW" or "cudnn" not in target.libs + ): # cuDNN requires a different kernel layout for NHWC inputs. if layout == "NCHW": assert kernel_layout == "OIHW" strategy.add_implementation( From 2c2e76cdf6d1de1f731eecac576cb8ba267490d0 Mon Sep 17 00:00:00 2001 From: masahi Date: Wed, 15 Dec 2021 05:45:03 +0900 Subject: [PATCH 5/8] Update python/tvm/contrib/cutlass/gen_conv2d.py Co-authored-by: Cody Yu --- python/tvm/contrib/cutlass/gen_conv2d.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/contrib/cutlass/gen_conv2d.py b/python/tvm/contrib/cutlass/gen_conv2d.py index 0696003951a9..ff7d571de9c6 100644 --- a/python/tvm/contrib/cutlass/gen_conv2d.py +++ b/python/tvm/contrib/cutlass/gen_conv2d.py @@ -191,7 +191,7 @@ def profile( for op in ops: out = self.engine.evaluate(op, args.split(" ")) op["runtime"] = out - if out > 0 and profile_all is False: + if out > 0 and not profile_all: break valid_ops = filter(lambda op: op["runtime"] > 0, ops) From a5dbea58a91efb2be1f512a117644a7a9c27a2ea Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 15 Dec 2021 05:57:45 +0900 Subject: [PATCH 6/8] simplify processing profiler outputs --- python/tvm/contrib/cutlass/gen_conv2d.py | 7 +++---- python/tvm/contrib/cutlass/gen_gemm.py | 9 ++++----- python/tvm/contrib/cutlass/gen_tensor_op.py | 2 +- 3 files changed, 8 insertions(+), 10 deletions(-) diff --git a/python/tvm/contrib/cutlass/gen_conv2d.py b/python/tvm/contrib/cutlass/gen_conv2d.py index ff7d571de9c6..c5205ac8c531 100644 --- a/python/tvm/contrib/cutlass/gen_conv2d.py +++ b/python/tvm/contrib/cutlass/gen_conv2d.py @@ -194,7 +194,6 @@ def profile( if out > 0 and not profile_all: break - valid_ops = filter(lambda op: op["runtime"] > 0, ops) - output = sorted(valid_ops, key=lambda i: i["runtime"]) - self.cache[workload] = output[0] - return output[0] + output = min(ops, key=lambda i: i["runtime"]) + self.cache[workload] = output + return output diff --git a/python/tvm/contrib/cutlass/gen_gemm.py b/python/tvm/contrib/cutlass/gen_gemm.py index 7287d2471817..3660a25b1488 100644 --- a/python/tvm/contrib/cutlass/gen_gemm.py +++ b/python/tvm/contrib/cutlass/gen_gemm.py @@ -201,10 +201,9 @@ def profile( for op in ops: out = self.engine.evaluate(op, [M, N, K]) op["runtime"] = out - if out > 0 and profile_all is False: + if out > 0 and not profile_all: break - valid_ops = filter(lambda op: op["runtime"] > 0, ops) - output = sorted(valid_ops, key=lambda i: i["runtime"]) - self.cache[(M, N, K)] = output[0] - return output[0] + output = min(ops, key=lambda i: i["runtime"]) + self.cache[(M, N, K)] = output + return output diff --git a/python/tvm/contrib/cutlass/gen_tensor_op.py b/python/tvm/contrib/cutlass/gen_tensor_op.py index 01d340912f16..9ccde37bfe91 100644 --- a/python/tvm/contrib/cutlass/gen_tensor_op.py +++ b/python/tvm/contrib/cutlass/gen_tensor_op.py @@ -217,5 +217,5 @@ def evaluate(self, op, args): rt = float(sp.stdout) logger.info("%s, %f", op_name, rt) except subprocess.CalledProcessError: - rt = -1 + rt = float("inf") return rt From 28101ff89e9d08169718bd9cd49e8361d8451ad5 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 15 Dec 2021 06:07:09 +0900 Subject: [PATCH 7/8] more simplify --- python/tvm/contrib/cutlass/gen_conv2d.py | 6 ++---- python/tvm/contrib/cutlass/gen_gemm.py | 6 ++---- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/python/tvm/contrib/cutlass/gen_conv2d.py b/python/tvm/contrib/cutlass/gen_conv2d.py index c5205ac8c531..e20f533f5598 100644 --- a/python/tvm/contrib/cutlass/gen_conv2d.py +++ b/python/tvm/contrib/cutlass/gen_conv2d.py @@ -177,9 +177,6 @@ def profile( ops = GENERATOR_FUNC_TABLE[self.sm](out_dtype, op_creator=create_conv2d_operator) ops = list(filter(lambda op: self.check_align(op["name"], IC, OC), ops)) - for op in ops: - op["runtime"] = -1 - if profile_all: self.engine.compile_all(ops, use_multiprocessing) @@ -192,7 +189,8 @@ def profile( out = self.engine.evaluate(op, args.split(" ")) op["runtime"] = out if out > 0 and not profile_all: - break + self.cache[workload] = op + return op output = min(ops, key=lambda i: i["runtime"]) self.cache[workload] = output diff --git a/python/tvm/contrib/cutlass/gen_gemm.py b/python/tvm/contrib/cutlass/gen_gemm.py index 3660a25b1488..9a521decab7a 100644 --- a/python/tvm/contrib/cutlass/gen_gemm.py +++ b/python/tvm/contrib/cutlass/gen_gemm.py @@ -192,9 +192,6 @@ def profile( ) ops = list(filter(lambda op: self.check_align(op["name"], M, N, K), ops)) - for op in ops: - op["runtime"] = -1 - if profile_all: self.engine.compile_all(ops, use_multiprocessing) @@ -202,7 +199,8 @@ def profile( out = self.engine.evaluate(op, [M, N, K]) op["runtime"] = out if out > 0 and not profile_all: - break + self.cache[(M, N, K)] = op + return op output = min(ops, key=lambda i: i["runtime"]) self.cache[(M, N, K)] = output From 00c161ddf3886f22343e2d9bc6f2b9a6fe1a6f32 Mon Sep 17 00:00:00 2001 From: Masahiro Masuda Date: Wed, 15 Dec 2021 06:53:35 +0900 Subject: [PATCH 8/8] fix runtime check --- python/tvm/contrib/cutlass/gen_conv2d.py | 2 +- python/tvm/contrib/cutlass/gen_gemm.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/contrib/cutlass/gen_conv2d.py b/python/tvm/contrib/cutlass/gen_conv2d.py index e20f533f5598..b0d0566d6fab 100644 --- a/python/tvm/contrib/cutlass/gen_conv2d.py +++ b/python/tvm/contrib/cutlass/gen_conv2d.py @@ -188,7 +188,7 @@ def profile( for op in ops: out = self.engine.evaluate(op, args.split(" ")) op["runtime"] = out - if out > 0 and not profile_all: + if out < float("inf") and not profile_all: self.cache[workload] = op return op diff --git a/python/tvm/contrib/cutlass/gen_gemm.py b/python/tvm/contrib/cutlass/gen_gemm.py index 9a521decab7a..c171c5e23a89 100644 --- a/python/tvm/contrib/cutlass/gen_gemm.py +++ b/python/tvm/contrib/cutlass/gen_gemm.py @@ -198,7 +198,7 @@ def profile( for op in ops: out = self.engine.evaluate(op, [M, N, K]) op["runtime"] = out - if out > 0 and not profile_all: + if out < float("inf") and not profile_all: self.cache[(M, N, K)] = op return op