/*!
 * \file tl/op/fill.h
 * \brief Fill operations for tensor initialization
 */

#ifndef TVM_TL_OP_FILL_H_
#define TVM_TL_OP_FILL_H_

#include "operator.h"
#include "parallel.h"
#include "support/check.h"

namespace tvm {
namespace tl {

using namespace tirx;

/// Node class for fill operations
class FillNode : public TileOperatorNode {
public:
  tirx::Buffer dst;         ///< Destination buffer to fill
  PrimExpr value;           ///< Value to fill with
  ffi::Array<Range> region; ///< Region to fill within the buffer
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tl.Fill", FillNode, TileOperatorNode);

  Stmt Lower(const LowerArgs &lower_args,
             arith::Analyzer *analyzer) const override;
  LayoutMap InferLayout(const LayoutInferArgs &layout_args,
                        InferLevel level) const override;
  static const Op &Get();

  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<FillNode>()
        .def_ro("dst", &FillNode::dst)
        .def_ro("value", &FillNode::value)
        .def_ro("region", &FillNode::region);
  }

  TileOperator Clone() const;

  /// Create SIMT-style parallel loop for filling
  For MakeSIMTLoop(arith::Analyzer *analyzer) const;
};

using FillTargetPredicate = bool (*)(Target target);

struct FillImpl {
  const char *name;
  FillTargetPredicate match_target;

  Stmt (*lower)(const FillNode &op, const LowerArgs &lower_args,
                arith::Analyzer *analyzer);
};

void RegisterFillImpl(FillImpl impl);

/// Wrapper class for fill operations
class Fill : public TileOperator {
public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Fill, TileOperator, FillNode);
  TVM_DLL Fill(ffi::Array<PrimExpr> args,
               ffi::Map<ffi::String, ffi::ObjectRef> annotations =
                   ffi::Map<ffi::String, ffi::ObjectRef>());
  static const Op &Get();
};

} // namespace tl
} // namespace tvm

#endif // TVM_TL_OP_FILL_H_
