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/10/12 16:54:03 UTC
[GitHub] [tvm] Laurawly opened a new pull request #9261: [BYOC] CUTLASS integration
Laurawly opened a new pull request #9261:
URL: https://github.com/apache/tvm/pull/9261
As discussed in RFC: https://discuss.tvm.apache.org/t/rfc-byoc-nvidia-cutlass-integration/9147, this PR is a CUTLASS integration of GEMM to TVM. It also includes a profiler to search for best params in CUTLASS.
@masahi Please take over this PR, Thanks!
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737949955
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
Review comment:
I see, I always thought that `accum_type == output_type`. Since we don't do the third variant in TVM, I think removing that logic is a correct thing to do. Is that right?
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737951454
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
+
+ ops.extend(out)
+
+ return ops
+
+
+def generate_sm75_tensor_op_1688(out_dtype):
+ """Generate GEMM kernels for Turing."""
+ assert out_dtype in ["float32", "float16"]
+ math_instructions = {
+ "float32": [
+ MathInstruction(
+ [16, 8, 8],
+ DataType.f16,
+ DataType.f16,
+ DataType.f32,
+ OpcodeClass.TensorOp,
+ MathOperation.multiply_add,
+ )
+ ],
+ "float16": [
+ MathInstruction(
+ [16, 8, 8],
+ DataType.f16,
+ DataType.f16,
+ DataType.f16,
+ OpcodeClass.TensorOp,
+ MathOperation.multiply_add,
+ )
+ ],
+ }[out_dtype]
Review comment:
A bit confused about your comment. Here, I'm selecting either fp32 or fp16 accum / output kernels based on `out_dtype` of Relay `dense` op, assuming that `accum_dtype == out_dtype`. If the user requests fp32 output dtype, we must select fp32 accumulation, don't we?
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
Laurawly commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r738768020
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
Review comment:
BTW, NT layout has best performance for cutlass (A: col major, B: row major)
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942064270
IMHO, CUTLASS doesn't naturally benefit dynamic workloads due to the exact reason you mentioned. We internally use CUTLASS for training and it works well because we JIT kernels with known shapes in runtime.
In the case of CUTLASS with BYOC in TVM for inference, my impression is we could leverage high performance kernel templates while 1) keeping the binary self-contained, 2) fusing ops, and 3) having lightweight tuning (e.g., ~10 trials similar to CUDNN). On the other hand, dynamic workloads are still challenging, and hopefully our ongoing efforts of dynamic kernel tuning and generation could be landed soon to make it happen.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942878845
Right. Probably just need some experiments to see how many kernels should be generated to cover a range of batch sizes, and we may just hardcore the dispatching logic first to make it work. With these restrictions for now, C source codegen is also doable and it also has less overhead.
--
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] zhiics commented on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
zhiics commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942200090
I think even C source code should be able to handle dynamic shapes as well since it only expects tensors at the runtime, or I might forget something here. But in general, json format is more recommended as it is easier to debug and maintain. It is also more friendly in handling constants.
--
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 edited a comment on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-941990346
One question for @comaniac: Currently, this integration is implemented with C source backend. But one big use case for cutlass BYOC is for dynamic workloads, for which we cannot do well at the moment. I think C source backend is not compatible with dynamic inputs, is that right? On the other hand, with json codegen/runtime, I can see a way to support dynamic inputs.
However, since cutlass is a template library, the only way it could work with the json codegen/runtime is to instantiate all possible templates that we care about ahead of time, and build them together with `libtvm.so`. I wonder if that is feasible cc @Laurawly.
One benefit of the above approach is that we don't need to commit the cutlass code generator into our repo, which is a big part of this PR.
--
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] junrushao1994 commented on a change in pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r727471927
##########
File path: .gitmodules
##########
@@ -13,3 +13,6 @@
[submodule "3rdparty/libbacktrace"]
path = 3rdparty/libbacktrace
url = https://github.com/tlc-pack/libbacktrace.git
+[submodule "cutlass_profiler/cutlass"]
+ path = cutlass_profiler/cutlass
+ url = git@github.com:NVIDIA/cutlass.git
Review comment:
Use https btw
```suggestion
[submodule "cutlass_profiler/cutlass"]
path = cutlass_profiler/cutlass
url = https://github.com/NVIDIA/cutlass.git
```
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-948543423
This is ready for review. The bulk of PR is a kernel generator taken from cutlass, simplified for our purpose + more epilogue support.
The followings are supported:
* `dense` with bias, bias + relu, and bias + gelu epilogue
* fp16 and fp32 accumulation
* Kernels for Turing and Ampere
The follow-up work will add more features:
* dynamic input (only basic one, the heuristic stuff is complicated so I defer it for future)
* conv2d and batched gemm
* More epilogues to support e2d models. I'm planning to test MaskRCNN, DeeplabV3, EfficientNet v2, and BERT-large. This require some work to extend cutlass (see https://github.com/NVIDIA/cutlass/discussions/347)
Low priority (for me) but things we should support eventually:
* Heuristic for dynamic inputs
* int8 and int4 kernels
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r733617521
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
Review comment:
@Laurawly This is one of the things I deviated from the original cutlass code and your initial code. https://github.com/NVIDIA/cutlass/blob/2e07c4cc2fc94f4cd396ecf1e9132caf1efba50e/tools/library/scripts/generator.py#L599-L606
The original code generates two sets of kernels for fp16 and fp32 accumulation. But in our case, generating both of them is not a good idea because
1) We need to compile and tune both set of kernels but we know we only need either fp16 or fp32 one
2) If we want fp32 accum but an fp16 accum kernel is selected by the profiler (which will always happen), the final compilation would fail due to type mismatch in the accumulation related template parameters (some say `float` while other say `cutlass::half_t`)
--
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] zhiics commented on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
zhiics commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942200090
I think even C source code should be able to handle dynamic shapes as well since it only expects tensors at the runtime, or I might forget something here. But in general, json format is more recommended as it is easier to debug and maintain. It is also more friendly in handling constants.
--
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] junrushao1994 commented on a change in pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r727471035
##########
File path: .gitmodules
##########
@@ -13,3 +13,6 @@
[submodule "3rdparty/libbacktrace"]
path = 3rdparty/libbacktrace
url = https://github.com/tlc-pack/libbacktrace.git
+[submodule "cutlass_profiler/cutlass"]
+ path = cutlass_profiler/cutlass
+ url = git@github.com:NVIDIA/cutlass.git
Review comment:
Shall we move this folder to 3rdparty as our convention in the repo?
##########
File path: .gitmodules
##########
@@ -13,3 +13,6 @@
[submodule "3rdparty/libbacktrace"]
path = 3rdparty/libbacktrace
url = https://github.com/tlc-pack/libbacktrace.git
+[submodule "cutlass_profiler/cutlass"]
+ path = cutlass_profiler/cutlass
+ url = git@github.com:NVIDIA/cutlass.git
Review comment:
Use https btw
```suggestion
[submodule "cutlass_profiler/cutlass"]
path = cutlass_profiler/cutlass
url = https://github.com/NVIDIA/cutlass.git
```
--
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] hwu36 commented on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
hwu36 commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942558026
> I assume they also generate a run time dispatching logic (also in the generated C source code) to determine which kernel should be used given the known shape in run time.
At this moment, CUTLASS does not provide heuristics to choose which kernel to use based on the runtime information. The user of CUTLASS needs to make the decision to pick kernels.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942895664
I think it is reasonable to ask users for an estimate of a range of dynamism. Without such estimate, we just fallback to a default kernel that works on any input shapes, which shouldn't be too slow anyway since it is cutlass, not TVM :) This should be a reasonable trade off between performance and convenience of dynamic models: if they want to run fast on dynamic models, they need to help us, otherwise we cannot guarantee good performance.
I believe even if a network contains many dynamic workloads, there is only one underlying dynamism, from which other dynamic shapes can be derived. For example, MaskRCNN contains some dynamic batch dense, conv2d and conv2d transpose, but all dynamic dims directly correspond to the number of boxes detected, or some numbers derived from it (such as the number of box times * pool window size * feature size, the latter two are fixed numbers). For transformers, I guess the true dynamism is only in sequence length. So I think it wouldn't be too much of an ask to give an estimate of such numbers.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737954113
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
Review comment:
Update: I realized that Relay does allow users specify the output dtype of dense ops. Then we have:
1. When output dtype is FP32, we must use the 1st kernel.
2. When output dtype is FP16, we could use either 2nd or 3rd kernel.
Again although we don't do 3rd variant in TVM, the accumulation dtype is just used internally as a part of the kernel implementation, so we could still consider it.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737956211
##########
File path: tests/python/contrib/test_cutlass.py
##########
@@ -0,0 +1,121 @@
+# 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.
+import math
+import pytest
+import tvm
+from tvm import relay
+import numpy as np
+from tvm.contrib.cutlass import profile_and_build
+
+
+def get_ref_rt_mod(mod, params):
+ with tvm.transform.PassContext(opt_level=3):
+ lib = relay.build(mod, target="cuda", params=params)
+ dev = tvm.device("cuda", 0)
+ rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ return rt_mod, dev
+
+
+def get_output(rt_mod, x):
+ rt_mod.set_input("data", x)
+ rt_mod.run()
+ return rt_mod.get_output(0).asnumpy()
+
+
+def get_dense(M, N, K, out_dtype="float16"):
+ data = relay.var("data", shape=(M, K), dtype="float16")
+ weight = relay.var("weight", shape=(N, K), dtype="float16")
+ return relay.nn.dense(data, weight, out_dtype=out_dtype)
+
+
+def get_dense_bias(M, N, K, out_dtype="float16"):
+ dense = get_dense(M, N, K, out_dtype=out_dtype)
+ bias = relay.var("bias", shape=(N,), dtype=out_dtype)
+ return relay.nn.bias_add(dense, bias)
+
+
+def get_dense_bias_relu(M, N, K, out_dtype="float16"):
+ return relay.nn.relu(get_dense_bias(M, N, K, out_dtype="float16"))
+
+
+def get_dense_bias_gelu(M, N, K, out_dtype="float16"):
+ bias_add = get_dense_bias(M, N, K, out_dtype)
+ mul = bias_add * relay.const((1.0 / math.sqrt(2.0)), dtype=out_dtype)
+ if out_dtype == "float16":
+ erf = relay.cast(relay.op.erf(relay.cast(mul, "float32")), "float16")
+ else:
+ erf = relay.op.erf(mul)
+ mul_half = erf * relay.const(0.5, dtype=out_dtype)
+ add = mul_half + relay.const(0.5, dtype=out_dtype)
+ return add * bias_add
+
+
+def verify(func, M, N, K, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False):
+ if not tvm.get_global_func("relay.ext.cutlass", True):
+ return
+ mod = tvm.IRModule.from_expr(func)
+ typ = relay.transform.InferType()(mod)
+ out_dtype = typ["main"].body.checked_type.dtype
+ np_data = np.random.uniform(-1, 1, (M, K)).astype("float16")
+ np_weight = np.random.uniform(-1, 1, (N, K)).astype("float16")
+ np_bias = np.random.uniform(-1, 1, (N,)).astype(out_dtype)
+
+ params = {"weight": np_weight, "bias": np_bias}
+
+ rt_mod_ref, dev = get_ref_rt_mod(mod, params)
+ rt_mod, dev, num_partition = profile_and_build(mod, params, sm, tmp_dir="tmp")
Review comment:
What I'm worrying about is actually the CI time. I believe the tuning time in this case would definitely be much shorter than AutoTVM/Ansor, so it should be fine in the real use cases. Maybe we could provide some heuristics. For example:
1. A level/mode indicating the number of candidates to be profiled. We could use the lowest level in the CI, and recommend the highest level for users.
2. Simply let users specify the maximum number of candidates.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737957304
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
+
+ ops.extend(out)
+
+ return ops
+
+
+def generate_sm75_tensor_op_1688(out_dtype):
+ """Generate GEMM kernels for Turing."""
+ assert out_dtype in ["float32", "float16"]
+ math_instructions = {
+ "float32": [
+ MathInstruction(
+ [16, 8, 8],
+ DataType.f16,
+ DataType.f16,
+ DataType.f32,
+ OpcodeClass.TensorOp,
+ MathOperation.multiply_add,
+ )
+ ],
+ "float16": [
+ MathInstruction(
+ [16, 8, 8],
+ DataType.f16,
+ DataType.f16,
+ DataType.f16,
+ OpcodeClass.TensorOp,
+ MathOperation.multiply_add,
+ )
+ ],
+ }[out_dtype]
Review comment:
Hmm...the accuracy might be the one, as accumulating to FP16 may result in overflow. This is sensitive for training, but it might be fine for inference.
--
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] hwu36 commented on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
hwu36 commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-954353094
This is so cool. Thank you everyone.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942580043
> At this moment, CUTLASS does not provide heuristics to choose which kernel to use based on the runtime information. The user of CUTLASS needs to make the decision to pick kernels.
That's what I meant, so we need to have generate this logic in the CUTLASS BYOC codegen.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737956211
##########
File path: tests/python/contrib/test_cutlass.py
##########
@@ -0,0 +1,121 @@
+# 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.
+import math
+import pytest
+import tvm
+from tvm import relay
+import numpy as np
+from tvm.contrib.cutlass import profile_and_build
+
+
+def get_ref_rt_mod(mod, params):
+ with tvm.transform.PassContext(opt_level=3):
+ lib = relay.build(mod, target="cuda", params=params)
+ dev = tvm.device("cuda", 0)
+ rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ return rt_mod, dev
+
+
+def get_output(rt_mod, x):
+ rt_mod.set_input("data", x)
+ rt_mod.run()
+ return rt_mod.get_output(0).asnumpy()
+
+
+def get_dense(M, N, K, out_dtype="float16"):
+ data = relay.var("data", shape=(M, K), dtype="float16")
+ weight = relay.var("weight", shape=(N, K), dtype="float16")
+ return relay.nn.dense(data, weight, out_dtype=out_dtype)
+
+
+def get_dense_bias(M, N, K, out_dtype="float16"):
+ dense = get_dense(M, N, K, out_dtype=out_dtype)
+ bias = relay.var("bias", shape=(N,), dtype=out_dtype)
+ return relay.nn.bias_add(dense, bias)
+
+
+def get_dense_bias_relu(M, N, K, out_dtype="float16"):
+ return relay.nn.relu(get_dense_bias(M, N, K, out_dtype="float16"))
+
+
+def get_dense_bias_gelu(M, N, K, out_dtype="float16"):
+ bias_add = get_dense_bias(M, N, K, out_dtype)
+ mul = bias_add * relay.const((1.0 / math.sqrt(2.0)), dtype=out_dtype)
+ if out_dtype == "float16":
+ erf = relay.cast(relay.op.erf(relay.cast(mul, "float32")), "float16")
+ else:
+ erf = relay.op.erf(mul)
+ mul_half = erf * relay.const(0.5, dtype=out_dtype)
+ add = mul_half + relay.const(0.5, dtype=out_dtype)
+ return add * bias_add
+
+
+def verify(func, M, N, K, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False):
+ if not tvm.get_global_func("relay.ext.cutlass", True):
+ return
+ mod = tvm.IRModule.from_expr(func)
+ typ = relay.transform.InferType()(mod)
+ out_dtype = typ["main"].body.checked_type.dtype
+ np_data = np.random.uniform(-1, 1, (M, K)).astype("float16")
+ np_weight = np.random.uniform(-1, 1, (N, K)).astype("float16")
+ np_bias = np.random.uniform(-1, 1, (N,)).astype(out_dtype)
+
+ params = {"weight": np_weight, "bias": np_bias}
+
+ rt_mod_ref, dev = get_ref_rt_mod(mod, params)
+ rt_mod, dev, num_partition = profile_and_build(mod, params, sm, tmp_dir="tmp")
Review comment:
What I'm worrying about is actually the CI time. I believe the tuning time in this case would definitely be much shorter than AutoTVM/Ansor, so it should be fine in the real use cases. Maybe we could provide some heuristics to control the CI time. For example:
1. A level/mode indicating the number of candidates to be profiled. We could use the lowest level in the CI, and recommend the highest level for users.
2. Simply let users specify the maximum number of candidates.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r738602868
##########
File path: python/tvm/contrib/cutlass/build.py
##########
@@ -0,0 +1,155 @@
+# 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
+"""Driver for partitioning and building a Relay module for CUTLASS offload."""
+import tvm
+from tvm import runtime, relay
+from tvm.relay.op.contrib.cutlass import partition_for_cutlass
+from .gen_gemm import CutlassGemmProfiler
+
+
+class GemmAnnotator(tvm.relay.ExprVisitor):
+ """Annotates partitioned functions with shape and dtype information."""
+
+ def __init__(self):
+ super().__init__()
+ self.signature = {}
+
+ def visit_call(self, call):
+ op = call.op
+ if isinstance(op, relay.Function) and "PartitionedFromPattern" in op.attrs:
+ self.signature["op_type"] = op.attrs["Composite"]
+ for i, arg in enumerate(op.params):
+ self.signature["arg%d_shape" % i] = arg.checked_type.shape
+ self.signature["arg%d_dtype" % i] = arg.checked_type.dtype
+ self.signature["ret_shape"] = op.ret_type.shape
+ self.signature["ret_dtype"] = op.ret_type.dtype
+
+
+def profile_and_build(mod, params, sm, tmp_dir="./tmp", lib_path="compile.so"):
Review comment:
This one looks much better to me :)
I guess many people would forget to call `build_cutlass_kernels`, but this is probably the best way we can achieve at this moment; otherwise we need to improve `relay.build`.
--
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] zhiics commented on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
zhiics commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942853797
> > At this moment, CUTLASS does not provide heuristics to choose which kernel to use based on the runtime information. The user of CUTLASS needs to make the decision to pick kernels.
>
>
>
> That's what I meant, so we need to have generate this logic in the CUTLASS BYOC codegen.
I agree that kernel picking should be handled by the custom runtime.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-941443685
Thanks @Laurawly, I'll work on this as my top priority!
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-941990346
One question for @comaniac: Currently, this integration is implemented with C source backend. But one big use case for cutlass BYOC is for dynamic workloads, for which we cannot do well at the moment. I think C source backend is not compatible with dynamic inputs, is that right? On the other hand, with json codegen/runtime, I can see a way to support dynamic inputs.
However, since cutlass is a template library, the only way it could work with the json codegen/runtime is to instantiate all possible templates that we care about ahead of time, and build them together with `libtvm.so`. I wonder if that is feasible cc @Laurawly.
--
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 edited a comment on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-941990346
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-941443685
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737956332
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
Review comment:
Yes, I have a test that exercises hetero output dtype, and it works https://github.com/apache/tvm/blob/ff66ef49ecb9d1946511b545fd4d92359f60a813/tests/python/contrib/test_cutlass.py#L102
Although we currently don't generate such op, I'm not ruling out such use cases (e.g. if we want to make our so-called "ToMixedPrecision" pass truly support mixed precision).
Cutlass may support the 3rd kernel, but we don't have a good reason to support it, as discussed in the thread above.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737958148
##########
File path: tests/python/contrib/test_cutlass.py
##########
@@ -0,0 +1,121 @@
+# 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.
+import math
+import pytest
+import tvm
+from tvm import relay
+import numpy as np
+from tvm.contrib.cutlass import profile_and_build
+
+
+def get_ref_rt_mod(mod, params):
+ with tvm.transform.PassContext(opt_level=3):
+ lib = relay.build(mod, target="cuda", params=params)
+ dev = tvm.device("cuda", 0)
+ rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ return rt_mod, dev
+
+
+def get_output(rt_mod, x):
+ rt_mod.set_input("data", x)
+ rt_mod.run()
+ return rt_mod.get_output(0).asnumpy()
+
+
+def get_dense(M, N, K, out_dtype="float16"):
+ data = relay.var("data", shape=(M, K), dtype="float16")
+ weight = relay.var("weight", shape=(N, K), dtype="float16")
+ return relay.nn.dense(data, weight, out_dtype=out_dtype)
+
+
+def get_dense_bias(M, N, K, out_dtype="float16"):
+ dense = get_dense(M, N, K, out_dtype=out_dtype)
+ bias = relay.var("bias", shape=(N,), dtype=out_dtype)
+ return relay.nn.bias_add(dense, bias)
+
+
+def get_dense_bias_relu(M, N, K, out_dtype="float16"):
+ return relay.nn.relu(get_dense_bias(M, N, K, out_dtype="float16"))
+
+
+def get_dense_bias_gelu(M, N, K, out_dtype="float16"):
+ bias_add = get_dense_bias(M, N, K, out_dtype)
+ mul = bias_add * relay.const((1.0 / math.sqrt(2.0)), dtype=out_dtype)
+ if out_dtype == "float16":
+ erf = relay.cast(relay.op.erf(relay.cast(mul, "float32")), "float16")
+ else:
+ erf = relay.op.erf(mul)
+ mul_half = erf * relay.const(0.5, dtype=out_dtype)
+ add = mul_half + relay.const(0.5, dtype=out_dtype)
+ return add * bias_add
+
+
+def verify(func, M, N, K, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False):
+ if not tvm.get_global_func("relay.ext.cutlass", True):
+ return
+ mod = tvm.IRModule.from_expr(func)
+ typ = relay.transform.InferType()(mod)
+ out_dtype = typ["main"].body.checked_type.dtype
+ np_data = np.random.uniform(-1, 1, (M, K)).astype("float16")
+ np_weight = np.random.uniform(-1, 1, (N, K)).astype("float16")
+ np_bias = np.random.uniform(-1, 1, (N,)).astype(out_dtype)
+
+ params = {"weight": np_weight, "bias": np_bias}
+
+ rt_mod_ref, dev = get_ref_rt_mod(mod, params)
+ rt_mod, dev, num_partition = profile_and_build(mod, params, sm, tmp_dir="tmp")
Review comment:
Yeah, for testing purpose there is no need to tune. I think "one or all" granularity is fine (the third one).
I'm still going to add an option for compiling with `multiprocessing`.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737958148
##########
File path: tests/python/contrib/test_cutlass.py
##########
@@ -0,0 +1,121 @@
+# 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.
+import math
+import pytest
+import tvm
+from tvm import relay
+import numpy as np
+from tvm.contrib.cutlass import profile_and_build
+
+
+def get_ref_rt_mod(mod, params):
+ with tvm.transform.PassContext(opt_level=3):
+ lib = relay.build(mod, target="cuda", params=params)
+ dev = tvm.device("cuda", 0)
+ rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ return rt_mod, dev
+
+
+def get_output(rt_mod, x):
+ rt_mod.set_input("data", x)
+ rt_mod.run()
+ return rt_mod.get_output(0).asnumpy()
+
+
+def get_dense(M, N, K, out_dtype="float16"):
+ data = relay.var("data", shape=(M, K), dtype="float16")
+ weight = relay.var("weight", shape=(N, K), dtype="float16")
+ return relay.nn.dense(data, weight, out_dtype=out_dtype)
+
+
+def get_dense_bias(M, N, K, out_dtype="float16"):
+ dense = get_dense(M, N, K, out_dtype=out_dtype)
+ bias = relay.var("bias", shape=(N,), dtype=out_dtype)
+ return relay.nn.bias_add(dense, bias)
+
+
+def get_dense_bias_relu(M, N, K, out_dtype="float16"):
+ return relay.nn.relu(get_dense_bias(M, N, K, out_dtype="float16"))
+
+
+def get_dense_bias_gelu(M, N, K, out_dtype="float16"):
+ bias_add = get_dense_bias(M, N, K, out_dtype)
+ mul = bias_add * relay.const((1.0 / math.sqrt(2.0)), dtype=out_dtype)
+ if out_dtype == "float16":
+ erf = relay.cast(relay.op.erf(relay.cast(mul, "float32")), "float16")
+ else:
+ erf = relay.op.erf(mul)
+ mul_half = erf * relay.const(0.5, dtype=out_dtype)
+ add = mul_half + relay.const(0.5, dtype=out_dtype)
+ return add * bias_add
+
+
+def verify(func, M, N, K, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False):
+ if not tvm.get_global_func("relay.ext.cutlass", True):
+ return
+ mod = tvm.IRModule.from_expr(func)
+ typ = relay.transform.InferType()(mod)
+ out_dtype = typ["main"].body.checked_type.dtype
+ np_data = np.random.uniform(-1, 1, (M, K)).astype("float16")
+ np_weight = np.random.uniform(-1, 1, (N, K)).astype("float16")
+ np_bias = np.random.uniform(-1, 1, (N,)).astype(out_dtype)
+
+ params = {"weight": np_weight, "bias": np_bias}
+
+ rt_mod_ref, dev = get_ref_rt_mod(mod, params)
+ rt_mod, dev, num_partition = profile_and_build(mod, params, sm, tmp_dir="tmp")
Review comment:
Yeah, for testing purpose there is no need to tune. I think "one or all" granularity is fine (the third one). Right now I'm only supporting sm75 or later GPUs, the one on CI might be too old.
I'm still going to add an option for compiling with `multiprocessing`.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737953139
##########
File path: tests/python/contrib/test_cutlass.py
##########
@@ -0,0 +1,121 @@
+# 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.
+import math
+import pytest
+import tvm
+from tvm import relay
+import numpy as np
+from tvm.contrib.cutlass import profile_and_build
+
+
+def get_ref_rt_mod(mod, params):
+ with tvm.transform.PassContext(opt_level=3):
+ lib = relay.build(mod, target="cuda", params=params)
+ dev = tvm.device("cuda", 0)
+ rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ return rt_mod, dev
+
+
+def get_output(rt_mod, x):
+ rt_mod.set_input("data", x)
+ rt_mod.run()
+ return rt_mod.get_output(0).asnumpy()
+
+
+def get_dense(M, N, K, out_dtype="float16"):
+ data = relay.var("data", shape=(M, K), dtype="float16")
+ weight = relay.var("weight", shape=(N, K), dtype="float16")
+ return relay.nn.dense(data, weight, out_dtype=out_dtype)
+
+
+def get_dense_bias(M, N, K, out_dtype="float16"):
+ dense = get_dense(M, N, K, out_dtype=out_dtype)
+ bias = relay.var("bias", shape=(N,), dtype=out_dtype)
+ return relay.nn.bias_add(dense, bias)
+
+
+def get_dense_bias_relu(M, N, K, out_dtype="float16"):
+ return relay.nn.relu(get_dense_bias(M, N, K, out_dtype="float16"))
+
+
+def get_dense_bias_gelu(M, N, K, out_dtype="float16"):
+ bias_add = get_dense_bias(M, N, K, out_dtype)
+ mul = bias_add * relay.const((1.0 / math.sqrt(2.0)), dtype=out_dtype)
+ if out_dtype == "float16":
+ erf = relay.cast(relay.op.erf(relay.cast(mul, "float32")), "float16")
+ else:
+ erf = relay.op.erf(mul)
+ mul_half = erf * relay.const(0.5, dtype=out_dtype)
+ add = mul_half + relay.const(0.5, dtype=out_dtype)
+ return add * bias_add
+
+
+def verify(func, M, N, K, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False):
+ if not tvm.get_global_func("relay.ext.cutlass", True):
+ return
+ mod = tvm.IRModule.from_expr(func)
+ typ = relay.transform.InferType()(mod)
+ out_dtype = typ["main"].body.checked_type.dtype
+ np_data = np.random.uniform(-1, 1, (M, K)).astype("float16")
+ np_weight = np.random.uniform(-1, 1, (N, K)).astype("float16")
+ np_bias = np.random.uniform(-1, 1, (N,)).astype(out_dtype)
+
+ params = {"weight": np_weight, "bias": np_bias}
+
+ rt_mod_ref, dev = get_ref_rt_mod(mod, params)
+ rt_mod, dev, num_partition = profile_and_build(mod, params, sm, tmp_dir="tmp")
Review comment:
We might want to compile all executables on all cores, rather than repeating compile -> profile sequentially. What do you think? At least we should enable such feature optionally. cc @Laurawly
--
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 merged pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi merged pull request #9261:
URL: https://github.com/apache/tvm/pull/9261
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-952990355
ping @Laurawly @comaniac @junrushao1994 @vinx13 for review
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r738601111
##########
File path: python/tvm/contrib/cutlass/build.py
##########
@@ -0,0 +1,172 @@
+# 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
+"""Driver for partitioning and building a Relay module for CUTLASS offload."""
+import tvm
+from tvm import runtime, relay
+from .gen_gemm import CutlassGemmProfiler
+
+
+class GemmAnnotator(tvm.relay.ExprVisitor):
+ """Annotates partitioned functions with shape and dtype information."""
+
+ def __init__(self):
+ super().__init__()
+ self.signature = {}
+
+ def visit_call(self, call):
+ op = call.op
+ if isinstance(op, relay.Function) and "PartitionedFromPattern" in op.attrs:
+ self.signature["op_type"] = op.attrs["Composite"]
+ for i, arg in enumerate(op.params):
+ self.signature["arg%d_shape" % i] = arg.checked_type.shape
+ self.signature["arg%d_dtype" % i] = arg.checked_type.dtype
+ self.signature["ret_shape"] = op.ret_type.shape
+ self.signature["ret_dtype"] = op.ret_type.dtype
+
+
+def profile_cutlass_kernels(mod, sm, profile_all=True, use_multiprocessing=False, tmp_dir="./tmp"):
Review comment:
I'm wondering whether we should use "profile" here. My impression to "profile" is more like an analysis rather than transformation. For example, `profile_executor` sounds getting a report of detail execution latencies. Maybe "tune_cutlass_kernels" (similar to AutoTVM/Ansor) or "search_cutlass_kernels" (similar to CuDNN) might be better? Would like to hear to others as well.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737952797
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
Review comment:
IMHO, if we follow the type function in Relay op, we could use either 2nd or 3rd kernels, because in Relay we assume output dtype is always the same as the input dtype, regardless the accumulation dtype.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r733621486
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
Review comment:
@Laurawly I removed the logic in the original cutlass code https://github.com/NVIDIA/cutlass/blob/2e07c4cc2fc94f4cd396ecf1e9132caf1efba50e/tools/library/scripts/generator.py#L809-L817 because I don't understand it (why use `math_inst.element_a` for `ElementC` type etc).
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942488532
Ah sorry I didn't make it clear. The interface of C source codegen does deal with dynamic workloads, because it takes the raw pointer which could be any size in run time.
What I meant was how to generate CUTLASS kernels that are able to perform well with all shapes. In @Laurawly's post, they generate lots of kernels to cover possible shapes, which result in 7GB binary. I assume they also generate a run time dispatching logic (also in the generated C source code) to determine which kernel should be used given the known shape in run time. Obviously, the binary size will definitely be an issue for this solution.
For JSON codegen/runtime, it would be similar to TensorRT: We simply dump a JSON graph in codegen without doing anything else. Meanwhile, we have a custom runtime that JITs/catches CUTLASS kernels based on known shapes. This results in a much smaller binary, but the first execution (or an execution with new shapes) may take several seconds or even a minute to JIT all kernels.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942146086
hmm, are you concerned about slow performance due to lack of tuning, or the integration problem that I brought up? I wouldn't worry about the former because cublas is fast without tuning. So I expect cutlass to perform equally well.
Since our cublas offload supports dynamic inputs, not supporting dynamic inputs for cutlass would be a big bummer imo. So I want to discuss the integration problem first, and investigate performance issues later.
--
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 edited a comment on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942895664
I think it is reasonable to ask users for an estimate of a range of dynamism. Without such estimate, or if the estimate turns out to be off, we just fallback to a default kernel that works on any input shapes, which shouldn't be too slow anyway since it is cutlass, not TVM :) This should be a reasonable trade off between performance and convenience of dynamic models: if they want to run fast on dynamic models, they need to help us, otherwise we cannot guarantee good performance.
I believe even if a network contains many dynamic workloads, there is only one underlying dynamism, from which other dynamic shapes can be derived. For example, MaskRCNN contains some dynamic batch dense, conv2d and conv2d transpose, but all dynamic dims directly correspond to the number of boxes detected, or some numbers derived from it (such as the number of box times * pool window size * feature size, the latter two are fixed numbers). For transformers, I guess the true dynamism is only in sequence length. So I think it wouldn't be too much of an ask to give an estimate of such numbers.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r738730575
##########
File path: python/tvm/contrib/cutlass/build.py
##########
@@ -0,0 +1,172 @@
+# 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
+"""Driver for partitioning and building a Relay module for CUTLASS offload."""
+import tvm
+from tvm import runtime, relay
+from .gen_gemm import CutlassGemmProfiler
+
+
+class GemmAnnotator(tvm.relay.ExprVisitor):
+ """Annotates partitioned functions with shape and dtype information."""
+
+ def __init__(self):
+ super().__init__()
+ self.signature = {}
+
+ def visit_call(self, call):
+ op = call.op
+ if isinstance(op, relay.Function) and "PartitionedFromPattern" in op.attrs:
+ self.signature["op_type"] = op.attrs["Composite"]
+ for i, arg in enumerate(op.params):
+ self.signature["arg%d_shape" % i] = arg.checked_type.shape
+ self.signature["arg%d_dtype" % i] = arg.checked_type.dtype
+ self.signature["ret_shape"] = op.ret_type.shape
+ self.signature["ret_dtype"] = op.ret_type.dtype
+
+
+def profile_cutlass_kernels(mod, sm, profile_all=True, use_multiprocessing=False, tmp_dir="./tmp"):
Review comment:
That's a fair concern and I agree with your plan.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r733619721
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
+
+ ops.extend(out)
+
+ return ops
+
+
+def generate_sm75_tensor_op_1688(out_dtype):
+ """Generate GEMM kernels for Turing."""
+ assert out_dtype in ["float32", "float16"]
+ math_instructions = {
+ "float32": [
+ MathInstruction(
+ [16, 8, 8],
+ DataType.f16,
+ DataType.f16,
+ DataType.f32,
+ OpcodeClass.TensorOp,
+ MathOperation.multiply_add,
+ )
+ ],
+ "float16": [
+ MathInstruction(
+ [16, 8, 8],
+ DataType.f16,
+ DataType.f16,
+ DataType.f16,
+ OpcodeClass.TensorOp,
+ MathOperation.multiply_add,
+ )
+ ],
+ }[out_dtype]
Review comment:
@Laurawly This is one of the things I deviated from the original cutlass code and your initial code. https://github.com/NVIDIA/cutlass/blob/2e07c4cc2fc94f4cd396ecf1e9132caf1efba50e/tools/library/scripts/generator.py#L767-L778
The original code generates two sets of kernels for fp16 and fp32 accumulation. But in our case, generating both of them is not a good idea because
1) We need to compile and tune both set of kernels but we know we only need either fp16 or fp32 one
2) If we want fp32 accum but an fp16 accum kernel is selected by the profiler (which will always happen), the final compilation would fail due to type mismatch in the accumulation related template parameters (some say `float` while other say `cutlass::half_t`)
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737954113
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
Review comment:
Update: I realized that Relay does allow users specify the output dtype of dense ops. Then we have:
1. When output dtype is FP32, we must use the 1st kernel.
2. When output dtype is FP16, we could use either 2nd or 3rd kernel.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942254874
> I think even C source code should be able to handle dynamic shapes as well since it only expects tensors at the runtime
Exactly, I've looked at the code a bit and it seems currently we generate an API that takes `DLTensor`s as inputs, but it passes only raw pointers to the generated backend calls. It is easy to change things a bit to pass shape information as welll.
https://github.com/apache/tvm/blob/3cb838de434f80d584a48a0b41e2a253bb4759ed/src/relay/backend/contrib/codegen_c/codegen_c.h#L176-L185
I'm glad to find at least one path that supports all use cases. For the json codegen/runtime path, I'm not sure how to integrate cutlass's JIT codegen + compile approach with it.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942064270
--
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 edited a comment on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942895664
I think it is reasonable to ask users for an estimate of a range of dynamism. Without such estimate, or if the estimate turns out to be off, we just fallback to a default kernel that works on any input shapes, which shouldn't be too slow anyway since it is cutlass, not TVM :) This should be a reasonable trade off between performance and convenience of dynamic models: if they want to run fast on dynamic models, they need to help us, otherwise we cannot guarantee good performance.
I believe even if a network contains many dynamic workloads, there is only one underlying dynamism, from which other dynamic shapes can be derived. For example, MaskRCNN contains some dynamic batch dense, conv2d and conv2d transpose, but all dynamic dims directly correspond to the number of boxes detected, or some numbers derived from it. For transformers, I guess the true dynamism is only in sequence length. So I think it wouldn't be too much of an ask to give an estimate of such numbers.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r733621486
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
Review comment:
I removed the logic in the original cutlass code https://github.com/NVIDIA/cutlass/blob/2e07c4cc2fc94f4cd396ecf1e9132caf1efba50e/tools/library/scripts/generator.py#L809-L817 because I don't understand it (why use `math_inst.element_a` for `ElementC` type etc).
--
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 edited a comment on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-948543423
This is ready for review. The bulk of PR is a kernel generator taken from cutlass, simplified for our purpose + more epilogue support.
The followings are supported:
* TensorCore kernels (no WMMA or SIMT)
* `dense` with bias, bias + relu, and bias + gelu epilogue
* fp16 and fp32 accumulation
* Kernels for Turing and Ampere
The follow-up work will add more features:
* dynamic input (only basic one, the heuristic stuff is complicated so I defer it for future)
* conv2d and batched gemm
* More epilogues to support e2d models. I'm planning to test MaskRCNN, DeeplabV3, EfficientNet v2, and BERT-large. This require some work to extend cutlass (see https://github.com/NVIDIA/cutlass/discussions/347)
Low priority (for me) but things we should support eventually:
* Heuristic for dynamic inputs
* int8 and int4 kernels
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737955095
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
+
+ ops.extend(out)
+
+ return ops
+
+
+def generate_sm75_tensor_op_1688(out_dtype):
+ """Generate GEMM kernels for Turing."""
+ assert out_dtype in ["float32", "float16"]
+ math_instructions = {
+ "float32": [
+ MathInstruction(
+ [16, 8, 8],
+ DataType.f16,
+ DataType.f16,
+ DataType.f32,
+ OpcodeClass.TensorOp,
+ MathOperation.multiply_add,
+ )
+ ],
+ "float16": [
+ MathInstruction(
+ [16, 8, 8],
+ DataType.f16,
+ DataType.f16,
+ DataType.f16,
+ OpcodeClass.TensorOp,
+ MathOperation.multiply_add,
+ )
+ ],
+ }[out_dtype]
Review comment:
Yes, when the output dtype is fp16, I don't see a good reason to want fp32 accumulation either.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737956211
##########
File path: tests/python/contrib/test_cutlass.py
##########
@@ -0,0 +1,121 @@
+# 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.
+import math
+import pytest
+import tvm
+from tvm import relay
+import numpy as np
+from tvm.contrib.cutlass import profile_and_build
+
+
+def get_ref_rt_mod(mod, params):
+ with tvm.transform.PassContext(opt_level=3):
+ lib = relay.build(mod, target="cuda", params=params)
+ dev = tvm.device("cuda", 0)
+ rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ return rt_mod, dev
+
+
+def get_output(rt_mod, x):
+ rt_mod.set_input("data", x)
+ rt_mod.run()
+ return rt_mod.get_output(0).asnumpy()
+
+
+def get_dense(M, N, K, out_dtype="float16"):
+ data = relay.var("data", shape=(M, K), dtype="float16")
+ weight = relay.var("weight", shape=(N, K), dtype="float16")
+ return relay.nn.dense(data, weight, out_dtype=out_dtype)
+
+
+def get_dense_bias(M, N, K, out_dtype="float16"):
+ dense = get_dense(M, N, K, out_dtype=out_dtype)
+ bias = relay.var("bias", shape=(N,), dtype=out_dtype)
+ return relay.nn.bias_add(dense, bias)
+
+
+def get_dense_bias_relu(M, N, K, out_dtype="float16"):
+ return relay.nn.relu(get_dense_bias(M, N, K, out_dtype="float16"))
+
+
+def get_dense_bias_gelu(M, N, K, out_dtype="float16"):
+ bias_add = get_dense_bias(M, N, K, out_dtype)
+ mul = bias_add * relay.const((1.0 / math.sqrt(2.0)), dtype=out_dtype)
+ if out_dtype == "float16":
+ erf = relay.cast(relay.op.erf(relay.cast(mul, "float32")), "float16")
+ else:
+ erf = relay.op.erf(mul)
+ mul_half = erf * relay.const(0.5, dtype=out_dtype)
+ add = mul_half + relay.const(0.5, dtype=out_dtype)
+ return add * bias_add
+
+
+def verify(func, M, N, K, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False):
+ if not tvm.get_global_func("relay.ext.cutlass", True):
+ return
+ mod = tvm.IRModule.from_expr(func)
+ typ = relay.transform.InferType()(mod)
+ out_dtype = typ["main"].body.checked_type.dtype
+ np_data = np.random.uniform(-1, 1, (M, K)).astype("float16")
+ np_weight = np.random.uniform(-1, 1, (N, K)).astype("float16")
+ np_bias = np.random.uniform(-1, 1, (N,)).astype(out_dtype)
+
+ params = {"weight": np_weight, "bias": np_bias}
+
+ rt_mod_ref, dev = get_ref_rt_mod(mod, params)
+ rt_mod, dev, num_partition = profile_and_build(mod, params, sm, tmp_dir="tmp")
Review comment:
What I'm worrying about is actually the CI time. I believe the tuning time in this case would definitely much shorter than AutoTVM/Ansor, so it should be fine in the real use cases. Maybe we could provide some heuristics. For example:
1. A level/mode indicating the number of candidates to be profiled. We could use the lowest level in the CI, and recommend the highest level for users.
2. Simply let users specify the maximum number of candidates.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737958090
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
Review comment:
I agree. We could just get rid of the 3rd kernel for now if the 2nd kernel doesn't hurt the accuracy.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737937439
##########
File path: tests/python/contrib/test_cutlass.py
##########
@@ -0,0 +1,121 @@
+# 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.
+import math
+import pytest
+import tvm
+from tvm import relay
+import numpy as np
+from tvm.contrib.cutlass import profile_and_build
+
+
+def get_ref_rt_mod(mod, params):
+ with tvm.transform.PassContext(opt_level=3):
+ lib = relay.build(mod, target="cuda", params=params)
+ dev = tvm.device("cuda", 0)
+ rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ return rt_mod, dev
+
+
+def get_output(rt_mod, x):
+ rt_mod.set_input("data", x)
+ rt_mod.run()
+ return rt_mod.get_output(0).asnumpy()
+
+
+def get_dense(M, N, K, out_dtype="float16"):
+ data = relay.var("data", shape=(M, K), dtype="float16")
+ weight = relay.var("weight", shape=(N, K), dtype="float16")
+ return relay.nn.dense(data, weight, out_dtype=out_dtype)
+
+
+def get_dense_bias(M, N, K, out_dtype="float16"):
+ dense = get_dense(M, N, K, out_dtype=out_dtype)
+ bias = relay.var("bias", shape=(N,), dtype=out_dtype)
+ return relay.nn.bias_add(dense, bias)
+
+
+def get_dense_bias_relu(M, N, K, out_dtype="float16"):
+ return relay.nn.relu(get_dense_bias(M, N, K, out_dtype="float16"))
+
+
+def get_dense_bias_gelu(M, N, K, out_dtype="float16"):
+ bias_add = get_dense_bias(M, N, K, out_dtype)
+ mul = bias_add * relay.const((1.0 / math.sqrt(2.0)), dtype=out_dtype)
+ if out_dtype == "float16":
+ erf = relay.cast(relay.op.erf(relay.cast(mul, "float32")), "float16")
+ else:
+ erf = relay.op.erf(mul)
+ mul_half = erf * relay.const(0.5, dtype=out_dtype)
+ add = mul_half + relay.const(0.5, dtype=out_dtype)
+ return add * bias_add
+
+
+def verify(func, M, N, K, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False):
+ if not tvm.get_global_func("relay.ext.cutlass", True):
+ return
+ mod = tvm.IRModule.from_expr(func)
+ typ = relay.transform.InferType()(mod)
+ out_dtype = typ["main"].body.checked_type.dtype
+ np_data = np.random.uniform(-1, 1, (M, K)).astype("float16")
+ np_weight = np.random.uniform(-1, 1, (N, K)).astype("float16")
+ np_bias = np.random.uniform(-1, 1, (N,)).astype(out_dtype)
+
+ params = {"weight": np_weight, "bias": np_bias}
+
+ rt_mod_ref, dev = get_ref_rt_mod(mod, params)
+ rt_mod, dev, num_partition = profile_and_build(mod, params, sm, tmp_dir="tmp")
Review comment:
Just curious, how long would this profile take?
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
+
+ ops.extend(out)
+
+ return ops
+
+
+def generate_sm75_tensor_op_1688(out_dtype):
+ """Generate GEMM kernels for Turing."""
+ assert out_dtype in ["float32", "float16"]
+ math_instructions = {
+ "float32": [
+ MathInstruction(
+ [16, 8, 8],
+ DataType.f16,
+ DataType.f16,
+ DataType.f32,
+ OpcodeClass.TensorOp,
+ MathOperation.multiply_add,
+ )
+ ],
+ "float16": [
+ MathInstruction(
+ [16, 8, 8],
+ DataType.f16,
+ DataType.f16,
+ DataType.f16,
+ OpcodeClass.TensorOp,
+ MathOperation.multiply_add,
+ )
+ ],
+ }[out_dtype]
Review comment:
IMHO, we should not select them simply based on the performance. My considerations are:
1. Currently TVM doesn't really support configurable accumulation dtype, so if we want to align the TVM semantic, we should always use the FP16 one.
2. However, TVM doesn't know what accumulation dtype we are using here, so it may not be an issue for us to choose between FP16 and FP32.
3. If both dtypes won't hurt the inference accuracy, always using FP16 seems no doubt, as I believe its performance will never be worse than FP32.
##########
File path: python/tvm/contrib/cutlass/build.py
##########
@@ -0,0 +1,155 @@
+# 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
+"""Driver for partitioning and building a Relay module for CUTLASS offload."""
+import tvm
+from tvm import runtime, relay
+from tvm.relay.op.contrib.cutlass import partition_for_cutlass
+from .gen_gemm import CutlassGemmProfiler
+
+
+class GemmAnnotator(tvm.relay.ExprVisitor):
+ """Annotates partitioned functions with shape and dtype information."""
+
+ def __init__(self):
+ super().__init__()
+ self.signature = {}
+
+ def visit_call(self, call):
+ op = call.op
+ if isinstance(op, relay.Function) and "PartitionedFromPattern" in op.attrs:
+ self.signature["op_type"] = op.attrs["Composite"]
+ for i, arg in enumerate(op.params):
+ self.signature["arg%d_shape" % i] = arg.checked_type.shape
+ self.signature["arg%d_dtype" % i] = arg.checked_type.dtype
+ self.signature["ret_shape"] = op.ret_type.shape
+ self.signature["ret_dtype"] = op.ret_type.dtype
+
+
+def profile_and_build(mod, params, sm, tmp_dir="./tmp", lib_path="compile.so"):
Review comment:
I am concerned about this API as it wraps too many things and hurts the flexibility. Specifically:
1. It runs partitioning, which means we cannot have the flow like `partition_for_cutlass -> partition_rest_for_trt`.
2. It builds the module, which means we cannot directly access/control the build API (e.g., pass context, tuning the rest ops, etc).
Since we still require users to explicitly run `partition_for_xxx` APIs for each BYOC backend, I would suggest keeping them separate:
```python
mod = partition_for_cutlass(mod)
# users may do something else for the unpartitioned parts here.
mod = tune_for_cutlass(mod) # we could improve this later
with tvm.transform.PassContext(opt_level=3):
# Maybe we still need to wrap relay.build to perform exporting (compile) and load for now...
lib = relay.build(mod, target="cuda", params=params)
```
##########
File path: python/tvm/contrib/cutlass/gemm_operation.py
##########
@@ -0,0 +1,262 @@
+# 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 GEMM kernels."""
+from .library import *
+
+
+class GemmOperation:
+ """Describes various attributes for instantiating GEMM kernels."""
+
+ def __init__(
+ self,
+ arch,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+ ):
+ self.operation_kind = OperationKind.Gemm
+ self.arch = arch
+ self.tile_description = tile_description
+ self.A = A
+ self.B = B
+ self.C = C
+ self.element_epilogue = element_epilogue
+ self.epilogue_functor = epilogue_functor
+ self.swizzling_functor = swizzling_functor
+
+ def accumulator_type(self):
+ return self.tile_description.math_instruction.element_accumulator
+
+ def short_math_name(self):
+ return ShortDataTypeNames[self.accumulator_type()]
+
+ def core_name(self):
+ """ The basic operation kind is prefixed with a letter indicating the accumulation type. """
+ inst_shape = ""
+ intermediate_type = ""
+
+ if (
+ self.tile_description.math_instruction.opcode_class == OpcodeClass.TensorOp
+ or self.tile_description.math_instruction.opcode_class == OpcodeClass.WmmaTensorOp
+ ):
+ inst_shape = "%d%d%d" % tuple(self.tile_description.math_instruction.instruction_shape)
+ if (
+ self.tile_description.math_instruction.element_a != self.A.element
+ and self.tile_description.math_instruction.element_a
+ != self.tile_description.math_instruction.element_accumulator
+ ):
+ intermediate_type = DataTypeNames[self.tile_description.math_instruction.element_a]
+
+ return "%s%s%s%s" % (
+ self.short_math_name(),
+ inst_shape,
+ intermediate_type,
+ "gemm",
+ )
+
+ def extended_name(self):
+ """ Append data types if they differ from compute type. """
+ if (
+ self.C.element != self.tile_description.math_instruction.element_accumulator
+ and self.A.element != self.tile_description.math_instruction.element_accumulator
+ ):
+ extended_name = "${element_c}_${core_name}_${element_a}"
+ elif (
+ self.C.element == self.tile_description.math_instruction.element_accumulator
+ and self.A.element != self.tile_description.math_instruction.element_accumulator
+ ):
+ extended_name = "${core_name}_${element_a}"
+ else:
+ extended_name = "${core_name}"
+
+ extended_name = substitute_template(
+ extended_name,
+ {
+ "element_a": DataTypeNames[self.A.element],
+ "element_c": DataTypeNames[self.C.element],
+ "core_name": self.core_name(),
+ },
+ )
+
+ return extended_name
+
+ def layout_name(self):
+ return "%s%s" % (ShortLayoutTypeNames[self.A.layout], ShortLayoutTypeNames[self.B.layout])
+
+ def procedural_name(self):
+ """The full procedural name indicates architecture, extended name, tile size,
+ and layout.
+ """
+ threadblock = self.tile_description.procedural_name()
+ opcode_class_name = OpcodeClassNames[self.tile_description.math_instruction.opcode_class]
+
+ return substitute_template(
+ "cutlass_${opcode_class}_${extended_name}_${threadblock}_${layout}_align${alignment}",
+ {
+ "opcode_class": opcode_class_name,
+ "extended_name": self.extended_name(),
+ "threadblock": threadblock,
+ "layout": self.layout_name(),
+ "alignment": "%d" % self.A.alignment,
+ },
+ )
+
+ def leading_dim(self):
+ """ lda, ldb, ldc, according to the leading dimension. """
+ if self.A.layout == LayoutType.RowMajor:
+ lda = "K"
+ elif self.A.layout == LayoutType.ColumnMajor:
+ lda = "M"
+ else:
+ ValueError("The layout of A is not implemented.")
+
+ if self.B.layout == LayoutType.RowMajor:
+ ldb = "N"
+ elif self.B.layout == LayoutType.ColumnMajor:
+ ldb = "K"
+ else:
+ ValueError("The layout of B is not implemented.")
+
+ if self.C.layout == LayoutType.RowMajor:
+ ldc = "N"
+ elif self.C.layout == LayoutType.ColumnMajor:
+ ldc = "M"
+ else:
+ ValueError("The layout of B is not implemented.")
+
+ return substitute_template(
+ "int lda = ${lda_val};\n\tint ldb = ${ldb_val};\n\tint ldc = ${ldc_val};\n",
+ {
+ "lda_val": lda,
+ "ldb_val": ldb,
+ "ldc_val": ldc,
+ },
+ )
+
+
+class EmitGemmInstance:
+ """ Responsible for emitting a CUTLASS template definition."""
+
+ def __init__(self):
+ self.epilogue_default = """
+ ${epilogue_functor}<
+ ${element_c},
+ ${epilogue_vector_length},
+ ${element_accumulator},
+ ${element_epilogue}
+ >"""
+ self.epilogue_no_beta_scaling = """
+ ${epilogue_functor}<
+ ${element_c},
+ ${epilogue_vector_length},
+ ${element_accumulator},
+ ${element_epilogue},
+ cutlass::epilogue::thread::ScaleType::NoBetaScaling
+ >"""
+ self.gemm_template = """
+ // Gemm operator ${operation_name}
+ using Operation_${operation_name} = cutlass::gemm::device::Gemm<
+ ${element_a}, ${layout_a},
+ ${element_b}, ${layout_b},
+ ${element_c}, ${layout_c},
+ ${element_accumulator},
+ ${opcode_class},
+ ${arch},
+ cutlass::gemm::GemmShape<${threadblock_shape_m}, ${threadblock_shape_n}, ${threadblock_shape_k}>,
+ cutlass::gemm::GemmShape<${warp_shape_m}, ${warp_shape_n}, ${warp_shape_k}>,
+ cutlass::gemm::GemmShape<${instruction_shape_m}, ${instruction_shape_n}, ${instruction_shape_k}>,
+ ${epilogue},
+ ${swizzling_functor},
+ ${stages},
+ ${align_a},
+ ${align_b},
+ false,
+ ${math_operation}
+ ${residual}
+ >;
+"""
+
+ def emit(self, operation, no_beta_scaling=False):
+ """Instantiate a GEMM kernel from given `operation`."""
+ warp_shape = [
+ operation.tile_description.threadblock_shape[idx]
+ // operation.tile_description.warp_count[idx]
+ for idx in range(3)
+ ]
+ epilogue_vector_length = int(
+ min(operation.C.alignment * DataTypeSize[operation.C.element], 128)
+ / DataTypeSize[operation.C.element]
+ )
Review comment:
nit: You might be able to use `//` to get the integer directly.
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
Review comment:
As the comment indicates, the logic you removed is emitting the kernel which output dtype is same as the first input. I believe this is to generate the kernel that takes two FP16 tensors and outputs one FP16 tensors, but uses FP32 for accumulation. Specifically, with two math instructions:
Inst1: (FP16, FP16, FP32)
Inst2: (FP16, FP16, FP16)
We will only have two kernels without the removed logic,
1. (from Inst1): (FP16, FP16) -> (acc: FP32) -> FP32
2. (from Inst2): (FP16, FP16) -> (acc: FP16) -> FP16
The logic generates the 3rd kernel from the first instruction:
3. (from Inst1): (FP16, FP16) -> (acc: FP32) -> FP16
However, we need the if-guard to prevent the duplicated kernel from the second instruction:
4. (from Inst2): (FP16, FP16) -> (acc: FP16) -> FP16 ... same as 2
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r738289311
##########
File path: tests/python/contrib/test_cutlass.py
##########
@@ -0,0 +1,121 @@
+# 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.
+import math
+import pytest
+import tvm
+from tvm import relay
+import numpy as np
+from tvm.contrib.cutlass import profile_and_build
+
+
+def get_ref_rt_mod(mod, params):
+ with tvm.transform.PassContext(opt_level=3):
+ lib = relay.build(mod, target="cuda", params=params)
+ dev = tvm.device("cuda", 0)
+ rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ return rt_mod, dev
+
+
+def get_output(rt_mod, x):
+ rt_mod.set_input("data", x)
+ rt_mod.run()
+ return rt_mod.get_output(0).asnumpy()
+
+
+def get_dense(M, N, K, out_dtype="float16"):
+ data = relay.var("data", shape=(M, K), dtype="float16")
+ weight = relay.var("weight", shape=(N, K), dtype="float16")
+ return relay.nn.dense(data, weight, out_dtype=out_dtype)
+
+
+def get_dense_bias(M, N, K, out_dtype="float16"):
+ dense = get_dense(M, N, K, out_dtype=out_dtype)
+ bias = relay.var("bias", shape=(N,), dtype=out_dtype)
+ return relay.nn.bias_add(dense, bias)
+
+
+def get_dense_bias_relu(M, N, K, out_dtype="float16"):
+ return relay.nn.relu(get_dense_bias(M, N, K, out_dtype="float16"))
+
+
+def get_dense_bias_gelu(M, N, K, out_dtype="float16"):
+ bias_add = get_dense_bias(M, N, K, out_dtype)
+ mul = bias_add * relay.const((1.0 / math.sqrt(2.0)), dtype=out_dtype)
+ if out_dtype == "float16":
+ erf = relay.cast(relay.op.erf(relay.cast(mul, "float32")), "float16")
+ else:
+ erf = relay.op.erf(mul)
+ mul_half = erf * relay.const(0.5, dtype=out_dtype)
+ add = mul_half + relay.const(0.5, dtype=out_dtype)
+ return add * bias_add
+
+
+def verify(func, M, N, K, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False):
+ if not tvm.get_global_func("relay.ext.cutlass", True):
+ return
+ mod = tvm.IRModule.from_expr(func)
+ typ = relay.transform.InferType()(mod)
+ out_dtype = typ["main"].body.checked_type.dtype
+ np_data = np.random.uniform(-1, 1, (M, K)).astype("float16")
+ np_weight = np.random.uniform(-1, 1, (N, K)).astype("float16")
+ np_bias = np.random.uniform(-1, 1, (N,)).astype(out_dtype)
+
+ params = {"weight": np_weight, "bias": np_bias}
+
+ rt_mod_ref, dev = get_ref_rt_mod(mod, params)
+ rt_mod, dev, num_partition = profile_and_build(mod, params, sm, tmp_dir="tmp")
Review comment:
Added `profile_all` and `use_multiprocessing` options.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r738777119
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
Review comment:
Is that something we can use here? By swapping A and B, for example?
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
Laurawly commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r738786893
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
Review comment:
Yeah, we can also use `(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.RowMajor)` here
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
Laurawly commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r738786893
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
Review comment:
Yeah, we can add `(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.RowMajor)` here
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737956332
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
Review comment:
Yes, I have a test that exercises hetero output dtype, and it works https://github.com/apache/tvm/blob/ff66ef49ecb9d1946511b545fd4d92359f60a813/tests/python/contrib/test_cutlass.py#L102
Although we currently don't generate such op, I'm not ruling out such use cases (e.g. if we want to make our so-called "ToMixedPrecision" pass truly support mixed precision).
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r738290218
##########
File path: python/tvm/contrib/cutlass/build.py
##########
@@ -0,0 +1,155 @@
+# 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
+"""Driver for partitioning and building a Relay module for CUTLASS offload."""
+import tvm
+from tvm import runtime, relay
+from tvm.relay.op.contrib.cutlass import partition_for_cutlass
+from .gen_gemm import CutlassGemmProfiler
+
+
+class GemmAnnotator(tvm.relay.ExprVisitor):
+ """Annotates partitioned functions with shape and dtype information."""
+
+ def __init__(self):
+ super().__init__()
+ self.signature = {}
+
+ def visit_call(self, call):
+ op = call.op
+ if isinstance(op, relay.Function) and "PartitionedFromPattern" in op.attrs:
+ self.signature["op_type"] = op.attrs["Composite"]
+ for i, arg in enumerate(op.params):
+ self.signature["arg%d_shape" % i] = arg.checked_type.shape
+ self.signature["arg%d_dtype" % i] = arg.checked_type.dtype
+ self.signature["ret_shape"] = op.ret_type.shape
+ self.signature["ret_dtype"] = op.ret_type.dtype
+
+
+def profile_and_build(mod, params, sm, tmp_dir="./tmp", lib_path="compile.so"):
Review comment:
Revised the API so that users can use the following flow:
```
mod = partition_for_cutlass(mod)
mod, num_cutlass_partition = profile_cutlass_kernels(
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)
lib = build_cutlass_kernels(lib, sm, tmp_dir, lib_path)
dev = tvm.device("cuda", 0)
rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
```
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737953139
##########
File path: tests/python/contrib/test_cutlass.py
##########
@@ -0,0 +1,121 @@
+# 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.
+import math
+import pytest
+import tvm
+from tvm import relay
+import numpy as np
+from tvm.contrib.cutlass import profile_and_build
+
+
+def get_ref_rt_mod(mod, params):
+ with tvm.transform.PassContext(opt_level=3):
+ lib = relay.build(mod, target="cuda", params=params)
+ dev = tvm.device("cuda", 0)
+ rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ return rt_mod, dev
+
+
+def get_output(rt_mod, x):
+ rt_mod.set_input("data", x)
+ rt_mod.run()
+ return rt_mod.get_output(0).asnumpy()
+
+
+def get_dense(M, N, K, out_dtype="float16"):
+ data = relay.var("data", shape=(M, K), dtype="float16")
+ weight = relay.var("weight", shape=(N, K), dtype="float16")
+ return relay.nn.dense(data, weight, out_dtype=out_dtype)
+
+
+def get_dense_bias(M, N, K, out_dtype="float16"):
+ dense = get_dense(M, N, K, out_dtype=out_dtype)
+ bias = relay.var("bias", shape=(N,), dtype=out_dtype)
+ return relay.nn.bias_add(dense, bias)
+
+
+def get_dense_bias_relu(M, N, K, out_dtype="float16"):
+ return relay.nn.relu(get_dense_bias(M, N, K, out_dtype="float16"))
+
+
+def get_dense_bias_gelu(M, N, K, out_dtype="float16"):
+ bias_add = get_dense_bias(M, N, K, out_dtype)
+ mul = bias_add * relay.const((1.0 / math.sqrt(2.0)), dtype=out_dtype)
+ if out_dtype == "float16":
+ erf = relay.cast(relay.op.erf(relay.cast(mul, "float32")), "float16")
+ else:
+ erf = relay.op.erf(mul)
+ mul_half = erf * relay.const(0.5, dtype=out_dtype)
+ add = mul_half + relay.const(0.5, dtype=out_dtype)
+ return add * bias_add
+
+
+def verify(func, M, N, K, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False):
+ if not tvm.get_global_func("relay.ext.cutlass", True):
+ return
+ mod = tvm.IRModule.from_expr(func)
+ typ = relay.transform.InferType()(mod)
+ out_dtype = typ["main"].body.checked_type.dtype
+ np_data = np.random.uniform(-1, 1, (M, K)).astype("float16")
+ np_weight = np.random.uniform(-1, 1, (N, K)).astype("float16")
+ np_bias = np.random.uniform(-1, 1, (N,)).astype(out_dtype)
+
+ params = {"weight": np_weight, "bias": np_bias}
+
+ rt_mod_ref, dev = get_ref_rt_mod(mod, params)
+ rt_mod, dev, num_partition = profile_and_build(mod, params, sm, tmp_dir="tmp")
Review comment:
We might want to compile all executables on all cores, rather than repeating compile -> profile sequentially. What do you think? cc @Laurawly
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r728768117
##########
File path: .gitmodules
##########
@@ -13,3 +13,6 @@
[submodule "3rdparty/libbacktrace"]
path = 3rdparty/libbacktrace
url = https://github.com/tlc-pack/libbacktrace.git
+[submodule "cutlass_profiler/cutlass"]
+ path = cutlass_profiler/cutlass
+ url = git@github.com:NVIDIA/cutlass.git
Review comment:
done
--
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] junrushao1994 commented on a change in pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r727471035
##########
File path: .gitmodules
##########
@@ -13,3 +13,6 @@
[submodule "3rdparty/libbacktrace"]
path = 3rdparty/libbacktrace
url = https://github.com/tlc-pack/libbacktrace.git
+[submodule "cutlass_profiler/cutlass"]
+ path = cutlass_profiler/cutlass
+ url = git@github.com:NVIDIA/cutlass.git
Review comment:
Shall we move this folder to 3rdparty as our convention in the repo?
--
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] junrushao1994 commented on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-941462765
Really cool work @Laurawly!
--
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] junrushao1994 commented on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
junrushao1994 commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-941462765
Really cool work @Laurawly!
--
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 edited a comment on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942254874
> I think even C source code should be able to handle dynamic shapes as well since it only expects tensors at the runtime
Exactly, I've looked at the code a bit and it seems currently we generate an API that takes `DLTensor`s as inputs, but it passes only raw pointers to the generated backend calls. It is easy to change things a bit to pass shape information as well.
https://github.com/apache/tvm/blob/3cb838de434f80d584a48a0b41e2a253bb4759ed/src/relay/backend/contrib/codegen_c/codegen_c.h#L176-L185
I'm glad to find at least one path that supports all use cases. For the json codegen/runtime path, I'm not sure how to integrate cutlass's JIT codegen + compile approach with it.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942877261
> I agree that kernel picking should be handled by the custom runtime.
It would be interesting if we can apply the same logic that AutoTIR would develop soon, since they need to solve the same (or harder) problem anyway. I believe most of our dynamic input use cases have dynamism only in the first (batch) dimension, so things shouldn't be complicated.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r733609880
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
Review comment:
`ColumnMajor` for the transposed weight in our `dense` 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] masahi commented on a change in pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r733622480
##########
File path: python/tvm/contrib/cutlass/build.py
##########
@@ -0,0 +1,155 @@
+# 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
+"""Driver for partitioning and building a Relay module for CUTLASS offload."""
+import tvm
+from tvm import runtime, relay
+from tvm.relay.op.contrib.cutlass import partition_for_cutlass
+from .gen_gemm import CutlassGemmProfiler
+
+
+class GemmAnnotator(tvm.relay.ExprVisitor):
+ """Annotates partitioned functions with shape and dtype information."""
+
+ def __init__(self):
+ super().__init__()
+ self.signature = {}
+
+ def visit_call(self, call):
+ op = call.op
+ if isinstance(op, relay.Function) and "PartitionedFromPattern" in op.attrs:
+ self.signature["op_type"] = op.attrs["Composite"]
+ for i, arg in enumerate(op.params):
+ self.signature["arg%d_shape" % i] = arg.checked_type.shape
+ self.signature["arg%d_dtype" % i] = arg.checked_type.dtype
+ self.signature["ret_shape"] = op.ret_type.shape
+ self.signature["ret_dtype"] = op.ret_type.dtype
+
+
+def profile_and_build(mod, params, sm, tmp_dir="./tmp", lib_path="compile.so"):
Review comment:
Suggestions for a better name and API welcome. I didn't put deep thought into it.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r738777119
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
Review comment:
Is that something we can use here, given that the input layout is always fixed? Maybe we can swap A and B?
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-954352449
Thanks @Laurawly @comaniac @junrushao1994 @zhiics @hwu36 this is merged!
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737956211
##########
File path: tests/python/contrib/test_cutlass.py
##########
@@ -0,0 +1,121 @@
+# 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.
+import math
+import pytest
+import tvm
+from tvm import relay
+import numpy as np
+from tvm.contrib.cutlass import profile_and_build
+
+
+def get_ref_rt_mod(mod, params):
+ with tvm.transform.PassContext(opt_level=3):
+ lib = relay.build(mod, target="cuda", params=params)
+ dev = tvm.device("cuda", 0)
+ rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ return rt_mod, dev
+
+
+def get_output(rt_mod, x):
+ rt_mod.set_input("data", x)
+ rt_mod.run()
+ return rt_mod.get_output(0).asnumpy()
+
+
+def get_dense(M, N, K, out_dtype="float16"):
+ data = relay.var("data", shape=(M, K), dtype="float16")
+ weight = relay.var("weight", shape=(N, K), dtype="float16")
+ return relay.nn.dense(data, weight, out_dtype=out_dtype)
+
+
+def get_dense_bias(M, N, K, out_dtype="float16"):
+ dense = get_dense(M, N, K, out_dtype=out_dtype)
+ bias = relay.var("bias", shape=(N,), dtype=out_dtype)
+ return relay.nn.bias_add(dense, bias)
+
+
+def get_dense_bias_relu(M, N, K, out_dtype="float16"):
+ return relay.nn.relu(get_dense_bias(M, N, K, out_dtype="float16"))
+
+
+def get_dense_bias_gelu(M, N, K, out_dtype="float16"):
+ bias_add = get_dense_bias(M, N, K, out_dtype)
+ mul = bias_add * relay.const((1.0 / math.sqrt(2.0)), dtype=out_dtype)
+ if out_dtype == "float16":
+ erf = relay.cast(relay.op.erf(relay.cast(mul, "float32")), "float16")
+ else:
+ erf = relay.op.erf(mul)
+ mul_half = erf * relay.const(0.5, dtype=out_dtype)
+ add = mul_half + relay.const(0.5, dtype=out_dtype)
+ return add * bias_add
+
+
+def verify(func, M, N, K, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False):
+ if not tvm.get_global_func("relay.ext.cutlass", True):
+ return
+ mod = tvm.IRModule.from_expr(func)
+ typ = relay.transform.InferType()(mod)
+ out_dtype = typ["main"].body.checked_type.dtype
+ np_data = np.random.uniform(-1, 1, (M, K)).astype("float16")
+ np_weight = np.random.uniform(-1, 1, (N, K)).astype("float16")
+ np_bias = np.random.uniform(-1, 1, (N,)).astype(out_dtype)
+
+ params = {"weight": np_weight, "bias": np_bias}
+
+ rt_mod_ref, dev = get_ref_rt_mod(mod, params)
+ rt_mod, dev, num_partition = profile_and_build(mod, params, sm, tmp_dir="tmp")
Review comment:
What I'm worrying about is actually the CI time. I believe the tuning time in this case would definitely be much shorter than AutoTVM/Ansor, so it should be fine in the real use cases. Maybe we could provide some heuristics to control the CI time. For example:
1. A level/mode indicating the number of candidates to be profiled. We could use the lowest level in the CI, and recommend the highest level for users.
2. Simply let users specify the maximum number of candidates.
3. A special flag. When true, we will stop right after we found the first available kernel regardless its performance.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
comaniac commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737953742
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
+ ]
+ for math_inst in math_instructions:
+ tile_descriptions = get_tile_descriptions(math_inst)
+ data_type = [
+ math_inst.element_a,
+ math_inst.element_b,
+ math_inst.element_accumulator,
+ math_inst.element_accumulator,
+ ]
+
+ out = create_gemm_operator(layouts, tile_descriptions, data_type, alignment_constraints)
+
+ ops.extend(out)
+
+ return ops
+
+
+def generate_sm75_tensor_op_1688(out_dtype):
+ """Generate GEMM kernels for Turing."""
+ assert out_dtype in ["float32", "float16"]
+ math_instructions = {
+ "float32": [
+ MathInstruction(
+ [16, 8, 8],
+ DataType.f16,
+ DataType.f16,
+ DataType.f32,
+ OpcodeClass.TensorOp,
+ MathOperation.multiply_add,
+ )
+ ],
+ "float16": [
+ MathInstruction(
+ [16, 8, 8],
+ DataType.f16,
+ DataType.f16,
+ DataType.f16,
+ OpcodeClass.TensorOp,
+ MathOperation.multiply_add,
+ )
+ ],
+ }[out_dtype]
Review comment:
Ah I see. Maybe I have a wrong impression. So we could let users specify the output dtype in Relay.
Yeah then when output dtype if FP32, we have no choice. When output dtype is FP16, we still have two options for accumulation dtype, and this is what I my previous comment about.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r737952736
##########
File path: tests/python/contrib/test_cutlass.py
##########
@@ -0,0 +1,121 @@
+# 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.
+import math
+import pytest
+import tvm
+from tvm import relay
+import numpy as np
+from tvm.contrib.cutlass import profile_and_build
+
+
+def get_ref_rt_mod(mod, params):
+ with tvm.transform.PassContext(opt_level=3):
+ lib = relay.build(mod, target="cuda", params=params)
+ dev = tvm.device("cuda", 0)
+ rt_mod = tvm.contrib.graph_executor.GraphModule(lib["default"](dev))
+ return rt_mod, dev
+
+
+def get_output(rt_mod, x):
+ rt_mod.set_input("data", x)
+ rt_mod.run()
+ return rt_mod.get_output(0).asnumpy()
+
+
+def get_dense(M, N, K, out_dtype="float16"):
+ data = relay.var("data", shape=(M, K), dtype="float16")
+ weight = relay.var("weight", shape=(N, K), dtype="float16")
+ return relay.nn.dense(data, weight, out_dtype=out_dtype)
+
+
+def get_dense_bias(M, N, K, out_dtype="float16"):
+ dense = get_dense(M, N, K, out_dtype=out_dtype)
+ bias = relay.var("bias", shape=(N,), dtype=out_dtype)
+ return relay.nn.bias_add(dense, bias)
+
+
+def get_dense_bias_relu(M, N, K, out_dtype="float16"):
+ return relay.nn.relu(get_dense_bias(M, N, K, out_dtype="float16"))
+
+
+def get_dense_bias_gelu(M, N, K, out_dtype="float16"):
+ bias_add = get_dense_bias(M, N, K, out_dtype)
+ mul = bias_add * relay.const((1.0 / math.sqrt(2.0)), dtype=out_dtype)
+ if out_dtype == "float16":
+ erf = relay.cast(relay.op.erf(relay.cast(mul, "float32")), "float16")
+ else:
+ erf = relay.op.erf(mul)
+ mul_half = erf * relay.const(0.5, dtype=out_dtype)
+ add = mul_half + relay.const(0.5, dtype=out_dtype)
+ return add * bias_add
+
+
+def verify(func, M, N, K, sm=80, atol=1e-5, rtol=1e-5, run_benchmark=False):
+ if not tvm.get_global_func("relay.ext.cutlass", True):
+ return
+ mod = tvm.IRModule.from_expr(func)
+ typ = relay.transform.InferType()(mod)
+ out_dtype = typ["main"].body.checked_type.dtype
+ np_data = np.random.uniform(-1, 1, (M, K)).astype("float16")
+ np_weight = np.random.uniform(-1, 1, (N, K)).astype("float16")
+ np_bias = np.random.uniform(-1, 1, (N,)).astype(out_dtype)
+
+ params = {"weight": np_weight, "bias": np_bias}
+
+ rt_mod_ref, dev = get_ref_rt_mod(mod, params)
+ rt_mod, dev, num_partition = profile_and_build(mod, params, sm, tmp_dir="tmp")
Review comment:
I can give the detailed numbers later, but roughly, compiling different profiler executables for each candidate kernels take a lot of time. Once compiled, actual profiling is very quick.
We compile all executables when we hit the first cutlass partition. After that, we can use the same compiled executables for profiling different workloads, so only profiling the first partition takes a lot of time.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r738714042
##########
File path: python/tvm/contrib/cutlass/build.py
##########
@@ -0,0 +1,172 @@
+# 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
+"""Driver for partitioning and building a Relay module for CUTLASS offload."""
+import tvm
+from tvm import runtime, relay
+from .gen_gemm import CutlassGemmProfiler
+
+
+class GemmAnnotator(tvm.relay.ExprVisitor):
+ """Annotates partitioned functions with shape and dtype information."""
+
+ def __init__(self):
+ super().__init__()
+ self.signature = {}
+
+ def visit_call(self, call):
+ op = call.op
+ if isinstance(op, relay.Function) and "PartitionedFromPattern" in op.attrs:
+ self.signature["op_type"] = op.attrs["Composite"]
+ for i, arg in enumerate(op.params):
+ self.signature["arg%d_shape" % i] = arg.checked_type.shape
+ self.signature["arg%d_dtype" % i] = arg.checked_type.dtype
+ self.signature["ret_shape"] = op.ret_type.shape
+ self.signature["ret_dtype"] = op.ret_type.dtype
+
+
+def profile_cutlass_kernels(mod, sm, profile_all=True, use_multiprocessing=False, tmp_dir="./tmp"):
Review comment:
ok I'll change to `tune_cutlass_kernels`. `profile` is what cutlass uses in their repo (e.g. `cutlass_profiler`), but yeah, might be confusing to TVM users.
My only concern with `tune` is, people might get a wrong impression that `tune_cutlass_kernels` is an optional thing, since in Ansor / AutoTVM tuning is optional. We can add "default kernels" later to avoid having tuning as a hard requirement. We need default kernels for dynamic workload anyway.
--
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 #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
Laurawly commented on a change in pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#discussion_r738786893
##########
File path: python/tvm/contrib/cutlass/gen_gemm.py
##########
@@ -0,0 +1,335 @@
+# 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
+"""Kernel generator and profiler for CUTLASS."""
+import os
+import re
+import tempfile
+import subprocess
+from .gemm_operation import GemmOperation, EmitGemmInstance
+from .gemm_profiler import GemmProfilerEmitter
+from .library import (
+ EpilogueFunctor,
+ SwizzlingFunctor,
+ TensorDescription,
+ DataTypeTag,
+ LayoutType,
+ MathInstruction,
+ DataType,
+ OpcodeClass,
+ MathOperation,
+ TileDescription,
+)
+
+
+def create_gemm_operator(
+ layouts,
+ tile_descriptions,
+ data_type,
+ alignment_constraints,
+ epilogue_functor=EpilogueFunctor.LinearCombination,
+ swizzling_functor=SwizzlingFunctor.Identity8,
+):
+ """Exhaustively instantiate all kernels from a given configuration."""
+ ret = []
+ kernel_emitter = EmitGemmInstance()
+ profiler_emitter = GemmProfilerEmitter()
+
+ element_a, element_b, element_c, element_epilogue = data_type
+
+ for layout in layouts:
+ for tile_description 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)
+
+ op_entry = {}
+ op = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ epilogue_functor,
+ swizzling_functor,
+ )
+ op_bias = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationBias,
+ swizzling_functor,
+ )
+ op_bias_relu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationRelu,
+ swizzling_functor,
+ )
+ op_bias_gelu = GemmOperation(
+ tile_description.minimum_compute_capability,
+ tile_description,
+ A,
+ B,
+ C,
+ element_epilogue,
+ EpilogueFunctor.LinearCombinationGelu,
+ swizzling_functor,
+ )
+
+ kernel_emitter = EmitGemmInstance()
+ op_entry["op"] = op
+ op_entry["name"] = op.procedural_name()
+ op_entry["opdef"] = kernel_emitter.emit(op)
+ op_entry["opdef_bias"] = kernel_emitter.emit(op_bias, no_beta_scaling=True)
+ op_entry["opdef_bias_relu"] = kernel_emitter.emit(
+ op_bias_relu, no_beta_scaling=True
+ )
+ op_entry["opdef_bias_gelu"] = kernel_emitter.emit(op_bias_gelu)
+ op_entry["src"] = profiler_emitter.emit(
+ op.procedural_name(),
+ op_entry["opdef"],
+ DataTypeTag[element_a],
+ DataTypeTag[element_b],
+ DataTypeTag[element_c],
+ op.leading_dim(),
+ )
+ op_entry["runtime"] = 9999999
+ ret.append(op_entry)
+ return ret
+
+
+def generate_tensor_op_common(math_instructions, alignment_constraints, get_tile_descriptions):
+ """Common kernel generator to be used by archtecture specific generators."""
+ ops = []
+ layouts = [
+ (LayoutType.RowMajor, LayoutType.ColumnMajor, LayoutType.RowMajor),
Review comment:
Yeah, we can also add `(LayoutType.ColumnMajor, LayoutType.RowMajor, LayoutType.RowMajor)` here
--
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] tqchen commented on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
tqchen commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-946239057
given we are bringing in cutlass as a submodule, please update
- https://github.com/apache/tvm/tree/main/licenses
- https://github.com/apache/tvm/blob/main/LICENSE#L240 (I believe cutlass is BSD 3-clause License)
To reflect cutlass license,
--
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 edited a comment on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
masahi edited a comment on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-941990346
One question for @comaniac: Currently, this integration is implemented with C source backend. But one big use case for cutlass BYOC is for dynamic workloads, for which we cannot do well at the moment. I think C source backend is not compatible with dynamic inputs, is that right? On the other hand, with json codegen/runtime, I can see a way to support dynamic inputs.
However, since cutlass is a template library, the only way it could work with the json codegen/runtime is to instantiate all possible templates that we care about ahead of time, and build them together with `libtvm.so`. I wonder if that is feasible. cc @Laurawly
One benefit of the above approach is that we don't need to commit the cutlass code generator into our repo, which is a big part of this PR.
--
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] hwu36 commented on pull request #9261: [BYOC] CUTLASS integration
Posted by GitBox <gi...@apache.org>.
hwu36 commented on pull request #9261:
URL: https://github.com/apache/tvm/pull/9261#issuecomment-942558026
> I assume they also generate a run time dispatching logic (also in the generated C source code) to determine which kernel should be used given the known shape in run time.
At this moment, CUTLASS does not provide heuristics to choose which kernel to use based on the runtime information. The user of CUTLASS needs to make the decision to pick kernels.
--
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