buffer.h

Go to the documentation of this file.

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.
 */
 
#ifndef TVM_TIRX_BUFFER_H_
#define TVM_TIRX_BUFFER_H_
 
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <tvm/ir/expr.h>
#include <tvm/script/printer/config.h>
#include <tvm/tirx/layout.h>
#include <tvm/tirx/var.h>
 
#include <string>
 
namespace tvm {
namespace tirx {
 
#ifndef TVM_INDEX_DEFAULT_I64
#define TVM_INDEX_DEFAULT_I64 1
#endif
inline DataType DefaultIndexType() {
#if TVM_INDEX_DEFAULT_I64
  return DataType::Int(64);
#else
  return DataType::Int(32);
#endif
}
 
// forward declare Stmt
class Stmt;
 
enum BufferType : int {
  kDefault = 1,
  // Maps buffer[i][j][k] -> buffer[i][0][k] if dimension i's shape equals 1.
  kAutoBroadcast = 2,
};
 
class BufferNode : public ffi::Object {
 public:
  // Data fields.
  Var data;
  DataType dtype;
  ffi::Array<PrimExpr> shape;
  ffi::Array<IntImm> axis_separators;
  ffi::Array<PrimExpr> strides;
  PrimExpr elem_offset;
  // Meta data
  ffi::String name;
  int data_alignment;
  int offset_factor;
  BufferType buffer_type;
  mutable Span span;
 
  ffi::Optional<Layout> layout;
 
  ffi::Array<PrimExpr> allocated_addr;
 
  BufferNode() {}
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<BufferNode>()
        .def_ro("data", &BufferNode::data, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("dtype", &BufferNode::dtype)
        .def_ro("shape", &BufferNode::shape, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("strides", &BufferNode::strides, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("axis_separators", &BufferNode::axis_separators,
                refl::AttachFieldFlag::SEqHashDef())
        .def_ro("elem_offset", &BufferNode::elem_offset, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("name", &BufferNode::name, refl::AttachFieldFlag::SEqHashIgnore())
        .def_ro("data_alignment", &BufferNode::data_alignment)
        .def_ro("offset_factor", &BufferNode::offset_factor)
        .def_ro("buffer_type", &BufferNode::buffer_type)
        .def_ro("span", &BufferNode::span, refl::AttachFieldFlag::SEqHashIgnore())
        .def_ro("layout", &BufferNode::layout)
        .def_ro("allocated_addr", &BufferNode::allocated_addr);
  }
 
  DataType DefaultIndexType() const {
    return shape.size() != 0 ? shape[0].dtype() : tvm::tirx::DefaultIndexType();
  }
 
  ffi::Array<PrimExpr> ElemOffset(ffi::Array<PrimExpr> index, bool inner = false) const;
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Buffer", BufferNode, ffi::Object);
  TVM_OBJECT_ENABLE_SCRIPT_PRINTER();
};
 
class Buffer : public ffi::ObjectRef {
 public:
  // User can specify data_alignment and offset_factor to be 0
  // A default value will be picked.
  TVM_DLL Buffer(Var data, DataType dtype, ffi::Array<PrimExpr> shape, ffi::Array<PrimExpr> strides,
                 PrimExpr elem_offset, ffi::String name, int data_alignment, int offset_factor,
                 BufferType buffer_type, ffi::Array<IntImm> axis_separators = {},
                 Span span = Span(), ffi::Optional<Layout> layout = std::nullopt,
                 ffi::Array<PrimExpr> allocated_addr = {});
 
  TVM_DLL Buffer MakeStrideView() const;
  TVM_DLL Buffer MakeSlice(ffi::Array<PrimExpr> begins, ffi::Array<PrimExpr> extents) const;
  TVM_DLL PrimExpr access_ptr(int access_mask, DataType ptr_type = DataType::Handle(),
                              int content_lanes = 1, PrimExpr offset = IntImm(DataType::Int(32), 0),
                              ffi::Optional<PrimExpr> input_extent = std::nullopt) const;
  TVM_DLL PrimExpr vload(ffi::Array<PrimExpr> begin, DataType dtype,
                         ffi::Optional<PrimExpr> predicate = std::nullopt) const;
  TVM_DLL Stmt vstore(ffi::Array<PrimExpr> begin, PrimExpr value,
                      ffi::Optional<PrimExpr> predicate = std::nullopt) const;
 
  Buffer GetFlattenedBuffer() const;
 
  ffi::Array<PrimExpr> OffsetOf(ffi::Array<PrimExpr> index) const;
 
  PrimExpr OffsetOf_p(const ffi::Array<PrimExpr>& indices) const;
 
  TVM_DLL ffi::String scope() const;
 
  TVM_DLL Buffer with_allocated_addr(ffi::Array<PrimExpr> allocated_addr) const;
 
  TVM_DLL bool IsScalar(bool alloc_or_decl = true) const;
 
  TVM_DLL Buffer with_dtype(DataType dtype) const;
 
  TVM_DLL Buffer with_data(Var data) const;
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Buffer, ffi::ObjectRef, BufferNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(BufferNode);
};
 
TVM_DLL Buffer decl_buffer(ffi::Array<PrimExpr> shape, DataType dtype = DataType::Float(32),
                           ffi::String name = "buffer", ffi::String storage_scope = "",
                           ffi::Optional<ffi::Array<IntImm>> axis_separators = std::nullopt,
                           Span span = Span());
 
class DataProducerNode : public PrimExprConvertibleNode {
 public:
  virtual ~DataProducerNode() {}
  virtual ffi::Array<PrimExpr> GetShape() const = 0;
  virtual DataType GetDataType() const = 0;
  virtual ffi::String GetNameHint() const = 0;
  TVM_FFI_DECLARE_OBJECT_INFO("tirx.DataProducer", DataProducerNode, PrimExprConvertibleNode);
};
 
class DataProducer : public PrimExprConvertible {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(DataProducer, PrimExprConvertible, DataProducerNode);
};
 
TVM_DLL tirx::Buffer BufferWithOffsetAlignment(ffi::Array<PrimExpr> shape, DataType dtype,
                                               std::string name, int data_alignment,
                                               int offset_factor, bool compact,
                                               std::string memory_scope = "");
}  // namespace tirx
}  // namespace tvm
#endif  // TVM_TIR_BUFFER_H_