Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[TOPI] Use f-strings for string formatting, NFC #14822

Merged
merged 2 commits into from
May 11, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 3 additions & 12 deletions python/tvm/topi/arm_cpu/conv2d.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,10 +33,7 @@
schedule_conv2d_spatial_pack_nchw,
schedule_conv2d_spatial_pack_nhwc,
)
from .mprofile.dsp.conv2d import (
conv2d_nhwc_dsp_compute,
conv2d_nhwc_dsp_schedule,
)
from .mprofile.dsp.conv2d import conv2d_nhwc_dsp_compute, conv2d_nhwc_dsp_schedule


@autotvm.register_topi_compute("conv2d_nchw_spatial_pack.arm_cpu")
Expand Down Expand Up @@ -267,13 +264,7 @@ def _schedule_winograd(cfg, s, output, last):
if isinstance(U.op, tvm.te.ComputeOp):
kernel, G = U.op.input_tensors
s[G].compute_inline()
(
eps,
nu,
k,
c,
kk,
) = s[U].op.axis
(eps, nu, k, c, kk) = s[U].op.axis
if autotvm.GLOBAL_SCOPE.in_tuning:
# kernel transformation will be pre-computed during compilation, so we skip
# this part to make tuning records correct
Expand Down Expand Up @@ -364,7 +355,7 @@ def conv2d_nchw_winograd_nnpack(cfg, data, kernel, strides, padding, dilation, o
tvm.contrib.nnpack.ConvolutionAlgorithm.WT_8x8_FP16,
)
else:
raise ValueError("Unsupported data type {} for conv2d winograd nnpack".format(dtype))
raise ValueError(f"Unsupported data type {dtype} for conv2d winograd nnpack")


@autotvm.register_topi_schedule("conv2d_nchw_winograd_nnpack.arm_cpu")
Expand Down
26 changes: 12 additions & 14 deletions python/tvm/topi/arm_cpu/conv2d_alter_op.py
Original file line number Diff line number Diff line change
Expand Up @@ -152,9 +152,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
)
dispatch_ctx.update(target, new_workload, cfg)
return relay.nn.conv2d(
inputs[0],
relay.Constant(tvm.nd.array(reshaped_new_kernel)),
**new_attrs,
inputs[0], relay.Constant(tvm.nd.array(reshaped_new_kernel)), **new_attrs
)

# Only microTVM does layout alteration for NHWC layout with real data types
Expand All @@ -167,7 +165,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
CO, _, KH, KW = get_const_tuple(kernel.shape)
VC = cfg["tile_co"].size[-1]

new_attrs["kernel_layout"] = "OIHW%do" % VC
new_attrs["kernel_layout"] = f"OIHW{VC}o"

new_data = data
new_kernel = te.placeholder((idxd(CO, VC), CI, KH, KW, VC), dtype=kernel.dtype)
Expand Down Expand Up @@ -275,7 +273,7 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
CO, M, KH, KW = get_const_tuple(kernel.shape)
VC = cfg["tile_co"].size[-1]

new_attrs["kernel_layout"] = "OIHW%do" % (cfg["tile_co"].size[-1])
new_attrs["kernel_layout"] = f"OIHW{cfg['tile_co'].size[-1]}o"

# Store the same config for the altered operator (workload)
new_data = data
Expand Down Expand Up @@ -309,10 +307,10 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):

# update new attrs
new_attrs["channels"] = out_channel
new_attrs["data_layout"] = "NCHW%dc" % ic_bn
new_attrs["data_layout"] = f"NCHW{ic_bn}c"
# (oc, ic, h, w) -> (OC, IC, h, w, ic, oc)
new_attrs["kernel_layout"] = "OIHW%di%do" % (ic_bn, oc_bn)
new_attrs["out_layout"] = "NCHW%dc" % oc_bn
new_attrs["kernel_layout"] = f"OIHW{ic_bn}i{oc_bn}o"
new_attrs["out_layout"] = f"NCHW{oc_bn}c"

# Store altered operator's config
new_data = te.placeholder(
Expand Down Expand Up @@ -353,9 +351,9 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):

