Skip to content

Commit 3692e08

Browse files
xzliu-optxzliu-opt
authored
Fix 831: split xor/ixor operators and register ixor (#5283)
1 parent 7eefdf4 commit 3692e08

5 files changed

Lines changed: 93 additions & 54 deletions

File tree

‎benchmark/test_ixor.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
import pytest
2+
import torch
3+
4+
from . import base, consts
5+
6+
7+
@pytest.mark.ixor
8+
def test_ixor():
9+
bench = base.BinaryPointwiseBenchmark(
10+
op_name="ixor",
11+
torch_op=torch.ops.aten.__ixor__,
12+
dtypes=consts.INT_DTYPES + consts.BOOL_DTYPES,
13+
)
14+
bench.run()

benchmark/test_xor__.py renamed to benchmark/test_xor.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,10 +4,10 @@
44
from . import base, consts
55

66

7-
@pytest.mark.xor__
8-
def test_xor__():
7+
@pytest.mark.xor
8+
def test_xor():
99
bench = base.BinaryPointwiseBenchmark(
10-
op_name="xor__",
10+
op_name="xor",
1111
torch_op=torch.bitwise_xor,
1212
dtypes=consts.INT_DTYPES + consts.BOOL_DTYPES,
1313
)

‎conf/operators.yaml‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11496,13 +11496,26 @@ ops:
1149611496
stages:
1149711497
- beta: '5.0'
1149811498
- stable: '5.3'
11499-
- id: xor__
11499+
- id: xor
1150011500
description: |
11501-
Computes the bitwise XOR of input tensors or a tensor and a scalar.
11501+
Computes the bitwise XOR of input tensors or a tensor and a scalar (non-inplace).
1150211502
for:
1150311503
- __xor__
1150411504
- __xor__.Tensor
1150511505
- __xor__.Scalar
11506+
labels:
11507+
- aten
11508+
- pointwise
11509+
- KernelGen
11510+
kind:
11511+
- Math
11512+
stages:
11513+
- alpha: '5.4'
11514+
- id: ixor
11515+
description: |
11516+
Computes the bitwise XOR of input tensors or a tensor and a scalar (inplace).
11517+
for:
11518+
- __ixor__
1150611519
- __ixor__.Tensor
1150711520
- __ixor__.Scalar
1150811521
labels:

‎tests/test_ixor.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
import random
2+
3+
import pytest
4+
import torch
5+
6+
import flag_gems
7+
8+
from . import accuracy_utils as utils
9+
10+
# __ixor__ only supports integer and boolean dtypes
11+
INT_DTYPES = [torch.int16, torch.int32, torch.int64]
12+
BOOL_TYPES = [torch.bool]
13+
14+
15+
@pytest.mark.ixor
16+
@pytest.mark.parametrize("shape", utils.POINTWISE_SHAPES)
17+
@pytest.mark.parametrize("dtype", INT_DTYPES + BOOL_TYPES)
18+
def test_ixor(shape, dtype):
19+
if dtype in BOOL_TYPES:
20+
inp1 = torch.randint(0, 2, size=shape, dtype=dtype, device=flag_gems.device)
21+
inp2 = torch.randint(0, 2, size=shape, dtype=dtype, device=flag_gems.device)
22+
else:
23+
inp1 = torch.randint(
24+
low=-0x7FFF, high=0x7FFF, size=shape, dtype=dtype, device="cpu"
25+
).to(flag_gems.device)
26+
inp2 = torch.randint(
27+
low=-0x7FFF, high=0x7FFF, size=shape, dtype=dtype, device="cpu"
28+
).to(flag_gems.device)
29+
ref_inp1 = utils.to_reference(inp1.clone())
30+
ref_inp2 = utils.to_reference(inp2)
31+
32+
ref_out = ref_inp1.__ixor__(ref_inp2)
33+
with flag_gems.use_gems():
34+
res_out = inp1.__ixor__(inp2)
35+
36+
utils.gems_assert_equal(res_out, ref_out)
37+
38+
39+
@pytest.mark.ixor
40+
@pytest.mark.parametrize("shape", utils.POINTWISE_SHAPES)
41+
@pytest.mark.parametrize("dtype", INT_DTYPES + BOOL_TYPES)
42+
def test_ixor_scalar(shape, dtype):
43+
if dtype in BOOL_TYPES:
44+
inp1 = torch.randint(0, 2, size=shape, dtype=dtype, device=flag_gems.device)
45+
inp2 = bool(random.randint(0, 2))
46+
else:
47+
inp1 = torch.randint(
48+
low=-0x7FFF, high=0x7FFF, size=shape, dtype=dtype, device="cpu"
49+
).to(flag_gems.device)
50+
inp2 = 0x00FF
51+
ref_inp1 = utils.to_reference(inp1.clone())
52+
53+
ref_out = ref_inp1.__ixor__(inp2)
54+
with flag_gems.use_gems():
55+
res_out = inp1.__ixor__(inp2)
56+
57+
utils.gems_assert_equal(res_out, ref_out)

tests/test_xor__.py renamed to tests/test_xor.py

Lines changed: 4 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -12,10 +12,10 @@
1212
BOOL_TYPES = [torch.bool]
1313

1414

15-
@pytest.mark.xor__
15+
@pytest.mark.xor
1616
@pytest.mark.parametrize("shape", utils.POINTWISE_SHAPES)
1717
@pytest.mark.parametrize("dtype", INT_DTYPES + BOOL_TYPES)
18-
def test_xor__(shape, dtype):
18+
def test_xor(shape, dtype):
1919
if dtype in BOOL_TYPES:
2020
inp1 = torch.randint(0, 2, size=shape, dtype=dtype, device="cpu").to(
2121
flag_gems.device
@@ -40,34 +40,10 @@ def test_xor__(shape, dtype):
4040
utils.gems_assert_equal(res_out, ref_out)
4141

4242

43-
@pytest.mark.xor__
43+
@pytest.mark.xor
4444
@pytest.mark.parametrize("shape", utils.POINTWISE_SHAPES)
4545
@pytest.mark.parametrize("dtype", INT_DTYPES + BOOL_TYPES)
46-
def test_xor__inplace(shape, dtype):
47-
if dtype in BOOL_TYPES:
48-
inp1 = torch.randint(0, 2, size=shape, dtype=dtype, device=flag_gems.device)
49-
inp2 = torch.randint(0, 2, size=shape, dtype=dtype, device=flag_gems.device)
50-
else:
51-
inp1 = torch.randint(
52-
low=-0x7FFF, high=0x7FFF, size=shape, dtype=dtype, device="cpu"
53-
).to(flag_gems.device)
54-
inp2 = torch.randint(
55-
low=-0x7FFF, high=0x7FFF, size=shape, dtype=dtype, device="cpu"
56-
).to(flag_gems.device)
57-
ref_inp1 = utils.to_reference(inp1.clone())
58-
ref_inp2 = utils.to_reference(inp2)
59-
60-
ref_out = ref_inp1.__ixor__(ref_inp2)
61-
with flag_gems.use_gems():
62-
res_out = inp1.__ixor__(inp2)
63-
64-
utils.gems_assert_equal(res_out, ref_out)
65-
66-
67-
@pytest.mark.xor__
68-
@pytest.mark.parametrize("shape", utils.POINTWISE_SHAPES)
69-
@pytest.mark.parametrize("dtype", INT_DTYPES + BOOL_TYPES)
70-
def test_xor__scalar(shape, dtype):
46+
def test_xor_scalar(shape, dtype):
7147
if dtype in BOOL_TYPES:
7248
inp1 = torch.randint(0, 2, size=shape, dtype=dtype, device="cpu").to(
7349
flag_gems.device
@@ -85,24 +61,3 @@ def test_xor__scalar(shape, dtype):
8561
res_out = inp1 ^ inp2
8662

8763
utils.gems_assert_equal(res_out, ref_out)
88-
89-
90-
@pytest.mark.xor__
91-
@pytest.mark.parametrize("shape", utils.POINTWISE_SHAPES)
92-
@pytest.mark.parametrize("dtype", INT_DTYPES + BOOL_TYPES)
93-
def test_xor__scalar_inplace(shape, dtype):
94-
if dtype in BOOL_TYPES:
95-
inp1 = torch.randint(0, 2, size=shape, dtype=dtype, device=flag_gems.device)
96-
inp2 = bool(random.randint(0, 2))
97-
else:
98-
inp1 = torch.randint(
99-
low=-0x7FFF, high=0x7FFF, size=shape, dtype=dtype, device="cpu"
100-
).to(flag_gems.device)
101-
inp2 = 0x00FF
102-
ref_inp1 = utils.to_reference(inp1.clone())
103-
104-
ref_out = ref_inp1.__ixor__(inp2)
105-
with flag_gems.use_gems():
106-
res_out = inp1.__ixor__(inp2)
107-
108-
utils.gems_assert_equal(res_out, ref_out)

0 commit comments

Comments
 (0)