postproc.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_S_TIR_META_SCHEDULE_POSTPROC_H_
#define TVM_S_TIR_META_SCHEDULE_POSTPROC_H_
 
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ir/prim/expr.h>
#include <tvm/runtime/base.h>
#include <tvm/s_tir/schedule/schedule.h>
 
namespace tvm {
namespace s_tir {
namespace meta_schedule {
 
class TuneContext;
class Postproc;
 
class PostprocObj : public ffi::Object {
 public:
  virtual ~PostprocObj() = default;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PostprocObj>();
  }
 
  virtual void InitializeWithTuneContext(const TuneContext& context) = 0;
 
  virtual bool Apply(const s_tir::Schedule& sch) = 0;
 
  virtual Postproc Clone() const = 0;
 
  static constexpr const bool _type_mutable = true;
  TVM_FFI_DECLARE_OBJECT_INFO("s_tir.meta_schedule.Postproc", PostprocObj, ffi::Object);
};
 
class Postproc : public ffi::ObjectRef {
 public:
  explicit Postproc(ffi::UnsafeInit tag, ffi::ObjectPtr<PostprocObj> data) : ffi::ObjectRef(tag) {
    data_ = std::move(data);
  }
 
  using FInitializeWithTuneContext = ffi::TypedFunction<void(const TuneContext&)>;
  using FApply = ffi::TypedFunction<bool(const s_tir::Schedule&)>;
  using FClone = ffi::TypedFunction<Postproc()>;
  TVM_DLL static Postproc PyPostproc(FInitializeWithTuneContext f_initialize_with_tune_context,  //
                                     FApply f_apply,                                             //
                                     FClone f_clone);
  TVM_DLL static Postproc DisallowDynamicLoop();
  TVM_DLL static Postproc DisallowAsyncStridedMemCopy();
  TVM_DLL static Postproc RewriteCooperativeFetch();
  TVM_DLL static Postproc RewriteParallelVectorizeUnroll();
  TVM_DLL static Postproc RewriteReductionBlock();
  TVM_DLL static Postproc RewriteUnboundBlock(int max_threadblocks);
  TVM_DLL static Postproc RewriteTensorize(bool vectorize_init_loop = false);
  TVM_DLL static Postproc VerifyGPUCode();
  TVM_DLL static Postproc VerifyVTCMLimit();
  TVM_DLL static Postproc RewriteLayout();
  TVM_DLL static ffi::Array<Postproc, void> DefaultLLVM();
  TVM_DLL static ffi::Array<Postproc, void> DefaultCPUTensorization();
  TVM_DLL static ffi::Array<Postproc, void> DefaultRISCV();
  TVM_DLL static ffi::Array<Postproc, void> DefaultCUDA();
  TVM_DLL static ffi::Array<Postproc, void> DefaultCUDATensorCore();
  TVM_DLL static ffi::Array<Postproc, void> DefaultHexagon();
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Postproc, ffi::ObjectRef, PostprocObj);
};
 
class PyPostprocObj : public PostprocObj {
 public:
  using FInitializeWithTuneContext = Postproc::FInitializeWithTuneContext;
  using FApply = Postproc::FApply;
  using FClone = Postproc::FClone;
  FInitializeWithTuneContext f_initialize_with_tune_context;
  FApply f_apply;
  FClone f_clone;
 
  static void RegisterReflection() {
    // `f_initialize_with_tune_context` is not registered
    // `f_apply` is not registered
    // `f_clone` is not registered
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<PyPostprocObj>();
  }
 
  void InitializeWithTuneContext(const TuneContext& context) final;
  bool Apply(const s_tir::Schedule& sch) final;
  Postproc Clone() const final;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("s_tir.meta_schedule.PyPostproc", PyPostprocObj, PostprocObj);
};
 
}  // namespace meta_schedule
}  // namespace s_tir
}  // namespace tvm
 
#endif  // TVM_S_TIR_META_SCHEDULE_POSTPROC_H_