# update new attrs
new_attrs["channels"] = out_channel
new_attrs["data_layout"] = "NCHW%dc" % ic_bn
new_attrs["kernel_layout"] = "OIHW1i%do" % oc_bn
new_attrs["out_layout"] = "NCHW%dc" % oc_bn
new_attrs["data_layout"] = f"NCHW{ic_bn}c"
new_attrs["kernel_layout"] = f"OIHW1i{oc_bn}o"
new_attrs["out_layout"] = f"NCHW{oc_bn}c"

# Store altered operator's config.
new_data = te.placeholder(
Expand Down Expand Up @@ -407,9 +405,9 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):

# update new attrs
new_attrs["channels"] = out_channel
new_attrs["data_layout"] = "NCHW%dc" % ic_bn
new_attrs["kernel_layout"] = "OIHW{:n}i{:n}o{:n}i".format(ic_bn // n_elems, oc_bn, n_elems)
new_attrs["out_layout"] = "NCHW%dc" % oc_bn
new_attrs["data_layout"] = f"NCHW{ic_bn}c"
new_attrs["kernel_layout"] = f"OIHW{ic_bn // n_elems:n}i{oc_bn:n}o{n_elems:n}i"
new_attrs["out_layout"] = f"NCHW{oc_bn}c"

# Store altered operator's config.
new_data = te.placeholder(
Expand Down
2 changes: 1 addition & 1 deletion python/tvm/topi/cuda/batch_matmul.py
Original file line number Diff line number Diff line change
Expand Up @@ -342,7 +342,7 @@ def _schedule_batch_matmul_int8(cfg, s, output):
_, N, _ = get_const_tuple(input_y.shape)

k_factor = 4
assert K % k_factor == 0, "Input dimension must divide {}".format(k_factor)
assert K % k_factor == 0, f"Input dimension must divide {k_factor}"
if K % 16 == 0:
k_factor = 16

Expand Down
4 changes: 2 additions & 2 deletions python/tvm/topi/cuda/conv2d_int8.py
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ def conv2d_NCHWc_int8(cfg, data, kernel, stride, padding, dilation, layout, out_
batch, channels, height, width = get_const_tuple(data.shape)
assert (
channels % ic_block_factor == 0
), "Number of input channels should be multiple of {}".format(ic_block_factor)
), f"Number of input channels should be multiple of {ic_block_factor}"
packed_data = te.compute(
(batch, channels // ic_block_factor, height, width, ic_block_factor),
lambda n, c, h, w, vc: data[n, c * ic_block_factor + vc, h, w],
Expand All @@ -100,7 +100,7 @@ def conv2d_NCHWc_int8(cfg, data, kernel, stride, padding, dilation, layout, out_
out_channels, in_channels, kernel_h, kernel_w = get_const_tuple(kernel.shape)
assert (
out_channels % oc_block_factor == 0
), "Number of output channels should be multiple of {}".format(oc_block_factor)
), f"Number of output channels should be multiple of {oc_block_factor}"
packed_kernel = te.compute(
(
out_channels // oc_block_factor,
Expand Down
2 changes: 1 addition & 1 deletion python/tvm/topi/cuda/conv3d_direct.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ def schedule_direct_conv3d_cuda(cfg, s, conv, layout, workload_name):
elif layout == "NDHWC":
n, d, y, x, f = s[conv].op.axis
else:
raise ValueError("not support this layout {} yet".format(layout))
raise ValueError(f"not support this layout {layout} yet")
rc, rd, ry, rx = s[conv].op.reduce_axis
cfg.define_split("tile_f", f, num_outputs=4)
cfg.define_split("tile_d", d, num_outputs=4)
Expand Down
18 changes: 3 additions & 15 deletions python/tvm/topi/cuda/dense.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,13 +29,7 @@


def _matmul_cublas_common(
cfg,
tensor_a,
tensor_b,
bias=None,
out_dtype=None,
transpose_a=False,
transpose_b=False,
cfg, tensor_a, tensor_b, bias=None, out_dtype=None, transpose_a=False, transpose_b=False
):
assert len(tensor_a.shape) == 2 and len(tensor_b.shape) == 2, "only support 2-dim matmul"
if bias is not None:
Expand All @@ -58,13 +52,7 @@ def _matmul_cublas_common(

@autotvm.register_topi_compute("matmul_cublas.cuda")
def matmul_cublas(
cfg,
tensor_a,
tensor_b,
bias=None,
out_dtype=None,
transpose_a=False,
transpose_b=False,
cfg, tensor_a, tensor_b, bias=None, out_dtype=None, transpose_a=False, transpose_b=False
):
"""Matmul operator on CUDA with CUBLAS"""
return _matmul_cublas_common(cfg, tensor_a, tensor_b, bias, out_dtype, transpose_a, transpose_b)
Expand Down Expand Up @@ -142,7 +130,7 @@ def _schedule_dense_int8(cfg, s, output):
out_dim, _ = get_const_tuple(weight.shape)

in_dim_factor = 4
assert in_dim % in_dim_factor == 0, "Input dimension must divide {}".format(in_dim_factor)
assert in_dim % in_dim_factor == 0, f"Input dimension must divide {in_dim_factor}"
if in_dim % 16 == 0:
in_dim_factor = 16

Expand Down
12 changes: 4 additions & 8 deletions python/tvm/topi/cuda/group_conv2d_nchw.py
Original file line number Diff line number Diff line change
Expand Up @@ -245,10 +245,10 @@ def group_conv2d_NCHWc_int8(
assert out_channels % groups == 0, "output channels must divide group size"
assert (
channels % ic_block_factor == 0
), "Number of input channels per group must divide {}".format(ic_block_factor)
), f"Number of input channels per group must divide {ic_block_factor}"
assert (
out_channels % oc_block_factor == 0
), "Number of output channels per group must divide {}".format(oc_block_factor)
), f"Number of output channels per group must divide {oc_block_factor}"

packed_data = te.compute(
(batch, channels // ic_block_factor, height, width, ic_block_factor),
Expand Down Expand Up @@ -282,14 +282,10 @@ def group_conv2d_NCHWc_int8(
# Shall we pad the channels to avoid raising assertions?
assert (
groups <= oc_chunk
), "Number of groups {} should be less than " "output channel chunk size {}".format(
groups, oc_chunk
)
), f"Number of groups {groups} should be less than output channel chunk size {oc_chunk}"
assert (
groups <= ic_chunk
), "Number of groups {} should be less than " "input channel chunk size {}".format(
groups, ic_chunk
)
), f"Number of groups {groups} should be less than input channel chunk size {ic_chunk}"

if isinstance(stride, int):
stride_h = stride_w = stride
Expand Down
2 changes: 1 addition & 1 deletion python/tvm/topi/cuda/scan.py
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@

def _get_thrust_func_name(tvmop):
tvmop_to_thrust_func_name = {tvm.tir.generic.add: "tvm.contrib.thrust.sum_scan"}
assert tvmop in tvmop_to_thrust_func_name, "{} not supported by thrust".format(tvmop)
assert tvmop in tvmop_to_thrust_func_name, f"{tvmop} not supported by thrust"
return tvmop_to_thrust_func_name[tvmop]


Expand Down
5 changes: 1 addition & 4 deletions python/tvm/topi/cuda/softmax.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,10 +44,7 @@ def _schedule_softmax(softmax_op, s, outs, tgt):
expsum = softmax_op.input_tensors[2]
else:
raise ValueError(
"Tag is expected to be softmax_output or log_softmax_output. \
Got {0}".format(
op_tag
)
f"Tag is expected to be softmax_output or log_softmax_output. Got {op_tag}"
)

# The nvptx and rocm backends only supports 32-bits warp shuffle
Expand Down
7 changes: 3 additions & 4 deletions python/tvm/topi/cuda/sparse.py
Original file line number Diff line number Diff line change
Expand Up @@ -159,10 +159,9 @@ def gen_ir(data, w_data, w_indices, w_indptr, out):
bs_m = bs_n
mb = m // bs_m
mi = warp_size
assert (
mb >= mi
), "Number of block rows in dense matrix must be larger than warp size: {} vs {}.".format(
warp_size, mb
assert mb >= mi, (
f"Number of block rows in dense matrix must be larger than warp size: "
f"{warp_size} vs {mb}."
)
mo = ceil_div(mb, mi)
ni = 1 # TODO(tkonolige): how do I compute the number of warps per block?
Expand Down
34 changes: 15 additions & 19 deletions python/tvm/topi/generic/conv2d.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,14 +45,12 @@ def fallback_schedule_cpu_common_int8(cfg, wkl, int32_lanes, num_int8_elements):
dilated_kernel_w = (wkl.kernel_w - 1) * wkl.dilation_w + 1
out_width = (wkl.width + pl + pr - dilated_kernel_w) // WSTR + 1

assert wkl.out_filter % int32_lanes == 0, "wkl.out_filter=%d, int32_lanes=%d" % (
wkl.out_filter,
int32_lanes,
)
assert wkl.in_filter % num_int8_elements == 0, "wkl.in_filter=%d, num_int8_elements=%d" % (
wkl.in_filter,
num_int8_elements,
)
assert (
wkl.out_filter % int32_lanes == 0
), f"wkl.out_filter={wkl.out_filter}, int32_lanes={int32_lanes}"
assert (
wkl.in_filter % num_int8_elements == 0
), f"wkl.in_filter={wkl.in_filter}, num_int8_elements={num_int8_elements}"

oc_bn = int32_lanes if int32_lanes >= num_int8_elements else num_int8_elements
ic_bn = 1
Expand Down Expand Up @@ -93,14 +91,12 @@ def fallback_schedule_cpu_1x1_int8(cfg, wkl, int32_lanes, num_int8_elements):
out_height = (wkl.height + pt + pb - wkl.kernel_h) // HSTR + 1
out_width = (wkl.width + pl + pr - wkl.kernel_w) // WSTR + 1

assert wkl.out_filter % int32_lanes == 0, "wkl.out_filter=%d, int32_lanes=%d" % (
wkl.out_filter,
int32_lanes,
)
assert wkl.in_filter % num_int8_elements == 0, "wkl.in_filter=%d, num_int8_elements=%d" % (
wkl.in_filter,
num_int8_elements,
)
assert (
wkl.out_filter % int32_lanes == 0
), f"wkl.out_filter={wkl.out_filter}, int32_lanes={int32_lanes}"
assert (
wkl.in_filter % num_int8_elements == 0
), f"wkl.in_filter={wkl.in_filter}, num_int8_elements={num_int8_elements}"

oc_bn = int32_lanes if int32_lanes >= num_int8_elements else num_int8_elements
ic_bn = 1
Expand All @@ -118,7 +114,7 @@ def fallback_schedule_cpu_1x1_int8(cfg, wkl, int32_lanes, num_int8_elements):
cfg["tile_oh"] = OtherOptionEntity(oh_factor)
cfg["tile_ow"] = SplitEntity([out_width // ow_factor, ow_factor])
return
raise ValueError("cannot decide default schedule for workload: {}".format(wkl))
raise ValueError(f"cannot decide default schedule for workload: {wkl}")


def schedule_conv_NCHWc_cpu_common_int8(
Expand Down Expand Up @@ -257,7 +253,7 @@ def schedule_conv_NCHWc_cpu_common_int8(
oc_chunk, oc_block = s[O].split(oc, factor=oc_bn)
s[O].reorder(oc_chunk, oh, ow_chunk, ow_block, oc_block)
else:
raise ValueError("Unsupported output ndim: %s" % out_ndim)
raise ValueError(f"Unsupported output ndim: {out_ndim}")
parallel_axis = s[O].fuse(batch, oc_chunk, oh)
if inline_fused:
s[C].compute_at(s[O], ow_block)
Expand Down Expand Up @@ -382,7 +378,7 @@ def schedule_conv_NCHWc_cpu_1x1_int8(
oh_outer, oh_inner = s[O].split(oh, factor=oh_factor)
ow_outer, ow_inner = s[O].split(ow, factor=ow_factor)
else:
raise ValueError("Unsupported output ndim: %s" % out_ndim)
raise ValueError(f"Unsupported output ndim: {out_ndim}")

s[O].reorder(oc_chunk, oh_outer, ow_outer, oh_inner, ow_inner, oc_block)
parallel_axis = s[O].fuse(batch, oc_chunk, oh_outer)
Expand Down
6 changes: 3 additions & 3 deletions python/tvm/topi/hexagon/conv2d_alter_op.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,9 +51,9 @@ def _alter_conv2d_layout(attrs, inputs, tinfos, out_type):
new_attrs = {k: attrs[k] for k in attrs.keys()}

new_attrs["channels"] = out_channel
new_attrs["data_layout"] = "NCHW%dc" % ic_bn
new_attrs["kernel_layout"] = "OIHW{:n}i{:n}o{:n}i".format(ic_bn // n_elems, oc_bn, n_elems)
new_attrs["out_layout"] = "NCHW%dc" % oc_bn
new_attrs["data_layout"] = f"NCHW{ic_bn}c"
new_attrs["kernel_layout"] = f"OIHW{ic_bn // n_elems:n}i{oc_bn:n}o{n_elems:n}i"
new_attrs["out_layout"] = f"NCHW{oc_bn}c"

return relay.nn.contrib_conv2d_nchwc(*inputs, **new_attrs)

Expand Down
6 changes: 3 additions & 3 deletions python/tvm/topi/hexagon/qnn/conv2d_alter_op.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,9 +44,9 @@ def _alter_qnn_conv2d_layout(attrs, inputs, tinfos, _out_type):

new_attrs = dict(attrs)
new_attrs["channels"] = out_channel
new_attrs["data_layout"] = "NCHW%dc" % ic_bn
new_attrs["kernel_layout"] = "OIHW{:n}i{:n}o{:n}i".format(ic_bn // n_elems, oc_bn, n_elems)
new_attrs["out_layout"] = "NCHW%dc" % oc_bn
new_attrs["data_layout"] = f"NCHW{ic_bn}c"
new_attrs["kernel_layout"] = f"OIHW{ic_bn // n_elems:n}i{oc_bn:n}o{n_elems:n}i"
new_attrs["out_layout"] = f"NCHW{oc_bn}c"

return relay.qnn.op.conv2d(*inputs, **new_attrs)

Expand Down
15 changes: 6 additions & 9 deletions python/tvm/topi/hls/nn.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ def traverse(OP):
Out = outs[0].op.output(0)
s[Conv2d].compute_at(s[Out], s[Out].op.axis[1])
else:
raise RuntimeError("Unsupported operator: %s" % OP.tag)
raise RuntimeError(f"Unsupported operator: {OP.tag}")

traverse(outs[0].op)

Expand Down Expand Up @@ -223,7 +223,7 @@ def traverse(OP):
Out = outs[0].op.output(0)
s[Reduce].compute_at(s[Out], s[Out].op.axis[0])
else:
raise RuntimeError("Unsupported operator: %s" % OP.tag)
raise RuntimeError(f"Unsupported operator: {OP.tag}")

traverse(outs[0].op)

Expand Down Expand Up @@ -264,10 +264,7 @@ def schedule_softmax(outs):
expsum = softmax.op.input_tensors[2]
else:
raise ValueError(
"Tag is expected to be softmax_output or log_softmax_output. \
Got {0}".format(
op_tag
)
f"Tag is expected to be softmax_output or log_softmax_output. Got {op_tag}"
)

if exp is not None:
Expand Down Expand Up @@ -315,7 +312,7 @@ def traverse(OP):
Out = outs[0].op.output(0)
s[Dense].compute_at(s[Out], s[Out].op.axis[1])
else:
raise RuntimeError("Unsupported operator: %s" % OP.tag)
raise RuntimeError(f"Unsupported operator: {OP.tag}")

traverse(outs[0].op)

Expand Down Expand Up @@ -358,7 +355,7 @@ def traverse(OP):
Out = outs[0].op.output(0)
s[Pool].compute_at(s[Out], s[Out].op.axis[1])
else:
raise RuntimeError("Unsupported operator: %s" % OP.tag)
raise RuntimeError(f"Unsupported operator: {OP.tag}")

traverse(outs[0].op)

Expand Down Expand Up @@ -401,7 +398,7 @@ def traverse(OP):
Out = outs[0].op.output(0)
s[Pool].compute_at(s[Out], s[Out].op.axis[1])
else:
raise RuntimeError("Unsupported operator: %s" % OP.tag)
raise RuntimeError(f"Unsupported operator: {OP.tag}")

traverse(outs[0].op)

Expand Down
Loading