Files
paddlepaddle--paddle/paddle/fluid/prim/utils/static/desc_tensor.h
T
2026-07-13 12:40:42 +08:00

69 lines
2.3 KiB
C++

// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
#pragma once
#include "paddle/common/ddim.h"
#include "paddle/fluid/framework/convert_utils.h"
#include "paddle/fluid/framework/var_desc.h"
#include "paddle/phi/core/extended_tensor.h"
#include "paddle/phi/core/utils/data_type.h"
#include "paddle/utils/any.h"
namespace paddle {
namespace prim {
class DescTensor : public phi::ExtendedTensor,
public phi::TypeInfoTraits<phi::TensorBase, DescTensor> {
public:
explicit DescTensor(framework::VarDesc* desc)
: desc_ptr_(desc), dims_(common::make_ddim(desc->GetShape())) {}
static const char* name() { return "DescTensor"; }
std::string Name() const { return desc_ptr_->Name(); }
std::vector<int64_t> shape() const { return desc_ptr_->GetShape(); }
const phi::DDim& dims() const override {
dims_ = common::make_ddim(desc_ptr_->GetShape());
return dims_;
}
int64_t numel() const override { return product(dims()); }
DataType dtype() const override {
return phi::TransToPhiDataType(desc_ptr_->GetDataType());
}
framework::VarDesc* get_ptr() { return desc_ptr_; }
const phi::Place& place() const override { return place_; }
bool initialized() const override { return desc_ptr_ != nullptr; }
// TODO(jiabin): override more operators here.
private:
// VarDesc's lifetime is held by block and it's program, so we just conceal
// its funcs instead of its life.
framework::VarDesc* desc_ptr_;
// TODO(jiabin): This is really ugly, but we have to hold a dims here so that
// we can inherit from ExtendedTensor Remove this when we make VarDesc's as
// same as Tensor, or make Tensor's dims more lightly.
mutable DDim dims_;
Place place_;
};
} // namespace prim
} // namespace paddle