[Mlir-commits] [mlir] [mlir] Fix use-after-free bugs in {RankedTensorType|VectorType}::Builder (PR #68969)
Benjamin Maxwell
llvmlistbot at llvm.org
Fri Oct 13 02:45:07 PDT 2023
https://github.com/MacDue created https://github.com/llvm/llvm-project/pull/68969
Previously, these would set their ArrayRef members to reference their storage SmallVectors after a copy-on-write operation. This leads to a use-after-free if the builder is copied and the original destroyed (as the new builder would still reference the old SmallVector).
The VectorType::Builder also set the ArrayRef<bool> scalableDims member to a temporary SmallVector when the provided scalableDims are empty. This again lead to a use-after-free, and is unnecessary as VectorType::get already handles being passed an empty scalableDims array.
These bugs were in-part caught by UBSAN, see:
https://lab.llvm.org/buildbot/#/builders/5/builds/37355
>From 081be10b57e31a785b2b8db84cb3447eababf36d Mon Sep 17 00:00:00 2001
From: Benjamin Maxwell <benjamin.maxwell at arm.com>
Date: Fri, 13 Oct 2023 09:33:25 +0000
Subject: [PATCH] [mlir] Fix use-after-free bugs in
{RankedTensorType|VectorType}::Builder
Previously, these would set their ArrayRef members to reference their
storage SmallVectors after a copy-on-write operation. This leads to a
use-after-free if the builder is copied and the original destroyed (as
the new builder would still reference the old SmallVector).
The VectorType::Builder also set the ArrayRef<bool> scalableDims member
to a temporary SmallVector when the provided scalableDims are empty.
This again lead to a use-after-free, and is unnecessary as
VectorType::get already handles being passed an empty scalableDims
array.
These bugs were in-part caught by UBSAN, see:
https://lab.llvm.org/buildbot/#/builders/5/builds/37355
---
mlir/include/mlir/IR/BuiltinTypes.h | 43 ++++++++++++++++-------------
1 file changed, 24 insertions(+), 19 deletions(-)
diff --git a/mlir/include/mlir/IR/BuiltinTypes.h b/mlir/include/mlir/IR/BuiltinTypes.h
index 9df5548cd5d939c..13fbae90b68c6cb 100644
--- a/mlir/include/mlir/IR/BuiltinTypes.h
+++ b/mlir/include/mlir/IR/BuiltinTypes.h
@@ -277,7 +277,7 @@ class RankedTensorType::Builder {
if (storage.empty())
storage.append(shape.begin(), shape.end());
storage.erase(storage.begin() + pos);
- shape = {storage.data(), storage.size()};
+ shape = {};
return *this;
}
@@ -287,12 +287,17 @@ class RankedTensorType::Builder {
if (storage.empty())
storage.append(shape.begin(), shape.end());
storage.insert(storage.begin() + pos, val);
- shape = {storage.data(), storage.size()};
+ shape = {};
return *this;
}
+ /// Returns the current shape.
+ ArrayRef<int64_t> getShape() {
+ return shape.empty() ? ArrayRef(storage) : shape;
+ }
+
operator RankedTensorType() {
- return RankedTensorType::get(shape, elementType, encoding);
+ return RankedTensorType::get(getShape(), elementType, encoding);
}
private:
@@ -319,20 +324,11 @@ class VectorType::Builder {
/// Build from scratch.
Builder(ArrayRef<int64_t> shape, Type elementType,
unsigned numScalableDims = 0, ArrayRef<bool> scalableDims = {})
- : shape(shape), elementType(elementType) {
- if (scalableDims.empty())
- scalableDims = SmallVector<bool>(shape.size(), false);
- else
- this->scalableDims = scalableDims;
- }
+ : shape(shape), elementType(elementType), scalableDims(scalableDims) {}
Builder &setShape(ArrayRef<int64_t> newShape,
ArrayRef<bool> newIsScalableDim = {}) {
- if (newIsScalableDim.empty())
- scalableDims = SmallVector<bool>(shape.size(), false);
- else
- scalableDims = newIsScalableDim;
-
+ scalableDims = newIsScalableDim;
shape = newShape;
return *this;
}
@@ -351,9 +347,8 @@ class VectorType::Builder {
storageScalableDims.append(scalableDims.begin(), scalableDims.end());
storage.erase(storage.begin() + pos);
storageScalableDims.erase(storageScalableDims.begin() + pos);
- shape = {storage.data(), storage.size()};
- scalableDims =
- ArrayRef<bool>(storageScalableDims.data(), storageScalableDims.size());
+ shape = {};
+ scalableDims = {};
return *this;
}
@@ -363,12 +358,22 @@ class VectorType::Builder {
storage.append(shape.begin(), shape.end());
assert(pos < storage.size() && "overflow");
storage[pos] = val;
- shape = {storage.data(), storage.size()};
+ shape = {};
return *this;
}
+ /// Returns the current shape.
+ ArrayRef<int64_t> getShape() {
+ return shape.empty() ? ArrayRef(storage) : shape;
+ }
+
+ /// Returns the current scalable dims.
+ ArrayRef<bool> getScalableDims() {
+ return scalableDims.empty() ? ArrayRef(storageScalableDims) : scalableDims;
+ }
+
operator VectorType() {
- return VectorType::get(shape, elementType, scalableDims);
+ return VectorType::get(getShape(), elementType, getScalableDims());
}
private:
More information about the Mlir-commits
mailing list