exec_scope.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_EXEC_SCOPE_H_
#define TVM_TIRX_EXEC_SCOPE_H_
 
#include <tvm/ffi/container/variant.h>
#include <tvm/ir/module.h>
#include <tvm/tirx/var.h>
 
#include <string>
#include <utility>
 
namespace tvm {
namespace tirx {
 
enum class ScopeKind : int {
  kWorld = 0,
  kKernel = 1,
  kCluster = 2,
  kCta = 3,
  kWarpgroup = 4,
  kWarp = 5,
  kThread = 6,
};
 
TVM_DLL std::string ScopeKindToString(ScopeKind kind);
 
TVM_DLL ScopeKind StringToScopeKind(const ffi::String& name);
 
enum class ScopeBinding : int {
  kKernelCluster = 0,
  kKernelCta = 1,
  kClusterCta = 2,
  kCtaWarpgroup = 3,
  kCtaWarp = 4,
  kWarpgroupWarp = 5,
  kWarpThread = 6,
  kCtaThread = 7,
  kWarpgroupThread = 8,
  kClusterCtaPair = 9,
};
 
TVM_DLL std::pair<ffi::String, ffi::String> ScopeBindingToStringPair(ScopeBinding binding);
 
TVM_DLL ScopeBinding StringPairToScopeBinding(const ffi::String& parent, const ffi::String& cur);
 
/******** Definition of ScopeId ********/
class ScopeIdDefNode : public ffi::Object {
 public:
  ffi::Array<Var> def_ids;
  ffi::Optional<ffi::Array<PrimExpr>> extents;
  ScopeBinding scope;
  ffi::Optional<ffi::Array<PrimExpr>> preferred_extents;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ScopeIdDefNode>()
        .def_ro("def_ids", &ScopeIdDefNode::def_ids, refl::AttachFieldFlag::SEqHashDef())
        .def_ro("extents", &ScopeIdDefNode::extents)
        .def_ro("scope", &ScopeIdDefNode::scope)
        .def_ro("preferred_extents", &ScopeIdDefNode::preferred_extents);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.ScopeIdDef", ScopeIdDefNode, ffi::Object);
};
 
class ScopeIdDef : public ffi::ObjectRef {
 public:
  TVM_DLL explicit ScopeIdDef(ffi::Array<Var> def_ids, ffi::Optional<ffi::Array<PrimExpr>> extents,
                              ScopeBinding scope,
                              ffi::Optional<ffi::Array<PrimExpr>> preferred_extents =
                                  ffi::Optional<ffi::Array<PrimExpr>>(std::nullopt));
 
  bool is_deferred() const { return !get()->extents.has_value(); }
 
  PrimExpr fused_extent() const;
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ScopeIdDef, ffi::ObjectRef, ScopeIdDefNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ScopeIdDefNode);
};
 
class ScopeIdDefVerifier {
 public:
  using ScopeIdSet = std::unordered_map<ScopeBinding, ScopeIdDef>;
 
  enum class Mode { kRelaxed, kStrict };
 
  bool Verify(const ffi::Array<ScopeIdDef>& defs, Mode mode = Mode::kStrict);
 
  ScopeIdSet id_set;
};
 
class ScopeIdResolve {
 public:
  using LaunchParams = std::unordered_map<ffi::String, IterVar>;
 
  TVM_DLL static ffi::Array<PrimExpr> Resolve(ScopeBinding binding,
                                              const ffi::Optional<ffi::Array<PrimExpr>>& extents,
                                              int out_dim, const ffi::String& target_kind,
                                              const LaunchParams& params);
 
  TVM_DLL static PrimExpr ComputeWarpIdInCta(const LaunchParams& params);
};
 
inline bool ScopeKindHigher(ScopeKind a, ScopeKind b) {
  return static_cast<int>(a) < static_cast<int>(b);
}
 
TVM_DLL bool ScopeNameHigher(const ffi::String& a, const ffi::String& b);
 
/******** Definition of Execution Scope ********/
class ExecScopeNode : public ffi::Object {
 public:
  ffi::Array<ScopeIdDef> scope_id_def;
 
  ScopeKind kind = ScopeKind::kKernel;
 
  ffi::String name() const { return ScopeKindToString(kind); }
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ExecScopeNode>()
        .def_ro("kind", &ExecScopeNode::kind)
        .def_ro("scope_id_def", &ExecScopeNode::scope_id_def);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO("tirx.ExecScope", ExecScopeNode, ffi::Object);
};
 
class ExecScope : public ffi::ObjectRef {
 public:
  TVM_DLL explicit ExecScope(ScopeKind kind, ffi::Array<ScopeIdDef> scope_id_def = {});
  TVM_DLL explicit ExecScope(const ffi::String& name, ffi::Array<ScopeIdDef> scope_id_def = {})
      : ExecScope(StringToScopeKind(name), std::move(scope_id_def)) {}
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(ExecScope, ffi::ObjectRef, ExecScopeNode);
};
 
}  // namespace tirx
}  // namespace tvm
 
#endif  // TVM_TIRX_EXEC_SCOPE_H_