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/tirx_stmt.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::Buffer;
using tvm::tirx::ExecScope;
using tvm::tirx::Layout;
using tvm::tirx::Var;
 
Buffer BufferDecl(ffi::Array<PrimExpr> shape, DataType dtype, ffi::String buffer_name,
                  ffi::Optional<Var> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
                  ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope, int align,
                  int offset_factor, ffi::String buffer_type,
                  ffi::Optional<ffi::Array<IntImm>> axis_separators,
                  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);
 
Buffer Arg(ffi::String name, Buffer buffer);
 
void FuncName(ffi::String name);
 
void FuncAttrs(ffi::Map<ffi::String, ffi::Any> attrs);
 
Type FuncRet(Type ret_type);
 
Buffer MatchBuffer(ffi::ObjectRef param, ffi::Array<PrimExpr> shape,
                   DataType dtype = DataType::Float(32), ffi::Optional<Var> data = std::nullopt,
                   ffi::Array<PrimExpr> strides = {}, PrimExpr elem_offset = PrimExpr(),
                   ffi::String storage_scope = "global", int align = -1, int offset_factor = 0,
                   ffi::String buffer_type = "default",
                   ffi::Optional<ffi::Array<IntImm>> axis_separators = std::nullopt,
                   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);
 
ExecScopeFrame ExecScopeBlock(ffi::String exec_scope_name,
                              ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
 
ExecScopeFrame Kernel(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
ExecScopeFrame Cluster(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
ExecScopeFrame WarpGroup(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
ExecScopeFrame CTA(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
ExecScopeFrame Warp(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
ExecScopeFrame Thread(ffi::Array<PrimExpr> guards = ffi::Array<PrimExpr>());
 
ffi::Array<tvm::tirx::Var> KernelId(ffi::Array<PrimExpr> extents, ffi::String parent);
 
ffi::Array<tvm::tirx::Var> CtaId(ffi::Array<PrimExpr> extents, ffi::String parent);
 
ffi::Array<tvm::tirx::Var> CtaIdInPair();
 
ffi::Array<tvm::tirx::Var> WarpId(ffi::Array<PrimExpr> extents, ffi::String parent);
 
ffi::Array<tvm::tirx::Var> ThreadId(ffi::Array<PrimExpr> extents, ffi::String parent);
 
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<Buffer, AllocBufferFrame> SBlockAllocBuffer(
    ffi::Array<PrimExpr> shape, DataType dtype = DataType::Float(32),
    ffi::Optional<Var> data = std::nullopt, ffi::Array<PrimExpr> strides = {},
    PrimExpr elem_offset = PrimExpr(), ffi::String storage_scope = "", int align = -1,
    int offset_factor = 0, ffi::String buffer_type = "default",
    ffi::Optional<ffi::Array<IntImm>> axis_separators = std::nullopt,
    ffi::Optional<Layout> layout = std::nullopt, ffi::Array<PrimExpr> allocated_addr = {});
 
namespace axis {
 
Var Spatial(Range dom, PrimExpr binding, DataType dtype = DataType::Int(32));
 
Var Reduce(Range dom, PrimExpr binding, DataType dtype = DataType::Int(32));
 
Var Scan(Range dom, PrimExpr binding, DataType dtype = DataType::Int(32));
 
Var Opaque(Range dom, PrimExpr binding, DataType dtype = DataType::Int(32));
 
ffi::Array<Var> Remap(ffi::String kinds, ffi::Array<PrimExpr> bindings,
                      DataType dtype = DataType::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);
ForFrame Parallel(PrimExpr start, PrimExpr stop,
                  ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
                  ffi::Optional<PrimExpr> step = std::nullopt);
ForFrame Vectorized(PrimExpr start, PrimExpr stop,
                    ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
                    ffi::Optional<PrimExpr> step = std::nullopt);
ForFrame Unroll(PrimExpr start, PrimExpr stop,
                ffi::Optional<ffi::Map<ffi::String, Any>> annotations = std::nullopt,
                ffi::Optional<PrimExpr> step = 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);
 
AssertFrame Assert(PrimExpr condition, ffi::String error_kind,
                   ffi::Array<ffi::String> message_parts);
 
Var Bind(PrimExpr 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);
 
WhileFrame While(PrimExpr condition);
 
void Break();
 
void Continue();
 
IfFrame If(PrimExpr condition);
 
ThenFrame Then();
 
ElseFrame Else();
 
DeclBufferFrame DeclBuffer(ffi::Array<PrimExpr> shape, DataType dtype, ffi::String buffer_name,
                           ffi::Optional<Var> data, ffi::Optional<ffi::Array<PrimExpr>> strides,
                           ffi::Optional<PrimExpr> elem_offset, ffi::String storage_scope,
                           int align, int offset_factor, ffi::String buffer_type,
                           ffi::Optional<ffi::Array<IntImm>> axis_separators,
                           ffi::Optional<Layout> layout = std::nullopt,
                           ffi::Optional<PrimExpr> allocated_addr = std::nullopt);
 
Buffer AllocBuffer(ffi::Array<PrimExpr> shape, DataType dtype = DataType::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, Buffer> workspace,
                         ffi::Map<ffi::String, ffi::Any> config,
                         ffi::Optional<ffi::String> dispatch = std::nullopt);
 
Var EnvThread(ffi::String thread_tag, DataType dtype = DataType::Int(32));
 
void BufferStore(Buffer buffer, PrimExpr value, ffi::Array<PrimExpr> indices,
                 ffi::Optional<PrimExpr> predicate);
 
void Evaluate(PrimExpr value);
 
inline Var Handle(runtime::DataType dtype = runtime::DataType::Void(),
                  ffi::String storage_scope = "global", bool is_size_var = false,
                  bool is_unknown_type = false) {
  Type type_annotation{nullptr};
  if (is_unknown_type && storage_scope == "global") {
    type_annotation = PrimType(runtime::DataType::Handle());
  } else {
    type_annotation = PointerType(PrimType(dtype), storage_scope);
  }
  return is_size_var ? tvm::tirx::SizeVar("", type_annotation)
                     : 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,                     \
                           bool is_size_var = false) {                                      \
    DataType dtype = DType;                                                                 \
    return expr.defined()                                                                   \
               ? tvm::cast(dtype, expr.value())                                             \
               : (is_size_var ? tvm::tirx::SizeVar("", dtype) : tvm::tirx::Var("", dtype)); \
  }
 
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(DType, FDType) \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##8, FDType(8));      \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##16, FDType(16));    \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##32, FDType(32));    \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##64, FDType(64));
 
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(BFloat, DataType::BFloat);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(Float, DataType::Float);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(UInt, DataType::UInt);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES(Int, DataType::Int);
 
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(FuncName, FDType, Size) \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x2, FDType(Size, 2))      \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x4, FDType(Size, 4));     \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x8, FDType(Size, 8));     \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x16, FDType(Size, 16));   \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x32, FDType(Size, 32));   \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(FuncName##x64, FDType(Size, 64));
 
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(DType, FDType) \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##8, FDType, 8);      \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##16, FDType, 16);    \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##32, FDType, 32);    \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES(DType##64, FDType, 64);
 
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(BFloat, DataType::BFloat);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(Float, DataType::Float);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(UInt, DataType::UInt);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_SIZES_LANES(Int, DataType::Int);
 
#define TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(DType, FDType) \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType, FDType(1));                    \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x2, FDType(2));                \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x4, FDType(4));                \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x8, FDType(8));                \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x16, FDType(16));              \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x32, FDType(32));              \
  TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(DType##x64, FDType(64));
 
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E3M4, DataType::Float8E3M4);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3, DataType::Float8E4M3);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3B11FNUZ, DataType::Float8E4M3B11FNUZ);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3FN, DataType::Float8E4M3FN);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E4M3FNUZ, DataType::Float8E4M3FNUZ);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E5M2, DataType::Float8E5M2);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E5M2FNUZ, DataType::Float8E5M2FNUZ);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float8E8M0FNU, DataType::Float8E8M0FNU);
 
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float6E2M3FN, DataType::Float6E2M3FN);
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float6E3M2FN, DataType::Float6E3M2FN);
 
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST_LANES_FIXED_SIZE(Float4E2M1FN, DataType::Float4E2M1FN);
 
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(Boolean, DataType::Bool());
TVM_TIRX_IR_BUILDER_DEF_DTYPE_CAST(Void, DataType::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_