layout.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_LAYOUT_H_
#define TVM_TIRX_LAYOUT_H_
 
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/tuple.h>
#include <tvm/ffi/enum.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/object.h>
#include <tvm/ir/expr.h>
#include <tvm/ir/module.h>
#include <tvm/target/target.h>
#include <tvm/tirx/tile/exec_scope.h>
 
namespace tvm {
 
namespace tirx {
 
class Layout;
class TileLayout;
class Iter;
using ffi::Array;
using ffi::Tuple;
 
// Base class for layout
class LayoutNode : public ffi::Object {
 public:
  virtual bool CompatibleWithShape(const ffi::Array<PrimExpr>& shape) const = 0;
 
  virtual bool VerifyWellFormed() const = 0;
 
  virtual PrimExpr GetSize(ffi::Optional<ffi::String> axis_name = std::nullopt) const = 0;
 
  virtual PrimExpr GetSpan(ffi::Optional<ffi::String> axis_name = std::nullopt) const = 0;
 
  virtual ffi::Map<ffi::String, PrimExpr> Apply(ffi::Array<PrimExpr> coord) const = 0;
  virtual ffi::Map<ffi::String, PrimExpr> Apply(PrimExpr coord) const = 0;
  virtual ffi::Map<ffi::String, PrimExpr> Apply(const ffi::Array<PrimExpr>& coord,
                                                const ffi::Array<PrimExpr>& shape) const;
 
  virtual Layout Canonicalize() const = 0;
 
  virtual Layout Tile(const TileLayout& outer, const ffi::Array<PrimExpr>& outer_shape,
                      const ffi::Array<PrimExpr>& inner_shape) const = 0;
 
  virtual ffi::Optional<Layout> Slice(const ffi::Array<PrimExpr>& shape,
                                      const ffi::Array<Range>& region) const = 0;
 
  virtual Layout DirectSum(const TileLayout& left, const ffi::Array<PrimExpr>& left_shape,
                           const ffi::Array<PrimExpr>& right_shape) const = 0;
 
  virtual ffi::Optional<TileLayout> IsTileInner(const Layout& tile_layout,
                                                const ffi::Array<PrimExpr>& tiled_shape,
                                                const ffi::Array<PrimExpr>& inner_shape) const = 0;
 
  virtual ffi::Optional<Layout> IsTileOuter(const Layout& tile_layout,
                                            const ffi::Array<PrimExpr>& tiled_shape,
                                            const ffi::Array<PrimExpr>& outer_shape) const = 0;
 
  virtual ffi::Optional<TileLayout> IsDirectSumRight(
      const Layout& sum_layout, const ffi::Array<PrimExpr>& interleaved_shape,
      const ffi::Array<PrimExpr>& right_shape) const = 0;
 
  virtual ffi::Optional<Layout> IsDirectSumLeft(const Layout& sum_layout,
                                                const ffi::Array<PrimExpr>& interleaved_shape,
                                                const ffi::Array<PrimExpr>& left_shape) const = 0;
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO("tirx.Layout", LayoutNode, ffi::Object);
};
 
class Layout : public ffi::ObjectRef {
 public:
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Layout, ffi::ObjectRef, LayoutNode);
};
 
// target, subscope, scope, iter -> fused_iter
using FAxisFuser = ffi::TypedFunction<ffi::Optional<Iter>(Target, ffi::String, ffi::String, Iter)>;
// target, scope, iter -> (outer_iter, inner_iter)
// Note(@bohao): use ffi::Array<Iter, void> to avoid incomplete type error (SFINAE)
using FAxisSplitter = ffi::TypedFunction<ffi::Array<Iter, void>(Target, ffi::String, Iter)>;
 
// Axis
class AxisNode : public ffi::EnumObj {
 public:
  bool IsThreadAxis() const;
 
  bool IsMemoryAxis() const;
 
  ffi::Optional<ffi::String> GetBackend() const;
 
  ffi::Optional<ExecScope> GetScope() const;
 
  ffi::Optional<ExecScope> GetSubscope() const;
 
  ffi::Optional<FAxisFuser> GetFuser() const;
 
  ffi::Optional<FAxisSplitter> GetSplitter() const;
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Axis", AxisNode, ffi::EnumObj);
};
 
class Axis : public ffi::Enum {
 public:
  const ffi::String& name() const { return get()->_str_index; }
 
