Skip to content

Commit f5cbb4e

Browse files
factnnfactnn
committed
merge: upstream/master
1 parent fa18f73 commit f5cbb4e

29 files changed

Lines changed: 984 additions & 186 deletions

‎.github/workflows/command.yaml‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
name: "ondemand test"
2+
3+
on:
4+
issue_comment:
5+
types: [created]
6+
7+
permissions:
8+
pull-requests: write
9+
checks: read
10+
contents: read
11+
12+
jobs:
13+
process-command:
14+
runs-on: ubuntu-latest
15+
outputs:
16+
op: ${{ steps.parse-command.outputs.op}}
17+
runner: ${{ steps.parse-command.outputs.runner }}
18+
continue: ${{ steps.check.outputs.continue }}
19+
steps:
20+
- id: check
21+
uses: github.qkg1.topmand@v2.0.3
22+
with:
23+
command: "/test"
24+
allowed_contexts: pull_request
25+
permissions: "write,admin"
26+
27+
- id: parse-command
28+
if: ${{ steps.check.outputs.continue == 'true' }}
29+
run: |
30+
# params is op:runner
31+
params=${{ steps.check.outputs.params }}
32+
op=$(echo $params | cut -d ':' -f 1)
33+
runner=$(echo $params | cut -d ':' -f 2)
34+
echo "op=${op}" >> $GITHUB_OUTPUT
35+
echo "runner=${op}" >> $GITHUB_OUTPUT
36+
37+
test-operator:
38+
needs: process-command
39+
if: needs.process-command.outputs.continue == 'true'
40+
runs-on: ${{ needs.process-command.outputs.runner }}
41+
steps:
42+
- uses: actions/checkout@v6
43+
- name: Setup FlagGems
44+
shell: bash
45+
env:
46+
RUNNER_LABEL: ${{ needs.process-command.outputs.runner }}
47+
run: |
48+
# Validate vendor names
49+
case "${RUNNER_LABEL}" in
50+
h20)
51+
VENDOR="nvidia"
52+
;;
53+
moore)
54+
VENDOR="mthreads"
55+
;;
56+
*)
57+
VENDOR="${RUNNER_LABEL}"
58+
;;
59+
esac
60+
61+
echo "VENDOR=${VENDOR}"
62+
echo "VENDOR=${VENDOR}" >> $GITHUB_ENV
63+
64+
./setup.sh ${VENDOR}
65+
66+
- name: Check GPU availability
67+
shell: bash
68+
run: |
69+
bash "tools/gpu_check_${VENDOR}.sh"
70+
71+
- name: Run tests
72+
shell: bash
73+
run: |
74+
source .venv/bin/activate
75+
source tools/set-env.sh ${VENDOR}
76+
tools/run_tests.py --ops ${{ needs.process-command.outputs.op }}
77+
78+
- name: Upload results
79+
uses: actions/upload-artifact@v7
80+
with:
81+
name: ondemand-test
82+
path: results/

‎.github/workflows/random-test.yaml‎

Lines changed: 0 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,6 @@ on:
1616
description: "GPUs (comma separated numbers)"
1717
required: true
1818
type: string
19-
secrets:
20-
RUNNER_SSH_KEY:
21-
required: true
22-
PROXYC:
23-
required: true
24-
PROXYG:
25-
required: true
2619
permissions:
2720
contents: read
2821

@@ -55,27 +48,7 @@ jobs:
5548
if: needs.check-permission.outputs.authorized == 'true'
5649
runs-on: ${{ inputs.runner }}
5750
steps:
58-
# - name: Set global proxy
59-
# if: False
60-
# shell: bash
61-
# env:
62-
# RUNNER_LABEL: ${{ inputs.runner }}
63-
# PROXY: ${{ secrets.PROXYG }}
64-
# run: |
65-
# # Setup proxy for git checkout
66-
# case "${RUNNER_LABEL}" in
67-
# h20|metax|hygon)
68-
# ;;
69-
# *)
70-
# echo "HTTP_PROXY=${PROXY}" >> $GITHUB_ENV
71-
# echo "HTTPS_PROXY=${PROXY}" >> $GITHUB_ENV
72-
# ;;
73-
# esac
74-
7551
- uses: actions/checkout@v6
76-
# with:
77-
# ssh-key: ${{ secrets.RUNNER_SSH_KEY }}
78-
7952
- name: Setup FlagGems
8053
shell: bash
8154
env:

