Skip to content

Commit

Permalink
use bitshifts for int to int in convert_dtype (#6978)
Browse files Browse the repository at this point in the history
  • Loading branch information
pmeier authored Nov 28, 2022
1 parent 51e8dac commit 346f6dd
Showing 1 changed file with 1 addition and 9 deletions.
10 changes: 1 addition & 9 deletions torchvision/prototype/transforms/functional/_meta.py
Original file line number Diff line number Diff line change
Expand Up @@ -379,15 +379,7 @@ def convert_dtype_image_tensor(image: torch.Tensor, dtype: torch.dtype = torch.f
if num_value_bits_input > num_value_bits_output:
return image.bitwise_right_shift(num_value_bits_input - num_value_bits_output).to(dtype)
else:
# The bitshift kernel is not vectorized
# https://github.com/pytorch/pytorch/blob/703c19008df4700b6a522b0ae5c4b6d5ffc0906f/aten/src/ATen/native/cpu/BinaryOpsKernel.cpp#L315-L322
# This results in the multiplication actually being faster.
# TODO: If the bitshift kernel is optimized in core, replace the computation below with
# `image.to(dtype).bitwise_left_shift_(num_value_bits_output - num_value_bits_input)`
max_value_input = float(_FT._max_value(dtype))
max_value_output = float(_FT._max_value(image.dtype))
factor = int((max_value_input + 1) // (max_value_output + 1))
return image.to(dtype).mul_(factor)
return image.to(dtype).bitwise_left_shift_(num_value_bits_output - num_value_bits_input)


# We changed the name to align it with the new naming scheme. Still, `convert_image_dtype` is
Expand Down

0 comments on commit 346f6dd

Please sign in to comment.