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

fix 6/7 dot tests

Browse Source
This commit is contained in:
Michael Melesse
2022-11-01 14:18:06 +00:00
parent 4f3e2d6ed7
commit 4fb9d4904e
3 changed files with 15 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -2973,19 +2973,20 @@ void generator::visit_dot_inst(ir::dot_inst* dot) {
size_t red_axis = 1;
unsigned NK = A_shapes[red_axis];
bool is_outer = NK == 1;
#ifdef USE_ROCM
bool is_mma = false;
return visit_fmadot(dot, A, B, D, NK, c_ty, f_mul_add);
#else
bool is_mma = layouts_->get(dot)->to_mma();
#endif
if(!is_outer && is_mma && tgt_->as_nvidia() && tgt_->as_nvidia()->sm() < 80)
return visit_mma884(dot, A, B, D, NK);
if(!is_outer && is_mma && tgt_->as_nvidia() && tgt_->as_nvidia()->sm() >= 80)
return visit_mma16816(dot, A, B, D, NK); // rename it as visit_mma_v2()?
// if (dot->get_type()->get_scalar_ty()->is_fp32_ty() &&
// A->get_type()->get_scalar_ty()->is_fp32_ty())
return visit_fmadot(dot, A, B, D, NK, c_ty, f_mul_add);
// throw std::runtime_error("dot has invalid operand type");
if (dot->get_type()->get_scalar_ty()->is_fp32_ty() &&
A->get_type()->get_scalar_ty()->is_fp32_ty())
return visit_fmadot(dot, A, B, D, NK, c_ty, f_mul_add);
throw std::runtime_error("dot has invalid operand type");
#endif
}
void generator::visit_trans_inst(ir::trans_inst* trans) {

2
python/test/unit/language/test_core.py
View File

@@ -1067,7 +1067,7 @@ def test_permute(dtype_str, shape, perm, device='cuda'):
[(epilogue, allow_tf32, dtype)
for epilogue in ['none', 'trans', 'add-matrix', 'add-rows', 'add-cols', 'softmax', 'chain-dot']
for allow_tf32 in [True, False]
for dtype in ['float32', 'float16']
for dtype in ['float16']
if not (allow_tf32 and (dtype in ['float16']))])
def test_dot(epilogue, allow_tf32, dtype, device='cuda'):
if torch.version.hip is not None:

7
python/triton/language/semantic.py
View File

@@ -6,6 +6,8 @@ from . import core as tl
from triton._C.libtriton.triton import ir
import torch
# Create custom exception that prints message "hello"
class IncompatibleTypeErrorimpl(Exception):
def __init__(self, type_a, type_b):
@@ -969,6 +971,11 @@ def dot(a: tl.tensor,
trans_b: bool,
allow_tf32: bool,
builder: ir.builder) -> tl.tensor:
if torch.version.hip is not None:
a = cast(a, tl.float32, builder)
b = cast(b, tl.float32, builder)
in_a = 1 if not trans_a else 0
in_b = 1 if trans_b else 0
assert a.type.is_block() and b.type.is_block()