  TVM_TIRX_DLL static Axis Get(const ffi::String& name);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Axis, ffi::Enum, AxisNode);
};
 
class IterNode : public ffi::Object {
 public:
  explicit IterNode(ffi::UnsafeInit tag) : extent(tag), stride(tag), axis(tag) {}
 
  IterNode(PrimExpr extent, PrimExpr stride, Axis axis)
      : extent(std::move(extent)), stride(std::move(stride)), axis(std::move(axis)) {}
 
  PrimExpr extent;
  PrimExpr stride;
  Axis axis;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<IterNode>()
        .def_ro("extent", &IterNode::extent)
        .def_ro("stride", &IterNode::stride)
        .def_ro("axis", &IterNode::axis);
  }
 
  static constexpr TVMFFISEqHashKind _type_s_eq_hash_kind = kTVMFFISEqHashKindTreeNode;
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.Iter", IterNode, ffi::Object);
};
 
class Iter : public ffi::ObjectRef {
 public:
  TVM_TIRX_DLL explicit Iter(PrimExpr extent, PrimExpr stride, Axis axis);
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(Iter, ffi::ObjectRef, IterNode);
};
 
class TileLayoutNode : public LayoutNode {
 public:
  ffi::Array<Iter> shard;
  ffi::Array<Iter> replica;
  ffi::Map<Axis, PrimExpr> offset;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<TileLayoutNode>()
        .def_ro("shard", &TileLayoutNode::shard)
        .def_ro("replica", &TileLayoutNode::replica)
        .def_ro("offset", &TileLayoutNode::offset);
  }
 
  bool CompatibleWithShape(const ffi::Array<PrimExpr>& shape) const final;
 
  bool VerifyWellFormed() const final;
 
  PrimExpr GetSize(ffi::Optional<ffi::String> axis_name = std::nullopt) const final;
 
  PrimExpr GetSpan(ffi::Optional<ffi::String> axis_name = std::nullopt) const final;
 
  ffi::Map<ffi::String, PrimExpr> Apply(ffi::Array<PrimExpr> coord) const final;
  ffi::Map<ffi::String, PrimExpr> Apply(PrimExpr coord) const final;
  ffi::Map<ffi::String, PrimExpr> Apply(const ffi::Array<PrimExpr>& coord,
                                        const ffi::Array<PrimExpr>& shape) const final;
 
  Layout Canonicalize() const final;
 
  Layout Tile(const TileLayout& outer, const ffi::Array<PrimExpr>& outer_shape,
              const ffi::Array<PrimExpr>& inner_shape) const final;
 
  Layout DirectSum(const TileLayout& left, const ffi::Array<PrimExpr>& left_shape,
                   const ffi::Array<PrimExpr>& right_shape) const final;
 
  ffi::Optional<TileLayout> IsTileInner(const Layout& tile_layout,
                                        const ffi::Array<PrimExpr>& tiled_shape,
                                        const ffi::Array<PrimExpr>& inner_shape) const final;
 
  ffi::Optional<Layout> IsTileOuter(const Layout& tile_layout,
                                    const ffi::Array<PrimExpr>& tiled_shape,
                                    const ffi::Array<PrimExpr>& outer_shape) const final;
 
  ffi::Optional<TileLayout> IsDirectSumRight(const Layout& sum_layout,
                                             const ffi::Array<PrimExpr>& interleaved_shape,
                                             const ffi::Array<PrimExpr>& right_shape) const final;
 
  ffi::Optional<Layout> IsDirectSumLeft(const Layout& sum_layout,
                                        const ffi::Array<PrimExpr>& interleaved_shape,
                                        const ffi::Array<PrimExpr>& left_shape) const final;
 
  ffi::Array<PrimExpr> GetShardShape() const;
 
  ffi::Optional<Layout> Slice(const ffi::Array<PrimExpr>& shape,
                              const ffi::Array<Range>& region) const final;
 
  bool IsTrivial() const;
 
  bool HasMemoryAxis() const;
 
  bool HasThreadAxis() const;
 
  ffi::Optional<Tuple<ExecScope, ExecScope>> GetScope() const;
 
  static TileLayout DefaultLayout(ffi::Array<PrimExpr> shape);
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.TileLayout", TileLayoutNode, LayoutNode);
};
 
