Skip to content

Commit 17e7e3d

Browse files
alex-weavertqchen
alex-weaver
authored and
tqchen
committed
Port build_module.py to C++ (apache#667)
* Port build_module.py to C++ * Fix lint errors * Fix more lint errors * Fix more lint errors * Fix more lint errors * Fix build error * Implemented style fixes * Fix lint errors * Added function to construct target from string lower now returns array * Fix lint error * Implemented review changes - style & Target options -> std::vector * Fixed lint, argument alignment and added unit test * Changed test to target LLVM, fixed sign compare warnings * Reverted unit test to CUDA, changed Jenkinsfile to enable GPU for C++ tests * Slight change to Jenkinsfile * Changed build_module test from CUDA to LLVM * Added function var() to construct a Var instance. Changed implementation of LLVMEnabled() * Reverted Jenkinsfile
1 parent f140ce4 commit 17e7e3d

File tree

6 files changed

+525
-5
lines changed

6 files changed

+525
-5
lines changed

‎include/tvm/build_module.h‎

Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
/*!
2+
* Copyright (c) 2017 by Contributors
3+
* \file build_module.h
4+
* \brief Functions for compiling ops.
5+
*/
6+
#ifndef TVM_BUILD_MODULE_H_
7+
#define TVM_BUILD_MODULE_H_
8+
9+
#include <string>
10+
#include <vector>
11+
#include "./tvm/runtime/packed_func.h"
12+
#include "./tvm/schedule_pass.h"
13+
#include "./tvm/lowered_func.h"
14+
15+
namespace tvm {
16+
17+
/*!
18+
* \brief Container for target device information.
19+
* Use target::llvm, target::cuda etc functions instead of constructing directly.
20+
*/
21+
struct Target {
22+
/*! \brief The name of the target device */
23+
std::string target_name;
24+
/*! \brief The type of the target device */
25+
DLDeviceType device_type;
26+
/*! \brief The maximum threads that a schedule should use for this device */
27+
int max_num_threads = 1;
28+
/*! \brief The warp size that should be used by the LowerThreadAllreduce pass */
29+
int thread_warp_size = 1;
30+
/*! \brief Keys for this target */
31+
std::unordered_set<std::string> keys;
32+
/*! \brief Options for this target */
33+
std::vector<std::string> options;
34+
35+
Target(const std::string& target_name,
36+
DLDeviceType device_type,
37+
int max_num_threads,
38+
int thread_warp_size,
39+
const std::unordered_set<std::string>& keys,
40+
const std::vector<std::string>& options) :
41+
target_name(target_name),
42+
device_type(device_type),
43+
max_num_threads(max_num_threads),
44+
thread_warp_size(thread_warp_size),
45+
keys(keys),
46+
options(options) {
47+
}
48+
49+
/*! \return the full device string to pass to codegen::Build */
50+
EXPORT std::string str() const;
51+
52+
/*!
53+
* \brief Create a Target given a string
54+
* \param target_str the string to parse
55+
*/
56+
EXPORT static Target create(const std::string& target_str);
57+
};
58+
59+
/*! \brief This namespace provides functions to construct Target instances */
60+
namespace target {
61+
/*! \return A target for LLVM */
62+
EXPORT Target llvm();
63+
64+
/*! \return A target for CUDA */
65+
EXPORT Target cuda();
66+
67+
/*! \return A target for ROCm */
68+
EXPORT Target rocm();
69+
70+
/*! \return A target for Metal */
71+
EXPORT Target metal();
72+
73+
/*! \return A target for rasp */
74+
EXPORT Target rasp();
75+
76+
/*! \return A target for stackvm */
77+
EXPORT Target stackvm();
78+
79+
} // namespace target
80+
81+
/*!
82+
* \brief Container for build configuration options
83+
*/
84+
struct BuildConfig {
85+
/*!
86+
* \brief The data alignment to use when constructing buffers. If this is set to
87+
* -1, then TVM's internal default will be used
88+
*/
89+
int data_alignment = -1;
90+
/*!
91+
* \brief The offset factor to use when constructing buffers. If this is set to
92+
* 0, then the offset field is not used.
93+
*/
94+
int offset_factor = 0;
95+
96+
/*!
97+
* \brief Splitting factor for loop splitting. If this is set to zero, no splitting will be
98+
* done. Otherwise, a split will be done with this factor and the inner loop will be unrolled.
99+
*/
100+
int double_buffer_split_loop = 1;
101+
/*! \brief Threshold of number of steps in the loop to be automatically unrolled */
102+
int auto_unroll_max_step = 0;
103+
/*! \brief The maximum nested level of loops that can be automatically unrolled */
104+
int auto_unroll_max_depth = 8;
105+
/*! \brief The maximum extent of loop that will be unrolled */
106+
int auto_unroll_max_extent = 0;
107+
/*!
108+
* \brief Whether to explicitly unroll the loop. If set to false, the unroll hint will
109+
* be passed to the CodeGen phase. Set to true if CodeGen supports unroll pragma.
110+
*/
111+
bool unroll_explicit = true;
112+
113+
/*! \brief Set to true if buffer arguments do not overlap. This enables more optimization. */
114+
bool restricted_func = true;
115+
116+
/*! \brief Whether to detect global barrier */
117+
bool detect_global_barrier = false;
118+
119+
BuildConfig() {
120+
}
121+
};
122+
123+
/*!
124+
* \brief Build a LoweredFunc given a schedule, args and binds
125+
* \param sch The schedule to lower.
126+
* \param args The arguments to the function.
127+
* \param name The name of the lowered function.
128+
* \param binds Buffer assignments.
129+
* \param config The build configuration.
130+
* \return The lowered function.
131+
*/
132+
EXPORT Array<LoweredFunc> lower(Schedule sch,
133+
const Array<Tensor>& args,
134+
const std::string& name,
135+
const std::unordered_map<Tensor, Buffer>& binds,
136+
const BuildConfig& config);
137+
138+
/*!
139+
* \brief Build a device and host module for a specific target from an array of lowered functions.
140+
* \param funcs The functions to be built.
141+
* \param target The target device to build for.
142+
* \param target_host The target for building host code. If null, a suitable default will be used.
143+
* \param config The build configuration.
144+
* \return The built module.
145+
*/
146+
EXPORT runtime::Module build(const Array<LoweredFunc>& funcs,
147+
const Target& target,
148+
Target* target_host,
149+
const BuildConfig& config);
150+
151+
} // namespace tvm
152+
153+
#endif // TVM_BUILD_MODULE_H_

‎include/tvm/expr.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -291,6 +291,13 @@ inline const char* IterVarType2String(IterVarType t) {
291291
return "Unknown";
292292
}
293293

294+
/*!
295+
* \brief Construct a new Var expression
296+
* \param name_hint The name hint for the expression
297+
* \param t The type of the expression
298+
*/
299+
TVM_DLL Var var(const std::string& name_hint, Type t = Int(32));
300+
294301
/*
295302
* \brief Template function to convert Map to unordered_map
296303
* Sometimes useful for API gluing when internal uses unordered_map

‎include/tvm/schedule.h‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ class Stage : public NodeRef {
8181
* \param thread_ivar The thread axis to be binded.
8282
* \return reference to self.
8383
*/
84-
Stage& bind(IterVar ivar, IterVar thread_ivar);
84+
EXPORT Stage& bind(IterVar ivar, IterVar thread_ivar);
8585
/*!
8686
* \brief Set predicate under which store to the array can be performed.
8787
* Use this when there are duplicated threads doing the same store and we only
@@ -110,7 +110,7 @@ class Stage : public NodeRef {
110110
* \param p_inner The result inner domain.
111111
* \return reference to self.
112112
*/
113-
Stage& split(IterVar parent, Expr factor, IterVar* p_outer, IterVar* p_inner); // NOLINT(*)
113+
EXPORT Stage& split(IterVar parent, Expr factor, IterVar* p_outer, IterVar* p_inner); // NOLINT(*)
114114
/*!
115115
* \brief Split the iteration with given number of parts.
116116
*
@@ -248,13 +248,13 @@ class Schedule : public NodeRef {
248248
* \brief Get the stage corresponds to the op
249249
* \param op The operation.
250250
*/
251-
Stage operator[](const Operation& op);
251+
EXPORT Stage operator[](const Operation& op);
252252
/*!
253253
* \brief Short hand for getting the stage of tensor's operation.
254254
* \param tensor The tensor
255255
* \return The stage corresponding to the tensor's op
256256
*/
257-
Stage operator[](const Tensor& tensor) {
257+
EXPORT Stage operator[](const Tensor& tensor) {
258258
return this->operator[](tensor->op);
259259
}
260260
/*!
@@ -493,7 +493,7 @@ class ScheduleNode : public Node {
493493
* \param ops The ops to be scheduled.
494494
* \return sch The created Schedule.
495495
*/
496-
static Schedule make(Array<Operation> ops);
496+
EXPORT static Schedule make(Array<Operation> ops);
497497

498498
static constexpr const char* _type_key = "Schedule";
499499
TVM_DECLARE_NODE_TYPE_INFO(ScheduleNode, Node);

0 commit comments

Comments
 (0)