function.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_TIR_FUNCTION_H_
#define TVM_TIR_FUNCTION_H_

#include <tvm/ir/function.h>
#include <tvm/runtime/ndarray.h>
#include <tvm/tir/buffer.h>
#include <tvm/tir/expr.h>
#include <tvm/tir/stmt.h>

#include <string>

namespace tvm {
namespace tir {

class PrimFuncNode : public BaseFuncNode {
 public:
  Array<tir::Var> params;
  tir::Stmt body;
  Type ret_type;
  Map<tir::Var, Buffer> buffer_map;

  Map<tir::Var, Buffer> preflattened_buffer_map;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("params", &params);
    v->Visit("body", &body);
    v->Visit("ret_type", &ret_type);
    v->Visit("buffer_map", &buffer_map);
    v->Visit("preflattened_buffer_map", &preflattened_buffer_map);
    v->Visit("attrs", &attrs);
    v->Visit("span", &span);
    v->Visit("_checked_type_", &checked_type_);
  }

  bool SEqualReduce(const PrimFuncNode* other, SEqualReducer equal) const {
    // visit params and buffer_map first as they contains defs.
    return equal.DefEqual(params, other->params) && equal(buffer_map, other->buffer_map) &&
           equal(preflattened_buffer_map, other->preflattened_buffer_map) &&
           equal(ret_type, other->ret_type) && equal(body, other->body) &&
           equal(attrs, other->attrs);
  }

  void SHashReduce(SHashReducer hash_reduce) const {
    hash_reduce.DefHash(params);
    hash_reduce(buffer_map);
    hash_reduce(preflattened_buffer_map);
    hash_reduce(ret_type);
    hash_reduce(body);
    hash_reduce(attrs);
  }
  TVM_DLL FuncType func_type_annotation() const;

  static constexpr const char* _type_key = "tir.PrimFunc";
  TVM_DECLARE_FINAL_OBJECT_INFO(PrimFuncNode, BaseFuncNode);
};

class PrimFunc : public BaseFunc {
 public:
  TVM_DLL PrimFunc(
      Array<tir::Var> params, Stmt body, Type ret_type = VoidType(),
      Map<tir::Var, Buffer> buffer_map = Map<tir::Var, Buffer>(),
      Optional<Map<tir::Var, Buffer>> preflattened_buffer_map = Optional<Map<tir::Var, Buffer>>(),
      DictAttrs attrs = NullValue<DictAttrs>(), Span span = Span());

  TVM_DEFINE_OBJECT_REF_METHODS(PrimFunc, BaseFunc, PrimFuncNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(PrimFuncNode);
};

class TensorIntrinNode : public Object {
 public:
  PrimFunc desc;
  PrimFunc impl;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("desc", &desc);
    v->Visit("impl", &impl);
  }

  static constexpr const char* _type_key = "tir.TensorIntrin";
  TVM_DECLARE_FINAL_OBJECT_INFO(TensorIntrinNode, Object);
};

class TensorIntrin : public ObjectRef {
 public:
  TVM_DLL explicit TensorIntrin(PrimFunc desc, PrimFunc impl);

  TVM_DLL static void Register(String name, TensorIntrin intrin, bool override = false);

  TVM_DLL static TensorIntrin Get(String name);

  TVM_DEFINE_OBJECT_REF_METHODS(TensorIntrin, ObjectRef, TensorIntrinNode)
};

/*
 * \brief Specialize parameters of PrimFunc.
 * \param func The PrimFunc to be specialized.
 * \param param_map The mapping from function params to the instance.
 * \return The new function with parameter specialized.
 * \note We can define a Meta TIR function with symbolic shape:
 *
 * \code
 *  @T.prim_func
 *  def mem_copy(a: T.handle, b: T.handle, m: T.int32, n: T.int32) -> None:
 *      A = T.match_buffer(a, (m, n), "float32")
 *      B = T.match_buffer(b, (m, n), "float32")
 *      for i, j in T.grid(m, n):
 *          with T.block():
 *              vi, vj = T.axis.remap("SS", [i, j])
 *              B[vi, vj] = A[vi, vj]
 * \endcode
 *
 * Then we can make it specialized with given shapes or buffers.
 *
 * \code
 *  a, _, m, n = mem_copy.params
 *  func = mem_copy.specialize({a: tir.decl_buffer((16, 16))})
 *  # or
 *  func = mem_copy.specialize({n: 16, m: 16})
 * \endcode
 *
 * \code {.language-id}
 *  @T.prim_func
 *  def mem_copy_16_16(a: T.handle, b: T.handle) -> None:
 *      A = T.match_buffer(a, (16, 16), "float32")
 *      B = T.match_buffer(b, (16, 16), "float32")
 *      for i, j in T.grid(16, 16):
 *          with T.block():
 *              vi, vj = T.axis.remap("SS", [i, j])
 *              B[vi, vj] = A[vi, vj]
 * \endcode
 */
PrimFunc Specialize(PrimFunc func, const Map<Var, ObjectRef>& param_map);

namespace attr {
constexpr const char* kDeviceThreadAxis = "tir.device_thread_axis";

constexpr const char* kDeviceUseDynSharedMemory = "tir.device_use_dyn_shared_memory";

constexpr const char* kNoAlias = "tir.noalias";

constexpr const char* kIsEntryFunc = "tir.is_entry_func";

constexpr const char* kIsGlobalFunc = "tir.is_global_func";

}  // namespace attr
}  // namespace tir
}  // namespace tvm
#endif  // TVM_TIR_FUNCTION_H_