Skip to content

Commit cf8521a

Browse files
authored
[Target] LLVM helper functions for any target info (#15761)
1 parent c318fa8 commit cf8521a

21 files changed

+538
-351
lines changed

cmake/modules/LLVM.cmake

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@ add_definitions(-DDMLC_USE_FOPEN64=0 -DNDEBUG=1)
2929
# It may be a boolean or a string
3030
if(NOT ${USE_LLVM} MATCHES ${IS_FALSE_PATTERN})
3131
find_llvm(${USE_LLVM})
32+
if (${TVM_LLVM_VERSION} LESS 60)
33+
message(FATAL_ERROR "LLVM version 6.0 or greater is required.")
34+
endif()
3235
include_directories(SYSTEM ${LLVM_INCLUDE_DIRS})
3336
add_definitions(${LLVM_DEFINITIONS})
3437
message(STATUS "Build with LLVM " ${LLVM_PACKAGE_VERSION})

python/tvm/target/codegen.py

Lines changed: 79 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
"""Code generation related functions."""
1818
from . import _ffi_api
1919
from .target import Target
20+
from ..ir.container import Array
2021

2122

2223
def build_module(mod, target):
@@ -39,6 +40,30 @@ def build_module(mod, target):
3940
return _ffi_api.Build(mod, target)
4041

4142

43+
def target_has_features(cpu_features, target=None):
44+
"""Check CPU features for the target's `-mtriple` and `-mcpu` and `-mattr`.
45+
46+
Parameters
47+
----------
48+
target : Target
49+
The TVM target.
50+
cpu_features : str or Array
51+
CPU Feature(s) to check.
52+
53+
Returns
54+
-------
55+
has_features : bool
56+
True if target has the feature(s).
57+
"""
58+
assert isinstance(target, Target) or target is None
59+
assert isinstance(cpu_features, (Array, list, tuple, str))
60+
has_feats = True
61+
cpu_features = [cpu_features] if isinstance(cpu_features, str) else cpu_features
62+
for feat in cpu_features:
63+
has_feats &= _ffi_api.target_has_feature(feat, target)
64+
return has_feats
65+
66+
4267
def llvm_lookup_intrinsic_id(name):
4368
"""Lookup LLVM intrinsic id by name.
4469
@@ -71,36 +96,76 @@ def llvm_get_intrinsic_name(intrin_id: int) -> str:
7196
return _ffi_api.llvm_get_intrinsic_name(intrin_id)
7297

7398

74-
def llvm_x86_get_archlist(only64bit=False):
75-
"""Get X86 CPU name list.
99+
def llvm_get_targets():
100+
"""Get LLVM target list.
101+
102+
Parameters
103+
----------
104+
105+
Returns
106+
-------
107+
llvm_targets : list[str]
108+
List of available LLVM targets.
109+
"""
110+
return _ffi_api.llvm_get_targets()
111+
112+
113+
def llvm_get_cpu_archlist(target=None):
114+
"""Get CPU architectures for the target's `-mtriple`.
115+
116+
Parameters
117+
----------
118+
target : Target
119+
The TVM target.
120+
121+
Returns
122+
-------
123+
cpu_archlist : list[str]
124+
List of available CPU architectures.
125+
"""
126+
assert isinstance(target, Target) or target is None
127+
return _ffi_api.llvm_get_cpu_archlist(target)
128+
129+
130+
def llvm_get_cpu_features(target=None):
131+
"""Get CPU features for the target's `-mtriple` and `-mcpu` and considering `-mattr`.
76132
77133
Parameters
78134
----------
79-
only64bit : bool
80-
Filter 64bit architectures.
135+
target : Target
136+
The TVM target.
81137
82138
Returns
83139
-------
84-
features : list[str]
85-
String list of X86 architectures.
140+
cpu_features : list[str]
141+
List of available CPU features.
86142
"""
87-
return _ffi_api.llvm_x86_get_archlist(only64bit)
143+
assert isinstance(target, Target) or target is None
144+
return _ffi_api.llvm_get_cpu_features(target)
88145

89146

90-
def llvm_x86_get_features(cpu_name):
91-
"""Get X86 CPU features.
147+
def llvm_cpu_has_features(cpu_features, target=None):
148+
"""Check CPU features for the target's `-mtriple` and `-mcpu` and considering `-mattr`.
92149
93150
Parameters
94151
----------
95-
cpu_name : string
96-
X86 CPU name (e.g. "skylake").
152+
target : Target
153+
The TVM target.
154+
cpu_features : str or Array
155+
CPU Feature(s) to check.
97156
98157
Returns
99158
-------
100-
features : list[str]
101-
String list of X86 CPU features.
159+
has_features : bool
160+
True if target CPU has the feature(s).
102161
"""
103-
return _ffi_api.llvm_x86_get_features(cpu_name)
162+
assert isinstance(target, Target) or target is None
163+
assert isinstance(cpu_features, (Array, list, tuple, str))
164+
has_feats = True
165+
cpu_features = [cpu_features] if isinstance(cpu_features, str) else cpu_features
166+
for feat in cpu_features:
167+
has_feats &= _ffi_api.llvm_cpu_has_feature(feat, target)
168+
return has_feats
104169

105170

106171
def llvm_version_major(allow_none=False):

python/tvm/target/x86.py

Lines changed: 1 addition & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -16,30 +16,7 @@
1616
# under the License.
1717
"""Common x86 related utilities"""
1818
from .._ffi import register_func
19-
from . import _ffi_api
20-
from ..ir.container import Array
21-
22-
23-
@register_func("tvm.target.x86.target_has_features")
24-
def target_has_features(features, target=None):
25-
"""Check X86 CPU features.
26-
Parameters
27-
----------
28-
features : str or Array
29-
Feature(s) to check.
30-
target : Target
31-
Optional TVM target, default `None` use the global context target.
32-
Returns
33-
-------
34-
has_feats : bool
35-
True if feature(s) are in the target arch.
36-
"""
37-
has_feats = True
38-
assert isinstance(features, (Array, str))
39-
features = [features] if isinstance(features, str) else features
40-
for feat in features:
41-
has_feats &= _ffi_api.llvm_x86_has_feature(feat, target)
42-
return has_feats
19+
from .codegen import target_has_features
4320

4421

4522
@register_func("tvm.topi.x86.utils.get_simd_32bit_lanes")
@@ -53,9 +30,6 @@ def get_simd_32bit_lanes():
5330
The optimal vector length of CPU from the global context target.
5431
"""
5532
vec_len = 4
56-
# avx512f: llvm.x86.avx512.addpd.w.512 (LLVM auto, added)
57-
# avx512bw: llvm.x86.avx512.pmaddubs.w.512" (TVM required)
58-
# + llvm.x86.avx512.pmaddw.d.512"
5933
if target_has_features(["avx512bw", "avx512f"]):
6034
vec_len = 16
6135
elif target_has_features("avx2"):

python/tvm/topi/x86/batch_matmul.py

Lines changed: 1 addition & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
from tvm import autotvm, te
2222
from tvm.autotvm.task.space import SplitEntity
2323
from tvm.contrib import cblas, mkl
24-
from tvm.target.x86 import target_has_features
24+
from tvm.target.codegen import target_has_features
2525

2626
from .. import generic, nn
2727
from ..transform import layout_transform
@@ -38,9 +38,6 @@ def batch_matmul_int8_compute(cfg, x, y, *_):
3838
packed_y = layout_transform(y, "BNK", packed_y_layout)
3939
_, n_o, _, n_i, _ = packed_y.shape
4040
ak = te.reduce_axis((0, k), name="k")
41-
# avx512f: llvm.x86.avx512.addpd.w.512 (LLVM auto, added)
42-
# avx512bw: llvm.x86.avx512.pmaddubs.w.512" (TVM required)
43-
# + llvm.x86.avx512.pmaddw.d.512"
4441
if target_has_features(["avx512bw", "avx512f"]):
4542
attrs_info = {"schedule_rule": "batch_matmul_int8"}
4643
else:
@@ -241,9 +238,6 @@ def _callback(op):
241238
layout_trans = op.input_tensors[1]
242239
if target_has_features("amx-int8"):
243240
batch_matmul_amx_schedule(cfg, s, op.output(0), outs[0], layout_trans)
244-
# avx512f: llvm.x86.avx512.addpd.w.512 (LLVM auto, added)
245-
# avx512bw: llvm.x86.avx512.pmaddubs.w.512" (TVM required)
246-
# + llvm.x86.avx512.pmaddw.d.512"
247241
elif target_has_features(["avx512bw", "avx512f"]):
248242
batch_matmul_int8_schedule(cfg, s, op.output(0), outs[0], layout_trans)
249243

python/tvm/topi/x86/dense.py

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,8 @@
2323
from tvm import autotvm, te
2424
from tvm.autotvm.task.space import SplitEntity
2525
from tvm.contrib import cblas, dnnl, mkl
26-
from tvm.target.x86 import get_simd_32bit_lanes, target_has_features
26+
from tvm.target.x86 import get_simd_32bit_lanes
27+
from tvm.target.codegen import target_has_features
2728

2829
from .. import generic, tag
2930
from ..utils import get_const_tuple, traverse_inline
@@ -303,9 +304,6 @@ def _callback(op):
303304
if "dense_int8" in op.tag:
304305
if target_has_features("amx-int8"):
305306
dense_amx_int8_schedule(cfg, s, op.output(0), outs[0])
306-
# avx512f: llvm.x86.avx512.addpd.w.512 (LLVM auto, added)
307-
# avx512bw: llvm.x86.avx512.pmaddubs.w.512" (TVM required)
308-
# + llvm.x86.avx512.pmaddw.d.512"
309307
elif target_has_features(["avx512bw", "avx512f"]):
310308
dense_int8_schedule(cfg, s, op.output(0), outs[0])
311309

@@ -318,9 +316,6 @@ def dense_int8_compute(cfg, X, packed_w, bias=None):
318316
m, k = X.shape
319317
n_o, _, n_i, _ = packed_w.shape
320318
ak = te.reduce_axis((0, k), name="k")
321-
# avx512f: llvm.x86.avx512.addpd.w.512 (LLVM auto, added)
322-
# avx512bw: llvm.x86.avx512.pmaddubs.w.512" (TVM required)
323-
# + llvm.x86.avx512.pmaddw.d.512"
324319
if target_has_features(["avx512bw", "avx512f"]):
325320
target_attr = {"schedule_rule": "meta_schedule.x86.dense_int8"}
326321
else:

python/tvm/topi/x86/dense_alter_op.py

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919

2020
import tvm
2121
from tvm import autotvm, relay, te
22-
from tvm.target.x86 import target_has_features
22+
from tvm.target.codegen import target_has_features
2323

2424
from .. import nn
2525
from ..nn import dense_alter_layout
@@ -28,9 +28,6 @@
2828

2929

3030
def check_int8_applicable(x, y, allow_padding=False):
31-
# avx512f: llvm.x86.avx512.addpd.w.512 (LLVM auto, added)
32-
# avx512bw: llvm.x86.avx512.pmaddubs.w.512" (TVM required)
33-
# + llvm.x86.avx512.pmaddw.d.512"
3431
simd_avai = target_has_features(["avx512bw", "avx512f"])
3532
simd_avai |= target_has_features("amx-int8")
3633
# TODO(vvchernov): may be also target_has_features("avx2") or lower?

src/meta_schedule/space_generator/space_generator.cc

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -24,20 +24,17 @@ namespace meta_schedule {
2424

2525
String GetRuleKindFromTarget(const Target& target) {
2626
if (target->kind->name == "llvm") {
27-
static const PackedFunc* llvm_x86_has_feature_fn_ptr =
28-
runtime::Registry::Get("target.llvm_x86_has_feature");
29-
ICHECK(llvm_x86_has_feature_fn_ptr != nullptr)
30-
<< "The `target.llvm_x86_has_feature` func is not in tvm registry.";
31-
bool have_avx512vnni = (*llvm_x86_has_feature_fn_ptr)("avx512vnni", target);
32-
bool have_avxvnni = (*llvm_x86_has_feature_fn_ptr)("avxvnni", target);
27+
static const PackedFunc* target_has_feature_fn_ptr =
28+
runtime::Registry::Get("target.target_has_feature");
29+
ICHECK(target_has_feature_fn_ptr != nullptr)
30+
<< "The `target.target_has_feature` func is not in tvm registry.";
31+
bool have_avx512vnni = (*target_has_feature_fn_ptr)("avx512vnni", target);
32+
bool have_avxvnni = (*target_has_feature_fn_ptr)("avxvnni", target);
3333
if (have_avx512vnni || have_avxvnni) {
3434
return "vnni";
3535
} else {
36-
// avx512f: llvm.x86.avx512.addpd.w.512 (LLVM auto, added)
37-
// avx512bw: llvm.x86.avx512.pmaddubs.w.512" (TVM required)
38-
// + llvm.x86.avx512.pmaddw.d.512"
39-
bool have_avx512f = (*llvm_x86_has_feature_fn_ptr)("avx512f", target);
40-
bool have_avx512bw = (*llvm_x86_has_feature_fn_ptr)("avx512bw", target);
36+
bool have_avx512f = (*target_has_feature_fn_ptr)("avx512f", target);
37+
bool have_avx512bw = (*target_has_feature_fn_ptr)("avx512bw", target);
4138
if (have_avx512bw && have_avx512f) {
4239
return "avx512";
4340
}

src/relay/qnn/op/requantize.cc

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -121,9 +121,9 @@ InferCorrectLayoutOutput RequantizeInferCorrectLayout(const Attrs& attrs,
121121
}
122122

123123
bool has_current_target_sse41_support() {
124-
auto llvm_x86_has_feature_fn_ptr = tvm::runtime::Registry::Get("target.llvm_x86_has_feature");
125-
ICHECK(llvm_x86_has_feature_fn_ptr) << "Function target.llvm_x86_has_feature not found";
126-
return (*llvm_x86_has_feature_fn_ptr)("sse4.1", Target::Current(true));
124+
auto target_has_feature_fn_ptr = tvm::runtime::Registry::Get("target.target_has_feature");
125+
ICHECK(target_has_feature_fn_ptr) << "Function target.target_has_feature not found";
126+
return (*target_has_feature_fn_ptr)("sse4.1", Target::Current(true));
127127
}
128128

129129
/*

src/relay/qnn/op/requantize_config.h

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -61,10 +61,10 @@ class RequantizeConfigNode : public Object {
6161
// For the x86 architecture, the float32 computation is expected to give significant speedup,
6262
// with little loss in the accuracy of the requantize operation.
6363
auto target = Target::Current(true);
64-
auto llvm_x86_has_feature_fn_ptr = tvm::runtime::Registry::Get("target.llvm_x86_has_feature");
65-
ICHECK(llvm_x86_has_feature_fn_ptr) << "Function target.llvm_x86_has_feature not found";
64+
auto target_has_feature_fn_ptr = tvm::runtime::Registry::Get("target.target_has_feature");
65+
ICHECK(target_has_feature_fn_ptr) << "Function target.target_has_feature not found";
6666
if (target.defined() && target->kind->name == "llvm") {
67-
if ((*llvm_x86_has_feature_fn_ptr)("sse4.1", target)) {
67+
if ((*target_has_feature_fn_ptr)("sse4.1", target)) {
6868
return "float32";
6969
}
7070
}

src/target/llvm/codegen_x86_64.cc

Lines changed: 2 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -29,9 +29,7 @@
2929
#if TVM_LLVM_VERSION >= 100
3030
#include <llvm/IR/IntrinsicsX86.h>
3131
#endif
32-
#include <llvm/MC/MCSubtargetInfo.h>
3332
#include <llvm/Support/Casting.h>
34-
#include <llvm/Target/TargetMachine.h>
3533
#include <tvm/runtime/registry.h>
3634

3735
#include <string>
@@ -43,38 +41,6 @@
4341
namespace tvm {
4442
namespace codegen {
4543

46-
namespace {
47-
bool TargetHasFeature(const llvm::TargetMachine& tm, const std::string& feature) {
48-
// MCSubTargetInfo::checkFeatures was added in LLVM 6.0
49-
#if TVM_LLVM_VERSION >= 60
50-
const auto* MCInfo = tm.getMCSubtargetInfo();
51-
return MCInfo->checkFeatures(std::string("+") + feature);
52-
#else
53-
return false;
54-
// TODO(tulloch) - enable this block, need to figure out how to reimplement
55-
// this given visibility constraints, similar to
56-
// https://github.com/rust-lang/rust/pull/31709
57-
58-
// Copied from
59-
// https://github.com/llvm-mirror/llvm/blob/5136df4/lib/MC/MCSubtargetInfo.cpp#L78-L88.
60-
61-
// auto checkFeatures = [&](const std::string FS) {
62-
// llvm::SubtargetFeatures T(FS);
63-
// llvm::FeatureBitset Set, All;
64-
// for (std::string F : T.getFeatures()) {
65-
// llvm::SubtargetFeatures::ApplyFeatureFlag(Set, F, MCInfo->ProcFeatures);
66-
// if (F[0] == '-') {
67-
// F[0] = '+';
68-
// }
69-
// llvm::SubtargetFeatures::ApplyFeatureFlag(All, F, MCInfo->ProcFeatures);
70-
// }
71-
// return (MCInfo->getFeatureBits() & All) == Set;
72-
// };
73-
// return checkFeatures(MCInfo, std::string("+") + feature);
74-
#endif
75-
}
76-
} // namespace
77-
7844
class CodeGenX86_64 final : public CodeGenCPU {
7945
public:
8046
llvm::Value* VisitExpr_(const CastNode* op) override;
@@ -92,9 +58,8 @@ llvm::Value* CodeGenX86_64::VisitExpr_(const CastNode* op) {
9258
const auto to = op->dtype;
9359
if (from.is_float() && to.is_float() && from.bits() == 16 && to.bits() == 32) {
9460
ICHECK_EQ(from.lanes(), to.lanes());
95-
llvm::TargetMachine* tm = llvm_target_->GetOrCreateTargetMachine();
9661

97-
const auto has_avx512 = TargetHasFeature(*tm, "avx512f");
62+
const auto has_avx512 = llvm_target_->TargetHasCPUFeature("avx512f");
9863

9964
if (from.lanes() >= 16 && has_avx512) {
10065
return CallVectorIntrin(
@@ -111,7 +76,7 @@ llvm::Value* CodeGenX86_64::VisitExpr_(const CastNode* op) {
11176

11277
#if TVM_LLVM_VERSION <= 100
11378
// The intrinsic x86_vcvtph2ps_256 was removed in LLVM 11.
114-
const auto has_f16c = TargetHasFeature(*tm, "f16c");
79+
const auto has_f16c = llvm_target_->TargetHasCPUFeature("f16c");
11580

11681
if (from.lanes() >= 8 && has_f16c) {
11782
return CallVectorIntrin(llvm::Intrinsic::x86_vcvtph2ps_256, 8,

0 commit comments

Comments
 (0)