[triton-mlir][BACKEND] Support masked load/store (#657)

This PR does

- fix some bugs to support masked load/store,
- refine frontend, and support the `and` and `or` syntax in mask(by
extending the BoolOp in python ast.visitor), e.g. `tl.store(...,
mask=offset<n and other_conditions)`,
- add `arith.cmpI` and `arith.cmpF` op conversion in backend(required by
mask),
- add more test cases in vecadd.
This commit is contained in:
Yan Chunwei
2022-10-10 13:29:53 +08:00
committed by GitHub
parent ccc5ab6ac9
commit 555f94f9b9
9 changed files with 396 additions and 74 deletions

View File

@@ -1,6 +1,6 @@
import pytest
import torch
from torch.testing import assert_allclose
from torch.testing import assert_close
import triton
import triton.language as tl
@@ -49,4 +49,4 @@ def test_gemm_impl(SIZE_M, SIZE_N, SIZE_K, NUM_WARPS):
num_warps=NUM_WARPS)
golden = torch.matmul(a, b)
torch.set_printoptions(profile="full")
assert_allclose(c, golden, rtol=1e-3, atol=1e-3)
assert_close(c, golden, rtol=1e-3, atol=1e-3, check_dtype=False)