You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by GitBox <gi...@apache.org> on 2021/11/26 10:27:49 UTC
[GitHub] [tvm] masahi opened a new pull request #9595: [CUTLASS] Initial conv2d support
masahi opened a new pull request #9595:
URL: https://github.com/apache/tvm/pull/9595
Thanks for contributing to TVM! Please refer to guideline https://tvm.apache.org/docs/contribute/ for useful information and tips. After the pull request is submitted, please request code reviews from [Reviewers](https://github.com/apache/incubator-tvm/blob/master/CONTRIBUTORS.md#reviewers) by @ them in the pull request thread.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on pull request #9595: [CUTLASS] Initial conv2d support
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #9595:
URL: https://github.com/apache/tvm/pull/9595#issuecomment-984151052
@comaniac good to go?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] comaniac commented on a change in pull request #9595: [CUTLASS] Initial conv2d support
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9595:
URL: https://github.com/apache/tvm/pull/9595#discussion_r759694173
##########
File path: python/tvm/contrib/cutlass/conv2d_operation.py
##########
@@ -0,0 +1,240 @@
+# 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=invalid-name, unused-wildcard-import, wildcard-import
+"""Generator for CUTLASS Conv2D kernels."""
+from .library import *
+
+
+class Conv2dOperation:
Review comment:
It might be better to have a base classes (Operation/Emit/Profiler) for both Conv2d and Gemm to enforce consistent APIs.
##########
File path: python/tvm/contrib/cutlass/gen_conv2d.py
##########
@@ -0,0 +1,132 @@
+# 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=invalid-name
+"""Conv2d kernel generator and profiler for CUTLASS."""
+from .conv2d_operation import Conv2dOperation, EmitConv2dInstance
+from .gen_gemm import CutlassGemmProfiler
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ LayoutType,
+ ConvKind,
+ StrideSupport,
+ IteratorAlgorithm,
+)
+
+
+def create_conv2d_operator(
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ swizzling_functor=SwizzlingFunctor.Identity4,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+
+ kernel_emitter = EmitConv2dInstance()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+ iterator_algorithms = [IteratorAlgorithm.Optimized]
+
+ layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC)
+ for tile in tile_descriptions:
+ for alignment in alignment_constraints:
+
+ alignment_c = min(8, alignment)
+
+ A = TensorDescription(element_a, layout[0], alignment)
+ B = TensorDescription(element_b, layout[1], alignment)
+ C = TensorDescription(element_c, layout[2], alignment_c)
+
+ swizzling_functor_ = swizzling_functor
+
+ for iterator_algorithm in iterator_algorithms:
+ op_entry = {}
+
+ for epilogue, opdef in zip(
+ [
+ EpilogueFunctor.LinearCombination,
+ EpilogueFunctor.LinearCombinationBias,
+ EpilogueFunctor.LinearCombinationRelu,
+ ],
+ ["opdef", "opdef_bias", "opdef_bias_relu"],
+ ):
+ op = Conv2dOperation(
+ ConvKind.Fprop,
+ iterator_algorithm,
+ tile.minimum_compute_capability,
+ tile,
+ A,
+ B,
+ C,
+ element_epilogue,
+ StrideSupport.Strided,
+ epilogue,
+ swizzling_functor_,
+ )
+
+ op_entry[opdef] = kernel_emitter.emit(op)
+
+ if epilogue == EpilogueFunctor.LinearCombination:
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["runtime"] = 9999999
Review comment:
What's this for? Maybe also adding comments for clarification.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #9595: [CUTLASS] Initial conv2d support
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9595:
URL: https://github.com/apache/tvm/pull/9595#discussion_r759712433
##########
File path: python/tvm/contrib/cutlass/gen_conv2d.py
##########
@@ -0,0 +1,132 @@
+# 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=invalid-name
+"""Conv2d kernel generator and profiler for CUTLASS."""
+from .conv2d_operation import Conv2dOperation, EmitConv2dInstance
+from .gen_gemm import CutlassGemmProfiler
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ LayoutType,
+ ConvKind,
+ StrideSupport,
+ IteratorAlgorithm,
+)
+
+
+def create_conv2d_operator(
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ swizzling_functor=SwizzlingFunctor.Identity4,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+
+ kernel_emitter = EmitConv2dInstance()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+ iterator_algorithms = [IteratorAlgorithm.Optimized]
+
+ layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC)
+ for tile in tile_descriptions:
+ for alignment in alignment_constraints:
+
+ alignment_c = min(8, alignment)
+
+ A = TensorDescription(element_a, layout[0], alignment)
+ B = TensorDescription(element_b, layout[1], alignment)
+ C = TensorDescription(element_c, layout[2], alignment_c)
+
+ swizzling_functor_ = swizzling_functor
+
+ for iterator_algorithm in iterator_algorithms:
+ op_entry = {}
+
+ for epilogue, opdef in zip(
+ [
+ EpilogueFunctor.LinearCombination,
+ EpilogueFunctor.LinearCombinationBias,
+ EpilogueFunctor.LinearCombinationRelu,
+ ],
+ ["opdef", "opdef_bias", "opdef_bias_relu"],
+ ):
+ op = Conv2dOperation(
+ ConvKind.Fprop,
+ iterator_algorithm,
+ tile.minimum_compute_capability,
+ tile,
+ A,
+ B,
+ C,
+ element_epilogue,
+ StrideSupport.Strided,
+ epilogue,
+ swizzling_functor_,
+ )
+
+ op_entry[opdef] = kernel_emitter.emit(op)
+
+ if epilogue == EpilogueFunctor.LinearCombination:
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["runtime"] = 9999999
Review comment:
This corresponds to the gemm generator counterpart in https://github.com/apache/tvm/blob/adf560ebed8465c22bf58f406d0a8d20663cdd1d/python/tvm/contrib/cutlass/gen_gemm.py#L109-L127
In addition to creating `opdef`, `opdef_bias` etc, we also need to set `op`, `name`, `runtime` etc. I tried to simplify that code and this is what I came up with.
I'll rewrite this code to make it easier to understand (by pulling the non-activation case, `EpilogueFunctor.LinearCombination`, out of the loop).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] comaniac commented on pull request #9595: [CUTLASS] Initial conv2d support
Posted by GitBox <gi...@apache.org>.
comaniac commented on pull request #9595:
URL: https://github.com/apache/tvm/pull/9595#issuecomment-984212313
Thanks @masahi. @Laurawly @ZihengJiang please feel free to continue the discussion here or at the forum for the potential issue.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] comaniac merged pull request #9595: [CUTLASS] Initial conv2d support
Posted by GitBox <gi...@apache.org>.
comaniac merged pull request #9595:
URL: https://github.com/apache/tvm/pull/9595
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #9595: [CUTLASS] Initial conv2d support
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9595:
URL: https://github.com/apache/tvm/pull/9595#discussion_r760170140
##########
File path: python/tvm/contrib/cutlass/gen_conv2d.py
##########
@@ -0,0 +1,132 @@
+# 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=invalid-name
+"""Conv2d kernel generator and profiler for CUTLASS."""
+from .conv2d_operation import Conv2dOperation, EmitConv2dInstance
+from .gen_gemm import CutlassGemmProfiler
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ LayoutType,
+ ConvKind,
+ StrideSupport,
+ IteratorAlgorithm,
+)
+
+
+def create_conv2d_operator(
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ swizzling_functor=SwizzlingFunctor.Identity4,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+
+ kernel_emitter = EmitConv2dInstance()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+ iterator_algorithms = [IteratorAlgorithm.Optimized]
+
+ layout = (LayoutType.TensorNHWC, LayoutType.TensorNHWC, LayoutType.TensorNHWC)
+ for tile in tile_descriptions:
+ for alignment in alignment_constraints:
+
+ alignment_c = min(8, alignment)
+
+ A = TensorDescription(element_a, layout[0], alignment)
+ B = TensorDescription(element_b, layout[1], alignment)
+ C = TensorDescription(element_c, layout[2], alignment_c)
+
+ swizzling_functor_ = swizzling_functor
+
+ for iterator_algorithm in iterator_algorithms:
+ op_entry = {}
+
+ for epilogue, opdef in zip(
+ [
+ EpilogueFunctor.LinearCombination,
+ EpilogueFunctor.LinearCombinationBias,
+ EpilogueFunctor.LinearCombinationRelu,
+ ],
+ ["opdef", "opdef_bias", "opdef_bias_relu"],
+ ):
+ op = Conv2dOperation(
+ ConvKind.Fprop,
+ iterator_algorithm,
+ tile.minimum_compute_capability,
+ tile,
+ A,
+ B,
+ C,
+ element_epilogue,
+ StrideSupport.Strided,
+ epilogue,
+ swizzling_functor_,
+ )
+
+ op_entry[opdef] = kernel_emitter.emit(op)
+
+ if epilogue == EpilogueFunctor.LinearCombination:
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["runtime"] = 9999999
Review comment:
It should be much clearer now.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #9595: [CUTLASS] Initial conv2d support
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9595:
URL: https://github.com/apache/tvm/pull/9595#discussion_r760174093
##########
File path: python/tvm/contrib/cutlass/conv2d_operation.py
##########
@@ -0,0 +1,240 @@
+# 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=invalid-name, unused-wildcard-import, wildcard-import
+"""Generator for CUTLASS Conv2D kernels."""
+from .library import *
+
+
+class Conv2dOperation:
Review comment:
Operation classes are only used by the corresponding Emit class, and there is not much to share between EmitGemm and EmitConv2d. I prefer to keep this way rather than introducing unnecessary abstractions.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] masahi commented on a change in pull request #9595: [CUTLASS] Initial conv2d support
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9595:
URL: https://github.com/apache/tvm/pull/9595#discussion_r758970801
##########
File path: src/relay/backend/contrib/cutlass/codegen.cc
##########
@@ -234,6 +238,112 @@ std::string BatchMatmulOp(std::string id, const Str2StrMap& attrs,
return gemm_decl.str();
}
+Str2StrMap Conv2dArgs(const Map<String, ObjectRef>& attrs) {
+ Str2StrMap args = ArgsCommon(attrs);
+ auto arg0_shape = attrs["arg0_shape"].as<ArrayNode>();
+ auto arg1_shape = attrs["arg1_shape"].as<ArrayNode>();
+ auto out_shape = attrs["ret_shape"].as<ArrayNode>();
+ args["N"] = GetDimAsStr(arg0_shape->at(0));
+ args["H"] = GetDimAsStr(arg0_shape->at(1));
+ args["W"] = GetDimAsStr(arg0_shape->at(2));
+ args["C"] = GetDimAsStr(arg0_shape->at(3));
+ args["K"] = GetDimAsStr(arg1_shape->at(0));
+ args["R"] = GetDimAsStr(arg1_shape->at(1));
+ args["S"] = GetDimAsStr(arg1_shape->at(1));
+ args["P"] = GetDimAsStr(out_shape->at(1));
+ args["Q"] = GetDimAsStr(out_shape->at(2));
+ args["pad_h"] = GetDimAsStr(attrs["padding"].as<ArrayNode>()->at(0));
+ args["pad_w"] = GetDimAsStr(attrs["padding"].as<ArrayNode>()->at(1));
+ args["stride_h"] = GetDimAsStr(attrs["strides"].as<ArrayNode>()->at(0));
+ args["stride_w"] = GetDimAsStr(attrs["strides"].as<ArrayNode>()->at(1));
+ args["dilation_h"] = GetDimAsStr(attrs["dilation"].as<ArrayNode>()->at(0));
+ args["dilation_w"] = GetDimAsStr(attrs["dilation"].as<ArrayNode>()->at(1));
+ return args;
+}
+
+std::string Conv2dOp(std::string id, const Str2StrMap& attrs,
+ const std::vector<std::string>& func_args) {
+ std::ostringstream conv2d_decl;
+ CutlassPrint(conv2d_decl, "using ElementInputA = " + attrs.at("ElementInputA") + ";\n");
+ CutlassPrint(conv2d_decl, "using ElementInputB = " + attrs.at("ElementInputB") + ";\n");
+ CutlassPrint(conv2d_decl, "using ElementOutput = " + attrs.at("ElementOutput") + ";\n");
+
+ CutlassPrint(conv2d_decl, "using ElementComputeEpilogue = " + attrs.at("ElementOutput") + ";\n");
+ CutlassPrint(conv2d_decl, attrs.at("op_def"));
+ CutlassPrint(conv2d_decl, "using Operation_" + attrs.at("op_name") +
+ " = cutlass::conv::device::ImplicitGemmConvolution<" +
+ attrs.at("op_name") + ">;\n");
+ CutlassPrint(conv2d_decl, "using Conv2d = Operation_" + attrs.at("op_name") + ";\n");
+
+ auto get_dim = [&attrs](const std::string& axis, const std::string& var_name, int axis_idx) {
+ if (attrs.at(axis) == kAnyDim) {
+ return var_name + "->shape[" + std::to_string(axis_idx) + "]";
+ } else {
+ return attrs.at(axis);
+ }
+ };
+
+ CutlassPrint(conv2d_decl, "int N = " + get_dim("N", func_args[0], 0) + ";\n");
+ CutlassPrint(conv2d_decl, "int H = " + get_dim("H", func_args[0], 1) + ";\n");
+ CutlassPrint(conv2d_decl, "int W = " + get_dim("W", func_args[0], 2) + ";\n");
+ CutlassPrint(conv2d_decl, "int C = " + attrs.at("C") + ";\n");
+ CutlassPrint(conv2d_decl, "int K = " + attrs.at("K") + ";\n");
+ CutlassPrint(conv2d_decl, "int R = " + attrs.at("R") + ";\n");
+ CutlassPrint(conv2d_decl, "int S = " + attrs.at("S") + ";\n");
+ CutlassPrint(conv2d_decl, "int P = " + get_dim("P", "out0", 1) + ";\n");
+ CutlassPrint(conv2d_decl, "int Q = " + get_dim("Q", "out0", 2) + ";\n");
+ CutlassPrint(conv2d_decl, "int pad_h = " + attrs.at("pad_h") + ";\n");
+ CutlassPrint(conv2d_decl, "int pad_w = " + attrs.at("pad_w") + ";\n");
+ CutlassPrint(conv2d_decl, "int stride_h = " + attrs.at("stride_h") + ";\n");
+ CutlassPrint(conv2d_decl, "int stride_w = " + attrs.at("stride_w") + ";\n");
+ CutlassPrint(conv2d_decl, "int dilation_h = " + attrs.at("dilation_h") + ";\n");
+ CutlassPrint(conv2d_decl, "int dilation_w = " + attrs.at("dilation_w") + ";\n");
+
+ CutlassPrint(
+ conv2d_decl,
+ "cutlass::conv::Conv2dProblemSize problem_size(N, H, W, C, K, R, S, P, Q, pad_h, pad_w, "
+ "stride_h, stride_w, dilation_h, dilation_w, cutlass::conv::Mode::kCrossCorrelation, 1);\n");
+
+ ICHECK(func_args.size() >= 2);
+ CutlassPrint(conv2d_decl, "void* ptr_a = (void*)(" + func_args[0] + "->data);\n");
+ CutlassPrint(conv2d_decl, "void* ptr_b = (void*)(" + func_args[1] + "->data);\n");
+ CutlassPrint(conv2d_decl, "void* ptr_out = (void*)(out0->data);\n");
+ CutlassPrint(conv2d_decl, "ElementComputeEpilogue alpha = ElementComputeEpilogue(1);\n");
+ CutlassPrint(conv2d_decl, "ElementComputeEpilogue beta = ElementComputeEpilogue(0);\n");
+
+ CutlassPrint(conv2d_decl, "using cutlass::layout::TensorNHWC;\n");
+ CutlassPrint(conv2d_decl,
+ "TensorNHWC layout_A(TensorNHWC::packed(cutlass::make_Coord(N, H, W, C)));\n");
+ CutlassPrint(conv2d_decl,
+ "TensorNHWC layout_B(TensorNHWC::packed(cutlass::make_Coord(K, R, S, C)));\n");
+ CutlassPrint(conv2d_decl,
+ "TensorNHWC layout_C(TensorNHWC::packed(cutlass::make_Coord(N, P, Q, K)));\n");
+ CutlassPrint(conv2d_decl, "typename Conv2d::Arguments arguments{\n");
+ CutlassPrint(conv2d_decl, " problem_size,\n");
+ CutlassPrint(conv2d_decl, " {static_cast<ElementInputA*>(ptr_a), layout_A},\n");
+ CutlassPrint(conv2d_decl, " {static_cast<ElementInputB*>(ptr_b), layout_B},\n");
+ CutlassPrint(conv2d_decl, " {static_cast<ElementOutput*>(ptr_out),layout_C},\n");
+ CutlassPrint(conv2d_decl, " {static_cast<ElementOutput*>(ptr_out),layout_C},\n");
+ CutlassPrint(conv2d_decl, "{alpha, beta}\n};\n");
+ CutlassPrint(conv2d_decl, "Conv2d conv2d_op;\n");
+
+ CutlassPrint(conv2d_decl, "size_t workspace_size = conv2d_op.get_workspace_size(arguments);\n");
+ // Allocate workspace memory
+ CutlassPrint(conv2d_decl,
+ "cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);\n");
Review comment:
Does the destructor clean itself up? https://github.com/NVIDIA/cutlass/blob/8f8a80cad57950be2d538ac0ead420740662318d/tools/util/include/cutlass/util/device_memory.h#L199
This is basically the same code as gemm op.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
[GitHub] [tvm] Laurawly commented on a change in pull request #9595: [CUTLASS] Initial conv2d support
Posted by GitBox <gi...@apache.org>.
Laurawly commented on a change in pull request #9595:
URL: https://github.com/apache/tvm/pull/9595#discussion_r758694948
##########
File path: src/relay/backend/contrib/cutlass/codegen.cc
##########
@@ -234,6 +238,112 @@ std::string BatchMatmulOp(std::string id, const Str2StrMap& attrs,
return gemm_decl.str();
}
+Str2StrMap Conv2dArgs(const Map<String, ObjectRef>& attrs) {
+ Str2StrMap args = ArgsCommon(attrs);
+ auto arg0_shape = attrs["arg0_shape"].as<ArrayNode>();
+ auto arg1_shape = attrs["arg1_shape"].as<ArrayNode>();
+ auto out_shape = attrs["ret_shape"].as<ArrayNode>();
+ args["N"] = GetDimAsStr(arg0_shape->at(0));
+ args["H"] = GetDimAsStr(arg0_shape->at(1));
+ args["W"] = GetDimAsStr(arg0_shape->at(2));
+ args["C"] = GetDimAsStr(arg0_shape->at(3));
+ args["K"] = GetDimAsStr(arg1_shape->at(0));
+ args["R"] = GetDimAsStr(arg1_shape->at(1));
+ args["S"] = GetDimAsStr(arg1_shape->at(1));
+ args["P"] = GetDimAsStr(out_shape->at(1));
+ args["Q"] = GetDimAsStr(out_shape->at(2));
+ args["pad_h"] = GetDimAsStr(attrs["padding"].as<ArrayNode>()->at(0));
+ args["pad_w"] = GetDimAsStr(attrs["padding"].as<ArrayNode>()->at(1));
+ args["stride_h"] = GetDimAsStr(attrs["strides"].as<ArrayNode>()->at(0));
+ args["stride_w"] = GetDimAsStr(attrs["strides"].as<ArrayNode>()->at(1));
+ args["dilation_h"] = GetDimAsStr(attrs["dilation"].as<ArrayNode>()->at(0));
+ args["dilation_w"] = GetDimAsStr(attrs["dilation"].as<ArrayNode>()->at(1));
+ return args;
+}
+
+std::string Conv2dOp(std::string id, const Str2StrMap& attrs,
+ const std::vector<std::string>& func_args) {
+ std::ostringstream conv2d_decl;
+ CutlassPrint(conv2d_decl, "using ElementInputA = " + attrs.at("ElementInputA") + ";\n");
+ CutlassPrint(conv2d_decl, "using ElementInputB = " + attrs.at("ElementInputB") + ";\n");
+ CutlassPrint(conv2d_decl, "using ElementOutput = " + attrs.at("ElementOutput") + ";\n");
+
+ CutlassPrint(conv2d_decl, "using ElementComputeEpilogue = " + attrs.at("ElementOutput") + ";\n");
+ CutlassPrint(conv2d_decl, attrs.at("op_def"));
+ CutlassPrint(conv2d_decl, "using Operation_" + attrs.at("op_name") +
+ " = cutlass::conv::device::ImplicitGemmConvolution<" +
+ attrs.at("op_name") + ">;\n");
+ CutlassPrint(conv2d_decl, "using Conv2d = Operation_" + attrs.at("op_name") + ";\n");
+
+ auto get_dim = [&attrs](const std::string& axis, const std::string& var_name, int axis_idx) {
+ if (attrs.at(axis) == kAnyDim) {
+ return var_name + "->shape[" + std::to_string(axis_idx) + "]";
+ } else {
+ return attrs.at(axis);
+ }
+ };
+
+ CutlassPrint(conv2d_decl, "int N = " + get_dim("N", func_args[0], 0) + ";\n");
+ CutlassPrint(conv2d_decl, "int H = " + get_dim("H", func_args[0], 1) + ";\n");
+ CutlassPrint(conv2d_decl, "int W = " + get_dim("W", func_args[0], 2) + ";\n");
+ CutlassPrint(conv2d_decl, "int C = " + attrs.at("C") + ";\n");
+ CutlassPrint(conv2d_decl, "int K = " + attrs.at("K") + ";\n");
+ CutlassPrint(conv2d_decl, "int R = " + attrs.at("R") + ";\n");
+ CutlassPrint(conv2d_decl, "int S = " + attrs.at("S") + ";\n");
+ CutlassPrint(conv2d_decl, "int P = " + get_dim("P", "out0", 1) + ";\n");
+ CutlassPrint(conv2d_decl, "int Q = " + get_dim("Q", "out0", 2) + ";\n");
+ CutlassPrint(conv2d_decl, "int pad_h = " + attrs.at("pad_h") + ";\n");
+ CutlassPrint(conv2d_decl, "int pad_w = " + attrs.at("pad_w") + ";\n");
+ CutlassPrint(conv2d_decl, "int stride_h = " + attrs.at("stride_h") + ";\n");
+ CutlassPrint(conv2d_decl, "int stride_w = " + attrs.at("stride_w") + ";\n");
+ CutlassPrint(conv2d_decl, "int dilation_h = " + attrs.at("dilation_h") + ";\n");
+ CutlassPrint(conv2d_decl, "int dilation_w = " + attrs.at("dilation_w") + ";\n");
+
+ CutlassPrint(
+ conv2d_decl,
+ "cutlass::conv::Conv2dProblemSize problem_size(N, H, W, C, K, R, S, P, Q, pad_h, pad_w, "
+ "stride_h, stride_w, dilation_h, dilation_w, cutlass::conv::Mode::kCrossCorrelation, 1);\n");
+
+ ICHECK(func_args.size() >= 2);
+ CutlassPrint(conv2d_decl, "void* ptr_a = (void*)(" + func_args[0] + "->data);\n");
+ CutlassPrint(conv2d_decl, "void* ptr_b = (void*)(" + func_args[1] + "->data);\n");
+ CutlassPrint(conv2d_decl, "void* ptr_out = (void*)(out0->data);\n");
+ CutlassPrint(conv2d_decl, "ElementComputeEpilogue alpha = ElementComputeEpilogue(1);\n");
+ CutlassPrint(conv2d_decl, "ElementComputeEpilogue beta = ElementComputeEpilogue(0);\n");
+
+ CutlassPrint(conv2d_decl, "using cutlass::layout::TensorNHWC;\n");
+ CutlassPrint(conv2d_decl,
+ "TensorNHWC layout_A(TensorNHWC::packed(cutlass::make_Coord(N, H, W, C)));\n");
+ CutlassPrint(conv2d_decl,
+ "TensorNHWC layout_B(TensorNHWC::packed(cutlass::make_Coord(K, R, S, C)));\n");
+ CutlassPrint(conv2d_decl,
+ "TensorNHWC layout_C(TensorNHWC::packed(cutlass::make_Coord(N, P, Q, K)));\n");
+ CutlassPrint(conv2d_decl, "typename Conv2d::Arguments arguments{\n");
+ CutlassPrint(conv2d_decl, " problem_size,\n");
+ CutlassPrint(conv2d_decl, " {static_cast<ElementInputA*>(ptr_a), layout_A},\n");
+ CutlassPrint(conv2d_decl, " {static_cast<ElementInputB*>(ptr_b), layout_B},\n");
+ CutlassPrint(conv2d_decl, " {static_cast<ElementOutput*>(ptr_out),layout_C},\n");
+ CutlassPrint(conv2d_decl, " {static_cast<ElementOutput*>(ptr_out),layout_C},\n");
+ CutlassPrint(conv2d_decl, "{alpha, beta}\n};\n");
+ CutlassPrint(conv2d_decl, "Conv2d conv2d_op;\n");
+
+ CutlassPrint(conv2d_decl, "size_t workspace_size = conv2d_op.get_workspace_size(arguments);\n");
+ // Allocate workspace memory
+ CutlassPrint(conv2d_decl,
+ "cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);\n");
Review comment:
There's memory leak by allocating workspace this way. @ZihengJiang
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@tvm.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org