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.

Sign up for GitHub

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

Special case bool dereferencing (consistent with PyTorch) #410

Open
wants to merge 1 commit into
base: develop
Choose a base branch
from
+27 −2
Open

Special case bool dereferencing (consistent with PyTorch) #410

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
29 changes: 27 additions & 2 deletions rocprim/include/rocprim/iterator/transform_iterator.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
// Copyright (c) 2017-2021 Advanced Micro Devices, Inc. All rights reserved.
// Copyright (c) 2017-2023 Advanced Micro Devices, Inc. All rights reserved.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
Expand Down Expand Up @@ -33,6 +33,29 @@

BEGIN_ROCPRIM_NAMESPACE

template <typename ITType, typename T>
struct LoadImpl {
ROCPRIM_HOST_DEVICE static T apply(ITType src) {
return *src;
}
};

template <typename ITType>
struct LoadImpl<ITType, bool> {
ROCPRIM_HOST_DEVICE static bool apply(ITType src) {
static_assert(sizeof(bool) == sizeof(char), "");
// Protect against invalid boolean values by loading as a byte
// first, then converting to bool.
return *reinterpret_cast<const unsigned char*>(src);
}
};


template <typename ITType, typename T>
ROCPRIM_HOST_DEVICE T load(ITType src) {
return LoadImpl<ITType, T>::apply(src);
}

/// \class transform_iterator
/// \brief A random-access input (read-only) iterator adaptor for transforming dereferenced values.
///
Expand Down Expand Up @@ -73,6 +96,8 @@ class transform_iterator
/// The type of unary function used to transform input range.
using unary_function = UnaryFunction;

using deref_type = typename std::iterator_traits<InputIterator>::value_type;

#ifndef DOXYGEN_SHOULD_SKIP_THIS
using self_type = transform_iterator;
#endif
Expand Down Expand Up @@ -125,7 +150,7 @@ class transform_iterator
ROCPRIM_HOST_DEVICE inline
value_type operator*() const
{
return transform_(*iterator_);
return transform_(load<InputIterator,deref_type>(iterator_));
}

ROCPRIM_HOST_DEVICE inline
Expand Down