gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[BACKEND][CODEGEN] vectorization bugfix (#502)

Browse Source
This commit is contained in:
Philippe Tillet
2022-04-23 13:18:33 -07:00
committed by GitHub
parent 0cc3b1129b
commit bda209002e
4 changed files with 51 additions and 101 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
include/triton/codegen/analysis/align.h
View File

@@ -12,6 +12,7 @@ namespace ir {
class phi_node;
class splat_inst;
class cast_inst;
class cmp_inst;
class reshape_inst;
class broadcast_inst;
class binary_operator;
@@ -35,6 +36,7 @@ private:
std::vector<cst_info> populate_is_constant_reshape(ir::reshape_inst* x);
std::vector<cst_info> populate_is_constant_broadcast(ir::broadcast_inst* x);
std::vector<cst_info> populate_is_constant_binop(ir::binary_operator* x);
std::vector<cst_info> populate_is_constant_cmp(ir::cmp_inst* x);
std::vector<cst_info> populate_is_constant_gep(ir::getelementptr_inst* x);
std::vector<cst_info> populate_is_constant_default(ir::value* v);
std::vector<cst_info> populate_is_constant(ir::value *v);
@@ -65,6 +67,7 @@ public:
void run(ir::module &mod);
unsigned get(ir::value* v, unsigned ax) const;
std::vector<unsigned> contiguous(ir::value* v) const;
std::vector<cst_info> get_cst_info(ir::value* v) const;
private:
std::map<ir::value*, std::vector<cst_info>> is_constant_;

41
lib/codegen/analysis/align.cc
View File

@@ -129,6 +129,33 @@ std::vector<align::cst_info> align::populate_is_constant_broadcast(ir::broadcast
return add_to_cache(x, result, is_constant_);
}
std::vector<align::cst_info> align::populate_is_constant_cmp(ir::cmp_inst* x) {
auto x_shapes = get_shapes(x);
std::vector<cst_info> result;
ir::value* lhs_op = x->get_operand(0);
ir::value* rhs_op = x->get_operand(1);
auto lhs = populate_is_constant(lhs_op);
auto rhs = populate_is_constant(rhs_op);
auto lhs_max_contiguous = populate_max_contiguous(lhs_op);
auto rhs_max_contiguous = populate_max_contiguous(rhs_op);
auto lhs_multiple_of = populate_starting_multiple(lhs_op);
auto rhs_multiple_of = populate_starting_multiple(rhs_op);
for(size_t d = 0; d < x_shapes.size(); d++) {
cst_info ax = {1, 0};
// if lhs (resp. rhs) is a range of M value starting at a multiple of N
// and rhs (resp. lhs) is made of M constants that are multiples of N
// then comparisons have M constants
int min_multiple = std::min(lhs_multiple_of[d], rhs_multiple_of[d]);
if(rhs[d].num_cst % lhs_max_contiguous[d] == 0)
ax = {std::min<int>(min_multiple, lhs_max_contiguous[d]), 0};
else if(lhs[d].num_cst % rhs_max_contiguous[d] == 0)
ax = {std::min<int>(min_multiple, rhs_max_contiguous[d]), 0};
result.push_back(ax);
}
return add_to_cache(x, result, is_constant_);
}
std::vector<align::cst_info> align::populate_is_constant_binop(ir::binary_operator* x) {
auto x_shapes = get_shapes(x);
std::vector<cst_info> result;
@@ -136,12 +163,15 @@ std::vector<align::cst_info> align::populate_is_constant_binop(ir::binary_operat
ir::value* rhs_op = x->get_operand(1);
auto lhs = populate_is_constant(lhs_op);
auto rhs = populate_is_constant(rhs_op);
auto max_contiguous = populate_max_contiguous(lhs_op);
auto lhs_max_contiguous = populate_max_contiguous(lhs_op);
auto rhs_max_contiguous = populate_max_contiguous(rhs_op);
auto lhs_multiple_of = populate_starting_multiple(lhs_op);
auto rhs_multiple_of = populate_starting_multiple(rhs_op);
for(size_t d = 0; d < x_shapes.size(); d++) {
cst_info ax;
if(lhs[d].num_cst==0 && rhs[d].value && x->is_int_div()){
// todo might not be entirely true
unsigned num_constants = gcd(max_contiguous[d], rhs[d].value);
unsigned num_constants = gcd(lhs_max_contiguous[d], rhs[d].value);
ax = {num_constants, 0};
}
else
@@ -184,6 +214,8 @@ std::vector<align::cst_info> align::populate_is_constant(ir::value *v) {
return populate_is_constant_broadcast(x);
if(auto *x = dynamic_cast<ir::binary_operator*>(v))
return populate_is_constant_binop(x);
if(auto *x = dynamic_cast<ir::cmp_inst*>(v))
return populate_is_constant_cmp(x);
if(auto *x = dynamic_cast<ir::getelementptr_inst*>(v))
return populate_is_constant_gep(x);
return populate_is_constant_default(v);
@@ -511,12 +543,15 @@ std::vector<unsigned> align::contiguous(ir::value* v) const {
return max_contiguous_.at(v);
}
std::vector<align::cst_info> align::get_cst_info(ir::value* v) const {
return is_constant_.at(v);
}
void align::populate(ir::value *v) {
populate_is_constant(v);
populate_starting_multiple(v);
populate_max_contiguous(v);
}
void align::run(ir::module &mod) {

10
lib/codegen/selection/generator.cc
View File

@@ -744,6 +744,11 @@ void generator::visit_load_inst(ir::load_inst* x){
if(op->get_type()->is_block_ty()){
auto ord = ords_.at(op);
size_t aln = alignment_->get(op, ord[0]);
if(mx){
size_t max_eq = alignment_->get_cst_info(mx->get_mask_operand())[ord[0]].num_cst;
max_eq = std::max<size_t>(max_eq, 1);
aln = std::min(aln, max_eq);
}
auto layout = layouts_->get(x)->to_scanline();
if(layout){
size_t nts = layout->nts(ord[0]);
@@ -912,6 +917,11 @@ void generator::visit_store_inst(ir::store_inst * x){
auto ord = ords_.at(x->get_pointer_operand());
size_t aln = alignment_->get(ptr_op, ord[0]);
size_t nts = axes_.at(a_axes_->get(x->get_pointer_operand(), ord[0])).contiguous;
if(mx){
size_t max_eq = alignment_->get_cst_info(mx->get_mask_operand())[ord[0]].num_cst;
max_eq = std::max<size_t>(max_eq, 1);
aln = std::min(aln, max_eq);
}
vec = std::min(nts, aln);
}
auto idxs = idxs_.at(val_op);

98
python/test/unit/runtime/test_comm.py
View File

@@ -1,98 +0,0 @@
import subprocess
import numpy as np
import pytest
import torch
import triton
import triton.language as tl
def get_p2p_matrix():
try:
stdout = subprocess.check_output(["nvidia-smi", "topo", "-p2p", "n"]).decode("ascii")
except subprocess.CalledProcessError:
return pytest.skip("No multi-GPU topology", allow_module_level=True)
lines = stdout.split("Legend")[0].split('\n')[1:]
matrix = np.array([line.split('\t')[1:-1] for line in lines][:-2])
if matrix.size <= 1:
return pytest.skip("No multi-GPU topology", allow_module_level=True)
else:
return matrix
def get_p2p_devices():
matrix = get_p2p_matrix()
idx = np.where(matrix == "OK")
return [f"cuda:{idx[0][0]}", f"cuda:{idx[1][0]}"] if len(idx[0]) > 0 else []
def get_non_p2p_devices():
matrix = get_p2p_matrix()
idx = np.where(matrix == "NS")
return [f"cuda:{idx[0][0]}", f"cuda:{idx[1][0]}"] if len(idx[0]) > 0 else []
p2p_devices = get_p2p_devices()
non_p2p_devices = get_non_p2p_devices()
@triton.jit
def _copy(from_ptr, to_ptr, N, **meta):
pid = tl.program_id(0)
offsets = pid * meta['BLOCK'] + tl.arange(0, meta['BLOCK'])
values = tl.load(from_ptr + offsets, mask=offsets < N)
tl.store(to_ptr + offsets, values, mask=offsets < N)
@pytest.mark.skipif(not p2p_devices, reason="No pair of device with P2P support")
@pytest.mark.parametrize("device_kernel, device_from, device_to, stream_from, stream_to",
[(device_kernel, device_from, device_to, stream_from, stream_to)
for device_kernel in p2p_devices
for device_from in p2p_devices
for device_to in p2p_devices
for stream_from in ['default', 'custom']
for stream_to in ['default', 'custom']
])
def test_p2p(device_kernel, device_from, device_to, stream_from, stream_to):
if device_to == device_from:
return pytest.skip()
torch.cuda.set_device(device_kernel)
N = 512
grid = lambda meta: (triton.cdiv(N, meta['BLOCK']),)
with torch.cuda.stream(None if stream_from == 'default' else torch.cuda.Stream(device_from)):
x_from = torch.randn(N, dtype=torch.float32, device=device_from)
with torch.cuda.stream(None if stream_to == 'default' else torch.cuda.Stream(device_to)):
x_to = torch.empty(N, dtype=torch.float32, device=device_to)
_copy[grid](x_from, x_to, N, BLOCK=1024)
assert torch.allclose(x_from, x_to.to(device_from))
@pytest.mark.skipif(not non_p2p_devices, reason="No pair of device with no P2P support")
@pytest.mark.parametrize("device_kernel, device_from, device_to, stream_from, stream_to",
[(device_kernel, device_from, device_to, stream_from, stream_to)
for device_kernel in non_p2p_devices
for device_from in non_p2p_devices
for device_to in non_p2p_devices
for stream_from in ['default', 'custom']
for stream_to in ['default', 'custom']
])
def test_non_p2p(device_kernel, device_from, device_to, stream_from, stream_to):
if device_to == device_from:
return pytest.skip()
with pytest.raises(RuntimeError):
torch.cuda.set_device(device_kernel)
N = 512
grid = lambda meta: (triton.cdiv(N, meta['BLOCK']),)
with torch.cuda.stream(None if stream_from == 'default' else torch.cuda.Stream(device_from)):
x_from = torch.randn(N, dtype=torch.float32, device=device_from)
with torch.cuda.stream(None if stream_to == 'default' else torch.cuda.Stream(device_to)):
x_to = torch.empty(N, dtype=torch.float32, device=device_to)
_copy[grid](x_from, x_to, N, BLOCK=1024)