ir.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_SCRIPT_IR_BUILDER_TIR_IR_H_
#define TVM_SCRIPT_IR_BUILDER_TIR_IR_H_
 
#include <tvm/ffi/container/tuple.h>
#include <tvm/ffi/container/variant.h>
#include <tvm/runtime/tensor.h>
#include <tvm/script/ir_builder/base.h>
#include <tvm/tirx/exec_scope.h>
#include <tvm/tirx/layout.h>
#include <tvm/tirx/op.h>
#include <tvm/tirx/script/builder/frame.h>
#include <tvm/tirx/tile_primitive.h>
 
namespace tvm {
namespace script {
namespace ir_builder {
namespace tirx {
 
using tvm::ffi::Tuple;
using tvm::ffi::Variant;
using tvm::runtime::Tensor;
using tvm::tirx::BufferVar;
using tvm::tirx::ExecScope;
using tvm::tirx::Layout;
using tvm::tirx::Var;
 
BufferVar BufferDecl(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::String buffer_name,
                     ffi::Optional<Expr> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
                     ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope, int align,
                     int offset_factor, ffi::Optional<Layout> layout = std::nullopt,
                     ffi::Array<PrimExpr> allocated_addr = {});
 
PrimFuncFrame PrimFunc(bool is_private, bool s_tir = false, bool persistent = false);
 
Var Arg(ffi::String name, Var var);
 
BufferVar Arg(ffi::String name, BufferVar buffer);
 
void FuncName(ffi::String name);
 
void FuncAttrs(ffi::Map<ffi::String, ffi::Any> attrs);
 
Type FuncRet(Type ret_type);
 
BufferVar MatchBuffer(ffi::ObjectRef param, ffi::Array<PrimExpr> shape,
                      PrimType dtype = PrimType::Float(32), ffi::Optional<Expr> data = std::nullopt,
                      ffi::Array<PrimExpr> strides = {}, PrimExpr elem_offset = PrimExpr(),
                      ffi::String storage_scope = "global", int align = -1, int offset_factor = 0,
                      ffi::Optional<Layout> layout = std::nullopt);
 
SBlockFrame Block(ffi::String name, bool no_realize = false, ffi::String exec_scope = "");
 
void TilePrimitiveCall(tvm::tirx::TilePrimitiveCall op_call);
 
ffi::Array<tvm::tirx::Var> ScopeId(ffi::Optional<ffi::Array<PrimExpr>> extents, ffi::String parent,
                                   ffi::String name, ffi::String cur,
                                   PrimType dtype = PrimType::Int(32));
 
ffi::Array<tvm::tirx::Var> ClusterId(ffi::Optional<ffi::Array<PrimExpr>> extents,
                                     ffi::String parent, PrimType dtype = PrimType::Int(32));
 
ffi::Array<tvm::tirx::Var> CtaId(ffi::Optional<ffi::Array<PrimExpr>> extents, ffi::String parent,
                                 ffi::Optional<ffi::Array<PrimExpr>> preferred = std::nullopt,
                                 PrimType dtype = PrimType::Int(32));
 
ffi::Array<tvm::tirx::Var> CtaIdInPair(PrimType dtype = PrimType::Int(32));
 
ffi::Array<tvm::tirx::Var> WarpgroupId(ffi::Optional<ffi::Array<PrimExpr>> extents,
                                       ffi::String parent, PrimType dtype = PrimType::Int(32));
 
ffi::Array<tvm::tirx::Var> WarpId(ffi::Optional<ffi::Array<PrimExpr>> extents, ffi::String parent,
                                  PrimType dtype = PrimType::Int(32));
 
ffi::Array<tvm::tirx::Var> ThreadId(ffi::Optional<ffi::Array<PrimExpr>> extents, ffi::String parent,
                                    PrimType dtype = PrimType::Int(32));
 
BlockInitFrame Init();
 
void Where(PrimExpr predicate);
 
void Reads(ffi::Array<ffi::ObjectRef> buffer_slices);
 
void Writes(ffi::Array<ffi::ObjectRef> buffer_slices);
 
void BlockAttrs(ffi::Map<ffi::String, ffi::Any> attrs);
 
ffi::Variant<BufferVar, AllocBufferFrame> SBlockAllocBuffer(
    ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
    ffi::Optional<Expr> data = std::nullopt, ffi::Array<PrimExpr> strides = {},
    PrimExpr elem_offset = PrimExpr(), ffi::String storage_scope = "", int align = -1,
    int offset_factor = 0, ffi::Optional<Layout> layout = std::nullopt,
    ffi::Array<PrimExpr> allocated_addr = {});
 
namespace axis {
 
Var Spatial(Range dom, PrimExpr binding, PrimType dtype = PrimType::Int(32));
 
Var Reduce(Range dom, PrimExpr binding, PrimType dtype = PrimType::Int(32));
 
Var Scan(Range dom, PrimExpr binding, PrimType dtype = PrimType::Int(32));
 
Var Opaque(Range dom, PrimExpr binding, PrimType dtype = PrimType::Int(32));
 
ffi::Array<Var> Remap(ffi::String kinds, ffi::Array<PrimExpr> bindings,
                      PrimType dtype = PrimType::Int(32));
 
}  // namespace axis
 
ForFrame Serial(PrimExpr start, PrimExpr stop,
                ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
                ffi::Optional<PrimExpr> step = std::nullopt,
                ffi::Optional<PrimType> dtype = std::nullopt);
ForFrame Parallel(PrimExpr start, PrimExpr stop,
                  ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
                  ffi::Optional<PrimExpr> step = std::nullopt,
                  ffi::Optional<PrimType> dtype = std::nullopt);
ForFrame Vectorized(PrimExpr start, PrimExpr stop,
                    ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
                    ffi::Optional<PrimExpr> step = std::nullopt,
                    ffi::Optional<PrimType> dtype = std::nullopt);
ForFrame Unroll(PrimExpr start, PrimExpr stop,
                ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
                ffi::Optional<PrimExpr> step = std::nullopt,
                ffi::Optional<PrimType> dtype = std::nullopt);
ForFrame ThreadBinding(PrimExpr start, PrimExpr stop, ffi::String thread,
                       ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt);
ForFrame Grid(ffi::Array<Variant<PrimExpr, ffi::Tuple<PrimExpr, PrimExpr>>> extents,
              ffi::Optional<PrimType> dtype = std::nullopt);
 
AssertFrame Assert(PrimExpr condition, ffi::String error_kind,
                   ffi::Array<ffi::String> message_parts);
 
Var Bind(Expr value, ffi::Optional<Type> type_annotation = std::nullopt,
         ffi::Optional<Var> var = std::nullopt);
 
AttrFrame Attr(ffi::Any node, ffi::String attr_key, PrimExpr value);
 
AttrFrame DeviceEntry();
 
WhileFrame While(PrimExpr condition);
 
void Return(Expr value);
 
void Break();
 
void Continue();
 
IfFrame If(PrimExpr condition);
 
ThenFrame Then();
 
ElseFrame Else();
 
DeclBufferFrame DeclBuffer(ffi::Array<PrimExpr> shape, PrimType dtype, ffi::String buffer_name,
                           ffi::Optional<Expr> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
                           ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope,
                           int align, int offset_factor,
                           ffi::Optional<Layout> layout = std::nullopt,
                           ffi::Optional<PrimExpr> allocated_addr = std::nullopt);
 
BufferVar AllocBuffer(ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
                      ffi::String storage_scope = "global",
                      ffi::Optional<ffi::Map<ffi::String, ffi::Any>> annotations = std::nullopt);
 
LaunchThreadFrame LaunchThread(Var var, PrimExpr extent);
 
LaunchThreadFrame LaunchThread(ffi::String thread_tag, PrimExpr extent);
 
ComposeOpFrame ComposeOp(ffi::Map<ffi::String, BufferVar> workspace,
                         ffi::Map<ffi::String, ffi::Any> config,
                         ffi::Optional<ffi::String> dispatch = std::nullopt);
 
Var EnvThread(ffi::String thread_tag, PrimType dtype = PrimType::Int(32));
 
void BufferStore(BufferVar buffer, PrimExpr value, ffi::Array<PrimExpr> indices);
 
void Evaluate(Expr value);
 
inline Var Handle(ffi::Optional<PrimType> dtype = std::nullopt,
                  ffi::String storage_scope = "global") {
  Type type_annotation = dtype.has_value() ? Type(PointerType(dtype.value(), storage_scope))
                                           : Type(PointerType::VoidPointerTy(storage_scope));
  return tvm::tirx::Var("", type_annotation);
}
 
inline Var TensorMap() { return tvm::tirx::Var("", PointerType(TensorMapType())); }
 
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName, DType)                      \
  inline PrimExpr FuncName(ffi::Optional<PrimExpr> expr = std::nullopt) {        \
    PrimType dtype = DType;                                                      \
    return expr.has_value() ? tvm::cast(dtype, expr.value())                     \
                            : tvm::tirx::Var("", dtype).as_or_throw<PrimExpr>(); \
  }
 
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(DType, Code)                         \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##8, (PrimType(DLDataType{Code, 8, 1})));   \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##16, (PrimType(DLDataType{Code, 16, 1}))); \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##32, (PrimType(DLDataType{Code, 32, 1}))); \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##64, (PrimType(DLDataType{Code, 64, 1})));
 
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(BFloat, kDLBfloat);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(Float, kDLFloat);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(UInt, kDLUInt);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(Int, kDLInt);
 
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(FuncName, Code, Size)                       \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x2, (PrimType(DLDataType{Code, Size, 2})))    \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x4, (PrimType(DLDataType{Code, Size, 4})));   \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x8, (PrimType(DLDataType{Code, Size, 8})));   \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x16, (PrimType(DLDataType{Code, Size, 16}))); \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x32, (PrimType(DLDataType{Code, Size, 32}))); \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x64, (PrimType(DLDataType{Code, Size, 64})));
 
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(DType, Code) \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##8, Code, 8);      \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##16, Code, 16);    \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##32, Code, 32);    \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##64, Code, 64);
 
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(BFloat, kDLBfloat);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(Float, kDLFloat);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(UInt, kDLUInt);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(Int, kDLInt);
 
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(DType, Code, Bits)            \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType, (PrimType(DLDataType{Code, Bits, 1})));       \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x2, (PrimType(DLDataType{Code, Bits, 2})));   \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x4, (PrimType(DLDataType{Code, Bits, 4})));   \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x8, (PrimType(DLDataType{Code, Bits, 8})));   \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x16, (PrimType(DLDataType{Code, Bits, 16}))); \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x32, (PrimType(DLDataType{Code, Bits, 32}))); \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x64, (PrimType(DLDataType{Code, Bits, 64})));
 
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E3M4, kDLFloat8_e3m4, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3, kDLFloat8_e4m3, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3B11FNUZ, kDLFloat8_e4m3b11fnuz, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3FN, kDLFloat8_e4m3fn, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3FNUZ, kDLFloat8_e4m3fnuz, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E5M2, kDLFloat8_e5m2, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E5M2FNUZ, kDLFloat8_e5m2fnuz, 8);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E8M0FNU, kDLFloat8_e8m0fnu, 8);
 
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float6E2M3FN, kDLFloat6_e2m3fn, 6);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float6E3M2FN, kDLFloat6_e3m2fn, 6);
 
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float4E2M1FN, kDLFloat4_e2m1fn, 4);
 
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(Boolean, PrimType::Bool());
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(Void, PrimType::Void());
 
#undef TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST
 
}  // namespace tirx
}  // namespace ir_builder
}  // namespace script
}  // namespace tvm
 
#endif  // TVM_TIRX_SCRIPT_BUILDER_IR_H_