search_task.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_AUTO_SCHEDULER_SEARCH_TASK_H_
#define TVM_AUTO_SCHEDULER_SEARCH_TASK_H_

#include <tvm/auto_scheduler/compute_dag.h>
#include <tvm/runtime/ndarray.h>
#include <tvm/target/target.h>

namespace tvm {
namespace auto_scheduler {

class HardwareParams;

class HardwareParamsNode : public Object {
 public:
  int num_cores;
  int vector_unit_bytes;
  int cache_line_bytes;

  // GPU related parameters got from device query API
  int max_shared_memory_per_block;
  int max_local_memory_per_block;
  int max_threads_per_block;
  int max_vthread_extent;
  int warp_size;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("num_cores", &num_cores);
    v->Visit("vector_unit_bytes", &vector_unit_bytes);
    v->Visit("cache_line_bytes", &cache_line_bytes);
    v->Visit("max_shared_memory_per_block", &max_shared_memory_per_block);
    v->Visit("max_local_memory_per_block", &max_local_memory_per_block);
    v->Visit("max_threads_per_block", &max_threads_per_block);
    v->Visit("max_vthread_extent", &max_vthread_extent);
    v->Visit("warp_size", &warp_size);
  }

  static HardwareParams GetDefaultHardwareParams(const Target& target, const Target& target_host);

  static constexpr const char* _type_key = "auto_scheduler.HardwareParams";
  TVM_DECLARE_FINAL_OBJECT_INFO(HardwareParamsNode, Object);
};

class HardwareParams : public ObjectRef {
 public:
  HardwareParams(int num_cores, int vector_unit_bytes, int cache_line_bytes,
                 int max_shared_memory_per_block, int max_local_memory_per_block,
                 int max_threads_per_block, int max_vthread_extent, int warp_size);

  TVM_DEFINE_OBJECT_REF_METHODS(HardwareParams, ObjectRef, HardwareParamsNode);
  TVM_DEFINE_OBJECT_REF_COW_METHOD(HardwareParamsNode);
};

class SearchTaskNode : public Object {
 public:
  ComputeDAG compute_dag;
  String workload_key;
  String desc;
  Target target;
  Target target_host;
  HardwareParams hardware_params;
  LayoutRewriteOption layout_rewrite_option;
  Array<String> task_input_names;

  void VisitAttrs(tvm::AttrVisitor* v) {
    v->Visit("compute_dag", &compute_dag);
    v->Visit("workload_key", &workload_key);
    v->Visit("desc", &desc);
    v->Visit("target", &target);
    v->Visit("target_host", &target_host);
    v->Visit("hardware_params", &hardware_params);
    v->Visit("layout_rewrite_option", &layout_rewrite_option);
    v->Visit("task_input_names", &task_input_names);
  }

  static constexpr const char* _type_key = "auto_scheduler.SearchTask";
  TVM_DECLARE_FINAL_OBJECT_INFO(SearchTaskNode, Object);
};

class SearchTask : public ObjectRef {
 public:
  SearchTask(ComputeDAG compute_dag, String workload_key, Target target, Target target_host,
             Optional<HardwareParams> hardware_params, LayoutRewriteOption layout_rewrite_option,
             Array<String> task_input_names, String desc = "");

  TVM_DEFINE_OBJECT_REF_METHODS(SearchTask, ObjectRef, SearchTaskNode);
};

}  // namespace auto_scheduler
}  // namespace tvm

#endif  // TVM_AUTO_SCHEDULER_SEARCH_TASK_H_