class TileLayout : public Layout {
 public:
  TVM_TIRX_DLL explicit TileLayout(ffi::Array<Iter> shard, ffi::Array<Iter> replica,
                                   ffi::Map<Axis, PrimExpr> offset);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(TileLayout, Layout, TileLayoutNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(TileLayoutNode);
};
 
// ComposeLayout
class ComposeLayoutNode : public LayoutNode {
 public:
  explicit ComposeLayoutNode(TileLayout tile_layout) : tile_layout(std::move(tile_layout)) {}
  explicit ComposeLayoutNode(ffi::UnsafeInit) : tile_layout(ffi::UnsafeInit{}) {}
 
  int per_element;
  int swizzle_len;
  int atom_len;
  bool swizzle_inner;
  TileLayout tile_layout;
  // Cached swizzle masks: inner_mask = (1 << swizzle_len) - 1;
  // outer_mask = inner_mask << atom_len. Set by the ComposeLayout ctor.
  int inner_mask;
  int outer_mask;
 
  static void RegisterReflection() {
    namespace refl = tvm::ffi::reflection;
    refl::ObjectDef<ComposeLayoutNode>()
        .def_ro("per_element", &ComposeLayoutNode::per_element)
        .def_ro("swizzle_len", &ComposeLayoutNode::swizzle_len)
        .def_ro("atom_len", &ComposeLayoutNode::atom_len)
        .def_ro("swizzle_inner", &ComposeLayoutNode::swizzle_inner)
        .def_ro("inner_mask", &ComposeLayoutNode::inner_mask)
        .def_ro("outer_mask", &ComposeLayoutNode::outer_mask)
        .def_ro("tile_layout", &ComposeLayoutNode::tile_layout);
  }
 
  bool CompatibleWithShape(const ffi::Array<PrimExpr>& shape) const final;
 
  bool VerifyWellFormed() const final;
 
  PrimExpr GetSize(ffi::Optional<ffi::String> axis_name = std::nullopt) const final;
 
  PrimExpr GetSpan(ffi::Optional<ffi::String> axis_name = std::nullopt) const final;
 
  ffi::Map<ffi::String, PrimExpr> Apply(ffi::Array<PrimExpr> coord) const final;
  ffi::Map<ffi::String, PrimExpr> Apply(PrimExpr coord) const final;
  ffi::Map<ffi::String, PrimExpr> Apply(const ffi::Array<PrimExpr>& coord,
                                        const ffi::Array<PrimExpr>& shape) const final;
 
  Layout Canonicalize() const final;
 
  Layout Tile(const TileLayout& outer, const ffi::Array<PrimExpr>& outer_shape,
              const ffi::Array<PrimExpr>& inner_shape) const final;
 
  Layout DirectSum(const TileLayout& left, const ffi::Array<PrimExpr>& left_shape,
                   const ffi::Array<PrimExpr>& right_shape) const final;
 
  ffi::Optional<TileLayout> IsTileInner(const Layout& tile_layout,
                                        const ffi::Array<PrimExpr>& tiled_shape,
                                        const ffi::Array<PrimExpr>& inner_shape) const final;
 
  ffi::Optional<Layout> IsTileOuter(const Layout& tile_layout,
                                    const ffi::Array<PrimExpr>& tiled_shape,
                                    const ffi::Array<PrimExpr>& outer_shape) const final;
 
  ffi::Optional<TileLayout> IsDirectSumRight(const Layout& sum_layout,
                                             const ffi::Array<PrimExpr>& interleaved_shape,
                                             const ffi::Array<PrimExpr>& right_shape) const final;
 
  ffi::Optional<Layout> IsDirectSumLeft(const Layout& sum_layout,
                                        const ffi::Array<PrimExpr>& interleaved_shape,
                                        const ffi::Array<PrimExpr>& left_shape) const final;
 
  ffi::Optional<Layout> Slice(const ffi::Array<PrimExpr>& shape,
                              const ffi::Array<Range>& region) const final;
 
  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("tirx.ComposeLayout", ComposeLayoutNode, LayoutNode);
};
 
class ComposeLayout : public Layout {
 public:
  TVM_TIRX_DLL explicit ComposeLayout(int per_element, int swizzle_len, int atom_len,
                                      TileLayout tile_layout, bool swizzle_inner = true);
 
  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(ComposeLayout, Layout, ComposeLayoutNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(ComposeLayoutNode);
};
 
constexpr int kPSUMMaxElemPerBank = 512;
constexpr int kPSUMBankNum = 8;
 
}  // namespace tirx
}  // namespace tvm
 
#endif  // TVM_TIRX_LAYOUT_H_