Skip to content

Commit

Permalink
#0: Clang-format buffer hpp/cpp files
Browse files Browse the repository at this point in the history 
- Add comment about struct Private in Buffer
#0: Switch to tt::stl::SmallVector container for DistributionType in DistributionSpec
  • Loading branch information
@TT-BrianLiu
TT-BrianLiu committed Mar 8, 2025
1 parent 9598f26 commit 20b8e74
Show file tree
Hide file tree
Showing 4 changed files with 128 additions and 87 deletions.
85 changes: 52 additions & 33 deletions tt_metal/api/tt-metalium/buffer.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,20 +53,33 @@ struct ShardSpec {
std::optional<std::array<uint32_t, 2>> physical_shard_shape = std::nullopt;

ShardSpec(
const CoreRangeSet &core_sets_,
const std::array<uint32_t, 2> &shard_shape_,
const ShardOrientation &shard_orientation_ = ShardOrientation::ROW_MAJOR,
const ShardMode &shard_mode_ = ShardMode::PHYSICAL) :
grid(core_sets_), shape(shard_shape_), orientation(shard_orientation_), mode(shard_mode_), physical_shard_shape(std::nullopt) {
}
const CoreRangeSet& core_sets_,
const std::array<uint32_t, 2>& shard_shape_,
const ShardOrientation& shard_orientation_ = ShardOrientation::ROW_MAJOR,
const ShardMode& shard_mode_ = ShardMode::PHYSICAL) :
grid(core_sets_),
shape(shard_shape_),
orientation(shard_orientation_),
mode(shard_mode_),
physical_shard_shape(std::nullopt) {}

ShardSpec(
const CoreRangeSet &core_sets_,
const std::array<uint32_t, 2> &shard_shape_,
const std::array<uint32_t, 2> &physical_shard_shape_,
const ShardOrientation &shard_orientation_ = ShardOrientation::ROW_MAJOR) :
grid(core_sets_), shape(shard_shape_), orientation(shard_orientation_), mode(ShardMode::LOGICAL), physical_shard_shape(physical_shard_shape_) {
TT_FATAL(physical_shard_shape_[0] >= shard_shape_[0] and physical_shard_shape_[1] >= shard_shape_[1], "Physical shard shape ({}, {}) must be greater or equal to logical shard shape ({}, {})!", physical_shard_shape_[0], physical_shard_shape_[1], shard_shape_[0], shard_shape_[1]);
const CoreRangeSet& core_sets_,
const std::array<uint32_t, 2>& shard_shape_,
const std::array<uint32_t, 2>& physical_shard_shape_,
const ShardOrientation& shard_orientation_ = ShardOrientation::ROW_MAJOR) :
grid(core_sets_),
shape(shard_shape_),
orientation(shard_orientation_),
mode(ShardMode::LOGICAL),
physical_shard_shape(physical_shard_shape_) {
TT_FATAL(
physical_shard_shape_[0] >= shard_shape_[0] and physical_shard_shape_[1] >= shard_shape_[1],
"Physical shard shape ({}, {}) must be greater or equal to logical shard shape ({}, {})!",
physical_shard_shape_[0],
physical_shard_shape_[1],
shard_shape_[0],
shard_shape_[1]);
}

const uint32_t num_cores() const { return this->grid.num_cores(); }
Expand All @@ -75,9 +88,11 @@ struct ShardSpec {
bool operator==(const ShardSpec& other) const;
bool operator!=(const ShardSpec& other) const;

static constexpr auto attribute_names = std::forward_as_tuple("grid", "shape", "orientation", "mode", "physical_shard_shape");
static constexpr auto attribute_names =
std::forward_as_tuple("grid", "shape", "orientation", "mode", "physical_shard_shape");
constexpr auto attribute_values() const {
return std::forward_as_tuple(this->grid, this->shape, this->orientation, this->mode, this->physical_shard_shape);
return std::forward_as_tuple(
this->grid, this->shape, this->orientation, this->mode, this->physical_shard_shape);
}
};

