Ensure that DenseElementAttr data is 64-bit aligned.

This allows for the raw data to be reinterpreted as the derived c++ type, e.g. ArrayRef<uint64_t>. This fixes a ubsan error for misaligned-pointer-use.

PiperOrigin-RevId: 259128031
This commit is contained in:
River Riddle 2019-07-20 09:22:36 -07:00 committed by A. Unique TensorFlower
parent 6b436eacbc
commit a47704e1e1
1 changed files with 13 additions and 6 deletions

View File

@ -480,15 +480,22 @@ struct DenseElementsAttributeStorage : public AttributeStorage {
/// Construct a new storage instance.
static DenseElementsAttributeStorage *
construct(AttributeStorageAllocator &allocator, KeyTy key) {
// If the data buffer is non-empty, we copy it into the allocator.
ArrayRef<char> data = allocator.copyInto(key.data);
// If the data buffer is non-empty, we copy it into the allocator with a
// 64-bit alignment.
ArrayRef<char> copy, data = key.data;
if (!data.empty()) {
char *rawData = reinterpret_cast<char *>(
allocator.allocate(data.size(), alignof(uint64_t)));
std::memcpy(rawData, data.data(), data.size());
// If this is a boolean splat, make sure only the first bit is used.
if (key.isSplat && key.type.getElementTypeBitWidth() == 1)
const_cast<char &>(data.front()) &= 1;
// If this is a boolean splat, make sure only the first bit is used.
if (key.isSplat && key.type.getElementTypeBitWidth() == 1)
rawData[0] &= 1;
copy = ArrayRef<char>(rawData, data.size());
}
return new (allocator.allocate<DenseElementsAttributeStorage>())
DenseElementsAttributeStorage(key.type, data, key.isSplat);
DenseElementsAttributeStorage(key.type, copy, key.isSplat);
}
ArrayRef<char> data;