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_S_TIR_SCRIPT_IR_BUILDER_IR_H_
#define TVM_S_TIR_SCRIPT_IR_BUILDER_IR_H_
 
#include <tvm/s_tir/script/ir_builder/frame.h>
#include <tvm/tirx/script/ir_builder/ir.h>
 
namespace tvm {
namespace script {
namespace ir_builder {
namespace s_tir {
 
using tvm::Var;
using tvm::tirx::Layout;
using tvm::tirx::TensorVar;
 
FunctionFrame Function(bool is_private = false, bool persistent = false);
FunctionFrame DeclFunction(bool is_private = false, bool persistent = false);
 
TensorVar MatchTensor(ffi::ObjectRef param, ffi::Array<PrimExpr> shape,
                      PrimType dtype = PrimType::Float(32), ffi::Optional<Expr> data = std::nullopt,
                      ffi::Array<PrimExpr> strides = {},
                      ffi::Optional<PrimExpr> elem_offset = std::nullopt,
                      ffi::String storage_scope = "global", int align = -1, int offset_factor = 0,
                      ffi::Optional<Layout> layout = std::nullopt,
                      ffi::Array<PrimExpr> allocated_addr = {});
 
SBlockFrame Block(ffi::String name, bool no_realize = false, ffi::String exec_scope = "");
 
BlockInitFrame Init();
 
void Where(PrimExpr predicate);
 
void Reads(ffi::Array<ffi::ObjectRef> tensor_slices);
 
void Writes(ffi::Array<ffi::ObjectRef> tensor_slices);
 
void BlockAttrs(ffi::Map<ffi::String, ffi::Any> attrs);
 
TensorVar SBlockAllocTensor(ffi::Array<PrimExpr> shape, PrimType dtype = PrimType::Float(32),
                            ffi::Optional<Expr> data = std::nullopt,
                            ffi::Array<PrimExpr> strides = {},
                            ffi::Optional<PrimExpr> elem_offset = std::nullopt,
                            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, ffi::Optional<PrimExpr> binding = std::nullopt,
            PrimType dtype = PrimType::Int(32));
 
Var Reduce(Range dom, ffi::Optional<PrimExpr> binding = std::nullopt,
           PrimType dtype = PrimType::Int(32));
 
Var Scan(Range dom, ffi::Optional<PrimExpr> binding = std::nullopt,
         PrimType dtype = PrimType::Int(32));
 
Var Opaque(Range dom, ffi::Optional<PrimExpr> binding = std::nullopt,
           PrimType dtype = PrimType::Int(32));
 
ffi::Array<Var> Remap(ffi::String kinds, ffi::Array<PrimExpr> bindings,
                      PrimType dtype = PrimType::Int(32));
 
}  // namespace axis
 
}  // namespace s_tir
}  // namespace ir_builder
}  // namespace script
}  // namespace tvm
 
#endif  // TVM_S_TIR_SCRIPT_IR_BUILDER_IR_H_