‎benchmark/test_scatter_reduce_two.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
import pytest
2+
import torch
3+
4+
from . import base, consts
5+
6+
7+
def _input_fn_factory(reduce):
8+
def inner(shape, dtype, device):
9+
inp = torch.randn(shape, dtype=dtype, device=device)
10+
dim = -1
11+
size_dim = shape[dim]
12+
index = torch.randint(0, size_dim, shape, dtype=torch.long, device=device)
13+
src = torch.randn(shape, dtype=dtype, device=device)
14+
yield inp, dim, index, src, {"reduce": reduce}
15+
16+
return inner
17+
18+
19+
@pytest.mark.scatter_reduce_two_
20+
def test_scatter_reduce_two_inplace_sum():
21+
bench = base.GenericBenchmark2DOnly(
22+
op_name="scatter_reduce_.sum",
23+
torch_op=torch.Tensor.scatter_reduce_,
24+
input_fn=_input_fn_factory("sum"),
25+
dtypes=consts.FLOAT_DTYPES,
26+
inplace=True,
27+
)
28+
bench.run()
29+
30+
31+
@pytest.mark.scatter_reduce_two_
32+
def test_scatter_reduce_two_inplace_amax():
33+
bench = base.GenericBenchmark2DOnly(
34+
op_name="scatter_reduce_.amax",
35+
torch_op=torch.Tensor.scatter_reduce_,
36+
input_fn=_input_fn_factory("amax"),
37+
dtypes=consts.FLOAT_DTYPES,
38+
inplace=True,
39+
)
40+
bench.run()
41+
42+
43+
@pytest.mark.scatter_reduce_two_
44+
def test_scatter_reduce_two_inplace_amin():
45+
bench = base.GenericBenchmark2DOnly(
46+
op_name="scatter_reduce_.amin",
47+
torch_op=torch.Tensor.scatter_reduce_,
48+
input_fn=_input_fn_factory("amin"),
49+
dtypes=consts.FLOAT_DTYPES,
50+
inplace=True,
51+
)
52+
bench.run()
53+
54+
55+
@pytest.mark.scatter_reduce_two_
56+
def test_scatter_reduce_two_inplace_mean():
57+
bench = base.GenericBenchmark2DOnly(
58+
op_name="scatter_reduce_.mean",
59+
torch_op=torch.Tensor.scatter_reduce_,
60+
input_fn=_input_fn_factory("mean"),
61+
dtypes=consts.FLOAT_DTYPES,
62+
inplace=True,
63+
)
64+
bench.run()

‎src/flag_gems/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -421,6 +421,7 @@ def torch_ge(v):
421421
("scatter_.reduce", scatter_),
422422
("scatter_.src", scatter_),
423423
("scatter_add_", scatter_add_),
424+
("scatter_reduce_.two", scatter_reduce_),
424425
("select_backward", select_backward),
425426
("select_scatter", select_scatter),
426427
("selu", selu),

‎src/flag_gems/ops/__init__.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -281,6 +281,7 @@
281281
from flag_gems.ops.scaled_softmax import scaled_softmax_backward, scaled_softmax_forward
282282
from flag_gems.ops.scatter import scatter, scatter_
283283
from flag_gems.ops.scatter_add_ import scatter_add_
284+
from flag_gems.ops.scatter_reduce_ import scatter_reduce_
284285
from flag_gems.ops.select_backward import select_backward
285286
from flag_gems.ops.select_scatter import select_scatter
286287
from flag_gems.ops.selu import selu
@@ -709,6 +710,7 @@
709710
"scatter",
710711
"scatter_",
711712
"scatter_add_",
713+
"scatter_reduce_",
712714
"select_backward",
713715
"select_scatter",
714716
"selu",

0 commit comments

Comments
 (0)