[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:
@@ -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
|
||||
@@ -44,4 +44,4 @@ def test_convert_layout_impl(NUM_WARPS, SIZE_M, SIZE_N):
|
||||
z = torch.empty((SIZE_N, SIZE_M), device=x.device, dtype=x.dtype)
|
||||
kernel[grid](x_ptr=x, stride_xm=x.stride(0), z_ptr=z, stride_zn=z.stride(0), SIZE_M=SIZE_M, SIZE_N=SIZE_N, num_warps=NUM_WARPS)
|
||||
golden_z = torch.t(x)
|
||||
assert_allclose(z, golden_z, rtol=1e-7, atol=1e-7)
|
||||
assert_close(z, golden_z, rtol=1e-7, atol=1e-7, check_dtype=False)
|
||||
|
Reference in New Issue
Block a user