Expand Down Expand Up @@ -118,7 +133,7 @@ struct ShardSpecBuffer {
inline namespace v0 {

struct BufferConfig {
IDevice *device;
IDevice* device;
DeviceAddr size; // Size in bytes
DeviceAddr page_size; // Size of unit being interleaved. For non-interleaved buffers: size == page_size
BufferType buffer_type;
Expand All @@ -130,7 +145,7 @@ typedef BufferConfig InterleavedBufferConfig;
// copied from above instead of using inheritance such that we can use
// designator constructor
struct ShardedBufferConfig {
IDevice *device;
IDevice* device;
DeviceAddr size; // Size in bytes
DeviceAddr page_size; // Size of unit being interleaved. For non-interleaved buffers: size == page_size
BufferType buffer_type = BufferType::L1;
Expand All @@ -140,7 +155,7 @@ struct ShardedBufferConfig {

} // namespace v0

bool is_sharded(const TensorMemoryLayout &layout);
bool is_sharded(const TensorMemoryLayout& layout);

struct BufferPageMapping {
std::vector<CoreCoord> all_cores_;
Expand All @@ -167,11 +182,15 @@ struct BufferRegion {
};

class Buffer final {
struct Private { explicit Private() = default; };
// Used in public Buffer constructors so they are only callable within Buffer
// Buffer constructors are public so we can call std::make_shared on Buffer
struct Private {
explicit Private() = default;
};

public:
public:
static std::shared_ptr<Buffer> create(
IDevice *device,
IDevice* device,
DeviceAddr size,
DeviceAddr page_size,
BufferType buffer_type,
Expand All @@ -180,7 +199,7 @@ class Buffer final {
std::optional<bool> bottom_up = std::nullopt,
std::optional<SubDeviceId> sub_device_id = std::nullopt);
static std::shared_ptr<Buffer> create(
IDevice *device,
IDevice* device,
DeviceAddr address,
DeviceAddr size,
DeviceAddr page_size,
Expand All @@ -190,10 +209,10 @@ class Buffer final {
std::optional<bool> bottom_up = std::nullopt,
std::optional<SubDeviceId> sub_device_id = std::nullopt);

Buffer(const Buffer &other) = delete;
Buffer &operator=(const Buffer &other) = delete;
Buffer(Buffer &&other) = delete;
Buffer &operator=(Buffer &&other) = delete;
Buffer(const Buffer& other) = delete;
Buffer& operator=(const Buffer& other) = delete;
Buffer(Buffer&& other) = delete;
Buffer& operator=(Buffer&& other) = delete;
~Buffer();

IDevice* device() const { return device_; }
Expand Down Expand Up @@ -263,7 +282,7 @@ class Buffer final {
bool owns_data,
Private);

private:
private:
enum class AllocationStatus : uint8_t {
ALLOCATION_REQUESTED,
ALLOCATION_FAILED,
Expand All @@ -277,20 +296,20 @@ class Buffer final {
void deallocate();
static void deleter(Buffer* buffer);
void deallocate_impl();
friend void DeallocateBuffer(Buffer &buffer);
friend void DeallocateBuffer(Buffer& buffer);

DeviceAddr translate_page_address(uint64_t offset, uint32_t bank_id) const;

IDevice* const device_;
const DeviceAddr size_; // Size in bytes
const DeviceAddr size_; // Size in bytes
const BufferType buffer_type_;
const TensorMemoryLayout buffer_layout_;
const bool bottom_up_;
const std::optional<SubDeviceId> sub_device_id_;
const bool owns_data_;

std::optional<SubDeviceManagerId> sub_device_manager_id_;
Allocator * allocator_;
Allocator* allocator_;

std::atomic<AllocationStatus> allocation_status_ = AllocationStatus::ALLOCATION_REQUESTED;
DeviceAddr address_ = 0;
Expand All @@ -300,7 +319,7 @@ class Buffer final {
std::atomic<bool> deallocation_requested_ = false;

// These members must be only accessed on the device worker thread
DeviceAddr page_size_; // Size of unit being interleaved. For non-interleaved buffers: size == page_size
DeviceAddr page_size_; // Size of unit being interleaved. For non-interleaved buffers: size == page_size
std::optional<ShardSpecBuffer> shard_parameters_;
std::shared_ptr<const BufferPageMapping> buffer_page_mapping_;

Expand All @@ -311,7 +330,7 @@ class Buffer final {

} // namespace v0

BufferPageMapping generate_buffer_page_mapping(const Buffer &buffer);
BufferPageMapping generate_buffer_page_mapping(const Buffer& buffer);

inline namespace v0 {

Expand All @@ -322,14 +341,14 @@ using HostDataType = std::variant<
const std::shared_ptr<std::vector<uint32_t>>,
const std::shared_ptr<std::vector<float>>,
const std::shared_ptr<std::vector<bfloat16>>,
const void *>;
const void*>;

} // namespace v0
} // namespace tt::tt_metal

namespace tt::stl::json {
template <>
struct from_json_t<tt_metal::ShardSpec> {
tt_metal::ShardSpec operator()(const nlohmann::json &json_object) const;
tt_metal::ShardSpec operator()(const nlohmann::json& json_object) const;
};
} // namespace tt::stl::json
8 changes: 5 additions & 3 deletions tt_metal/api/tt-metalium/distribution_spec.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,13 +61,15 @@ class DistributionSpec {
struct ReplicateCount {
size_t count = 0;
};
using DistributionType = std::variant<ShardSize, ReplicateCount>;
using DistributionType = std::variant<std::monostate, ShardSize, ReplicateCount>;

DistributionSpec(
const tt::tt_metal::Shape& tensor_shape, const std::vector<DistributionType>& spec, size_t num_targets);
const tt::tt_metal::Shape& tensor_shape,
const tt::stl::SmallVector<DistributionType>& spec,
size_t num_targets);

tt::tt_metal::Shape tensor_shape_;
std::vector<DistributionType> spec_;
tt::stl::SmallVector<DistributionType> spec_;
tt::tt_metal::Shape shard_shape_; // Determined based on spec_
size_t num_targets_ = 0;
};
Expand Down
Loading

0 comments on commit 20b8e74

Please sign in to comment.