chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,392 @@
|
||||
# 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, dangerous-default-value
|
||||
# ruff: noqa: E501
|
||||
"""Conv2d kernel generator and profiler for CUTLASS."""
|
||||
|
||||
import os
|
||||
import pickle
|
||||
from functools import partial
|
||||
|
||||
from .conv2d_operation import Conv2dOperation, EmitConv2dInstance
|
||||
from .conv2d_profiler import Conv2dProfilerEmitter
|
||||
from .gen_gemm import CutlassGemmProfiler
|
||||
from .gen_tensor_op import EPILOGUE_MAP, GENERATOR_FUNC_TABLE, ProfilerEngine
|
||||
from .library import (
|
||||
ConvKind,
|
||||
DataType,
|
||||
EpilogueFunctor,
|
||||
IteratorAlgorithm,
|
||||
LayoutType,
|
||||
StrideSupport,
|
||||
SwizzlingFunctor,
|
||||
TensorDescription,
|
||||
)
|
||||
|
||||
|
||||
def create_conv2d_operator_with_epilogue(
|
||||
conv_kind,
|
||||
stride_support,
|
||||
op_type,
|
||||
tile_description,
|
||||
data_type,
|
||||
alignment,
|
||||
alignment_epilogue,
|
||||
swizzling_functor,
|
||||
split_k_slices,
|
||||
):
|
||||
"""
|
||||
Instantiate a cutlass kernel from the given configuration,
|
||||
along with the epilouge functor
|
||||
"""
|
||||
if "residual" in op_type:
|
||||
activation_map = {
|
||||
"cutlass.conv2d_bias_hardswish": "cutlass::epilogue::thread::HardSwish",
|
||||
"cutlass.conv2d_bias_silu": "cutlass::epilogue::thread::SiLu",
|
||||
"cutlass.conv2d_bias_sigmoid": "cutlass::epilogue::thread::Sigmoid",
|
||||
"cutlass.conv2d_bias_relu": "cutlass::epilogue::thread::ReLu",
|
||||
"cutlass.conv2d_bias": "cutlass::epilogue::thread::Identity",
|
||||
}
|
||||
prefix = op_type[: op_type.find("_residual")]
|
||||
activation = activation_map[prefix]
|
||||
binary_op = "cutlass::multiplies" if "residual_multiply" in op_type else "cutlass::plus"
|
||||
unary_op = (
|
||||
"cutlass::epilogue::thread::ReLu"
|
||||
if op_type.endswith("relu")
|
||||
else "cutlass::epilogue::thread::Identity"
|
||||
)
|
||||
residual_block_info = {
|
||||
"activation": activation,
|
||||
"binary_op": binary_op,
|
||||
"unary_op": unary_op,
|
||||
}
|
||||
epilogue = EpilogueFunctor.LinearCombinationResidualBlock
|
||||
no_beta_scaling = False
|
||||
else:
|
||||
residual_block_info = None
|
||||
epilogue, no_beta_scaling = EPILOGUE_MAP[op_type]
|
||||
|
||||
element_a, element_b, element_c, element_epilogue = data_type
|
||||
|
||||
A = TensorDescription(element_a, LayoutType.TensorNHWC, alignment)
|
||||
B = TensorDescription(element_b, LayoutType.TensorNHWC, alignment)
|
||||
C = TensorDescription(element_c, LayoutType.TensorNHWC, alignment_epilogue)
|
||||
|
||||
op = Conv2dOperation(
|
||||
conv_kind,
|
||||
IteratorAlgorithm.Optimized,
|
||||
tile_description.minimum_compute_capability,
|
||||
tile_description,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
element_epilogue,
|
||||
stride_support,
|
||||
epilogue,
|
||||
swizzling_functor,
|
||||
split_k_slices,
|
||||
)
|
||||
|
||||
name = op.procedural_name()
|
||||
opdef = EmitConv2dInstance().emit(
|
||||
op,
|
||||
no_beta_scaling=no_beta_scaling,
|
||||
residual_block_info=residual_block_info,
|
||||
emit_reduction=split_k_slices > 1,
|
||||
)
|
||||
|
||||
return name, opdef
|
||||
|
||||
|
||||
def enumerate_conv2d_operators(
|
||||
conv_kind,
|
||||
stride_support,
|
||||
split_k_slices,
|
||||
alignment_c,
|
||||
tile_descriptions,
|
||||
data_type,
|
||||
alignment_constraints,
|
||||
swizzling_functor=SwizzlingFunctor.Identity4,
|
||||
):
|
||||
"""Exhaustively instantiate all kernels from a given configuration."""
|
||||
ret = []
|
||||
|
||||
kernel_emitter = EmitConv2dInstance()
|
||||
profiler_emitter = Conv2dProfilerEmitter()
|
||||
|
||||
element_a, element_b, element_c, element_epilogue = data_type
|
||||
|
||||
if conv_kind == ConvKind.Dgrad and stride_support == StrideSupport.Strided:
|
||||
swizzling_functor = SwizzlingFunctor.StridedDgradIdentity1
|
||||
|
||||
for split_k_slice in split_k_slices:
|
||||
for tile in tile_descriptions:
|
||||
for alignmentAB in alignment_constraints:
|
||||
for alignmentC in alignment_c:
|
||||
A = TensorDescription(element_a, LayoutType.TensorNHWC, alignmentAB)
|
||||
B = TensorDescription(element_b, LayoutType.TensorNHWC, alignmentAB)
|
||||
C = TensorDescription(element_c, LayoutType.TensorNHWC, alignmentC)
|
||||
|
||||
if element_c == DataType.s32 and A.alignment == 1:
|
||||
tile.threadblock_shape[0] = min(tile.threadblock_shape[0], 128)
|
||||
tile.threadblock_shape[1] = min(tile.threadblock_shape[1], 128)
|
||||
|
||||
op = Conv2dOperation(
|
||||
conv_kind,
|
||||
IteratorAlgorithm.Optimized,
|
||||
tile.minimum_compute_capability,
|
||||
tile,
|
||||
A,
|
||||
B,
|
||||
C,
|
||||
element_epilogue,
|
||||
stride_support,
|
||||
EpilogueFunctor.LinearCombination,
|
||||
swizzling_functor,
|
||||
split_k_slice,
|
||||
)
|
||||
|
||||
ret.append(
|
||||
{
|
||||
"src": profiler_emitter.emit(
|
||||
kernel_emitter.emit(op, emit_reduction=split_k_slice > 1),
|
||||
op.procedural_name(),
|
||||
element_output=element_c,
|
||||
split_k_slices=split_k_slice,
|
||||
),
|
||||
"name": op.procedural_name(),
|
||||
"tile_description": tile,
|
||||
"alignment": alignmentAB,
|
||||
"alignment_epilogue": alignmentC,
|
||||
"data_type": data_type,
|
||||
"swizzle_functor": swizzling_functor,
|
||||
"split_k_slices": split_k_slice,
|
||||
}
|
||||
)
|
||||
|
||||
return ret
|
||||
|
||||
|
||||
class CutlassConv2DProfiler:
|
||||
"""Profile all candidate kernels and select the best one."""
|
||||
|
||||
def __init__(self, sm, cutlass_path, binary_path):
|
||||
self.gemm_profiler = CutlassGemmProfiler(sm, cutlass_path, binary_path)
|
||||
self.sm = sm
|
||||
assert sm in GENERATOR_FUNC_TABLE, f"sm{sm} not supported yet."
|
||||
self.engine = ProfilerEngine(sm, cutlass_path, binary_path)
|
||||
self.cache_path = os.path.join(binary_path, "cutlass_conv2d_cache.pickle")
|
||||
if os.path.exists(self.cache_path):
|
||||
self.cache = pickle.load(open(self.cache_path, "rb"))
|
||||
else:
|
||||
self.cache = {}
|
||||
|
||||
def get_default(
|
||||
self,
|
||||
op_type,
|
||||
out_dtype,
|
||||
arg0_dtype,
|
||||
arg1_dtype,
|
||||
use_3xtf32,
|
||||
conv_kind=ConvKind.Fprop,
|
||||
stride=(1, 1),
|
||||
):
|
||||
"""Return the default kernel for the requested architecture.
|
||||
For now, the default kernel was picked arbitrary.
|
||||
"""
|
||||
gemm_profile_result = self.gemm_profiler.get_default(
|
||||
op_type, out_dtype, arg0_dtype, arg1_dtype, use_3xtf32
|
||||
)
|
||||
tile_description = gemm_profile_result["tile_description"]
|
||||
alignment = gemm_profile_result["alignment"]
|
||||
data_type = gemm_profile_result["data_type"]
|
||||
stride_support = StrideSupport.Strided if stride[0] > 1 else StrideSupport.Unity
|
||||
|
||||
if conv_kind == ConvKind.Dgrad and stride_support == StrideSupport.Strided:
|
||||
swizzling_functor = SwizzlingFunctor.StridedDgradIdentity1
|
||||
else:
|
||||
swizzling_functor = SwizzlingFunctor.Identity4
|
||||
|
||||
name, opdef = create_conv2d_operator_with_epilogue(
|
||||
conv_kind,
|
||||
stride_support,
|
||||
op_type,
|
||||
tile_description,
|
||||
data_type,
|
||||
alignment,
|
||||
alignment,
|
||||
swizzling_functor,
|
||||
split_k_slices=1,
|
||||
)
|
||||
return {"name": name, "opdef": opdef}
|
||||
|
||||
def select_op(
|
||||
self,
|
||||
d_shape,
|
||||
w_shape,
|
||||
padding,
|
||||
stride,
|
||||
dilation,
|
||||
out_dtype,
|
||||
data_dtype,
|
||||
weight_dtype,
|
||||
use_3xtf32,
|
||||
conv_kind,
|
||||
stride_support,
|
||||
split_k_slices,
|
||||
profile_all_alignments=False,
|
||||
find_first_valid=False,
|
||||
use_multiprocessing=False,
|
||||
):
|
||||
"""
|
||||
Profile and select the best kernel from candidate kernels.
|
||||
See the documentation for the profile method below.
|
||||
"""
|
||||
N, H, W, IC = d_shape
|
||||
OC, R, S, _ = w_shape
|
||||
|
||||
workload = (
|
||||
N,
|
||||
H,
|
||||
W,
|
||||
IC,
|
||||
OC,
|
||||
R,
|
||||
S,
|
||||
padding[0],
|
||||
padding[1],
|
||||
stride[0],
|
||||
stride[1],
|
||||
dilation[0],
|
||||
dilation[1],
|
||||
)
|
||||
|
||||
if workload in self.cache:
|
||||
return self.cache[workload]
|
||||
|
||||
def alignments(dtype):
|
||||
if dtype in ["float16"]:
|
||||
alignments = [8, 4, 2, 1]
|
||||
elif dtype in ["float", "float32"]:
|
||||
alignments = [4, 2, 1]
|
||||
else:
|
||||
raise ValueError(f"Unsupported data type: {dtype}")
|
||||
return alignments
|
||||
|
||||
alignments_c = [align for align in alignments(out_dtype) if OC % align == 0]
|
||||
|
||||
if not profile_all_alignments:
|
||||
alignments_c = [alignments_c[0]]
|
||||
|
||||
ops = GENERATOR_FUNC_TABLE[self.sm](
|
||||
out_dtype,
|
||||
data_dtype,
|
||||
weight_dtype,
|
||||
partial(
|
||||
enumerate_conv2d_operators,
|
||||
conv_kind,
|
||||
stride_support,
|
||||
split_k_slices,
|
||||
alignments_c,
|
||||
),
|
||||
lambda align: all([dim % align == 0 for dim in [IC]]),
|
||||
use_3xtf32,
|
||||
profile_all_alignments,
|
||||
# Use fp32 accumulation for wgrad to align with cuDNN
|
||||
accumlator_dtype="float32" if conv_kind == ConvKind.Wgrad else out_dtype,
|
||||
)
|
||||
|
||||
if not find_first_valid:
|
||||
self.engine.compile_all(ops, use_multiprocessing)
|
||||
|
||||
args = "--n={} --h={} --w={} --c={} --k={} --r={} --s={} --pad_h={} --pad_w={} --stride_h={} --stride_w={} --dilation_h={} --dilation_w={}".format(
|
||||
*workload
|
||||
)
|
||||
|
||||
for op in ops:
|
||||
out = self.engine.evaluate(op, args.split(" "))
|
||||
op["runtime"] = out
|
||||
if out < float("inf") and find_first_valid:
|
||||
self.cache[workload] = op
|
||||
return op
|
||||
|
||||
op = min(ops, key=lambda i: i["runtime"])
|
||||
self.cache[workload] = op
|
||||
with open(self.cache_path, "wb") as f:
|
||||
pickle.dump(self.cache, f)
|
||||
return op
|
||||
|
||||
def profile(
|
||||
self,
|
||||
op_type,
|
||||
d_shape,
|
||||
w_shape,
|
||||
padding,
|
||||
stride,
|
||||
dilation,
|
||||
out_dtype,
|
||||
data_dtype,
|
||||
weight_dtype,
|
||||
use_3xtf32=True,
|
||||
conv_kind=ConvKind.Fprop,
|
||||
split_k_slices=[1],
|
||||
profile_all_alignments=False,
|
||||
find_first_valid=False,
|
||||
use_multiprocessing=False,
|
||||
):
|
||||
"""Profile and select the best kernel from candidate kernels.
|
||||
If find_first_valid is True, return immediately after the first applicable kernel is found.
|
||||
If use_multiprocessing is True, compile all profiler executables in parallel.
|
||||
"""
|
||||
# Dgrad requires Unity stride when stride == (1, 1)
|
||||
stride_support = (
|
||||
StrideSupport.Unity
|
||||
if stride[0] == 1 and stride[1] == 1 and conv_kind == ConvKind.Dgrad
|
||||
else StrideSupport.Strided
|
||||
)
|
||||
|
||||
op = self.select_op(
|
||||
d_shape,
|
||||
w_shape,
|
||||
padding,
|
||||
stride,
|
||||
dilation,
|
||||
out_dtype,
|
||||
data_dtype,
|
||||
weight_dtype,
|
||||
use_3xtf32,
|
||||
conv_kind,
|
||||
stride_support,
|
||||
split_k_slices,
|
||||
profile_all_alignments,
|
||||
find_first_valid,
|
||||
use_multiprocessing,
|
||||
)
|
||||
|
||||
name, opdef = create_conv2d_operator_with_epilogue(
|
||||
conv_kind,
|
||||
stride_support,
|
||||
op_type,
|
||||
op["tile_description"],
|
||||
op["data_type"],
|
||||
op["alignment"],
|
||||
op["alignment_epilogue"],
|
||||
op["swizzle_functor"],
|
||||
op["split_k_slices"],
|
||||
)
|
||||
|
||||
return name, opdef, op["runtime"]
|
||||
Reference in New Issue
Block a user