From c0582ee326189dca751ae6509055fe0a5bba44f8 Mon Sep 17 00:00:00 2001 From: Yichen Yan Date: Wed, 13 Mar 2024 11:10:03 +0800 Subject: [PATCH 1/5] [CI] Run examples in CI --- .github/workflows/integration-tests.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml index 58ae52c132da..7afbce55f309 100644 --- a/.github/workflows/integration-tests.yml +++ b/.github/workflows/integration-tests.yml @@ -108,6 +108,9 @@ jobs: TRITON_DISABLE_LINE_INFO=0 python3 -m pytest -vvv language/test_line_info.py # Run hopper/test_flashattention.py separately to avoid out of gpu memory python3 -m pytest -vs hopper/test_flashattention.py + # Test example + cd ../../examples + find . -type f -name '*.py' -exec echo {} \; -exec python3 {} \; - name: Clear cache run: | From 47d29395948852273c1a835c07fc742981388733 Mon Sep 17 00:00:00 2001 From: Yichen Yan Date: Wed, 13 Mar 2024 11:29:10 +0800 Subject: [PATCH 2/5] Merge `get_program_id` and `get_num_programs`'s def --- include/triton/Dialect/Triton/IR/TritonOps.td | 24 ++++++------------- 1 file changed, 7 insertions(+), 17 deletions(-) diff --git a/include/triton/Dialect/Triton/IR/TritonOps.td b/include/triton/Dialect/Triton/IR/TritonOps.td index 75200dde8e8a..6250cb58ad9c 100644 --- a/include/triton/Dialect/Triton/IR/TritonOps.td +++ b/include/triton/Dialect/Triton/IR/TritonOps.td @@ -528,21 +528,8 @@ def TT_TransOp : TT_Op<"trans", [Pure, // // SPMD Ops // -def TT_GetProgramIdOp : TT_Op<"get_program_id", [Pure]> { - let arguments = (ins TT_ProgramDim:$axis); - - let results = (outs I32:$result); - - let assemblyFormat = "$axis attr-dict `:` type($result)"; - let extraClassDeclaration = [{ - int32_t getAxisAsInt() { - return static_cast(getAxis()); - } - }]; -} - -def TT_GetNumProgramsOp : TT_Op<"get_num_programs", [Pure]> { +class program_id_inst : TT_Op { let arguments = (ins TT_ProgramDim:$axis); let results = (outs I32:$result); @@ -550,12 +537,15 @@ def TT_GetNumProgramsOp : TT_Op<"get_num_programs", [Pure]> { let assemblyFormat = "$axis attr-dict `:` type($result)"; let extraClassDeclaration = [{ - int32_t getAxisAsInt() { - return static_cast(getAxis()); - } + int32_t getAxisAsInt() { + return static_cast(getAxis()); + } }]; } +def TT_GetProgramIdOp : program_id_inst<"get_program_id">; +def TT_GetNumProgramsOp : program_id_inst<"get_num_programs">; + // // Dot Op // From 168f706bab2e511b1af5a360ffcfb2882f1da26a Mon Sep 17 00:00:00 2001 From: Yichen Yan Date: Wed, 13 Mar 2024 20:34:24 +0800 Subject: [PATCH 3/5] Update based on review --- .github/workflows/integration-tests.yml | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml index 7afbce55f309..ba52359e5d48 100644 --- a/.github/workflows/integration-tests.yml +++ b/.github/workflows/integration-tests.yml @@ -108,8 +108,14 @@ jobs: TRITON_DISABLE_LINE_INFO=0 python3 -m pytest -vvv language/test_line_info.py # Run hopper/test_flashattention.py separately to avoid out of gpu memory python3 -m pytest -vs hopper/test_flashattention.py + + - name: Run python examples on CUDA + # Execute on H100 due to its ~30s speed advantage over A100; + # example duration matches this gain and is device-agnostic. + if: ${{(matrix.runner[0] == 'self-hosted') && (matrix.runner[1] == 'H100')}} + run: | + cd python/examples # Test example - cd ../../examples find . -type f -name '*.py' -exec echo {} \; -exec python3 {} \; - name: Clear cache From 84feb767e764d12e631d7ff6ecb823d1d5e83496 Mon Sep 17 00:00:00 2001 From: Yichen Yan Date: Thu, 14 Mar 2024 10:22:27 +0800 Subject: [PATCH 4/5] Revert "Merge `get_program_id` and `get_num_programs`'s def" This reverts commit 47d29395948852273c1a835c07fc742981388733. --- include/triton/Dialect/Triton/IR/TritonOps.td | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/include/triton/Dialect/Triton/IR/TritonOps.td b/include/triton/Dialect/Triton/IR/TritonOps.td index 6250cb58ad9c..75200dde8e8a 100644 --- a/include/triton/Dialect/Triton/IR/TritonOps.td +++ b/include/triton/Dialect/Triton/IR/TritonOps.td @@ -528,8 +528,7 @@ def TT_TransOp : TT_Op<"trans", [Pure, // // SPMD Ops // - -class program_id_inst : TT_Op { +def TT_GetProgramIdOp : TT_Op<"get_program_id", [Pure]> { let arguments = (ins TT_ProgramDim:$axis); let results = (outs I32:$result); @@ -537,14 +536,25 @@ class program_id_inst : TT_Op { let assemblyFormat = "$axis attr-dict `:` type($result)"; let extraClassDeclaration = [{ - int32_t getAxisAsInt() { - return static_cast(getAxis()); - } + int32_t getAxisAsInt() { + return static_cast(getAxis()); + } }]; } -def TT_GetProgramIdOp : program_id_inst<"get_program_id">; -def TT_GetNumProgramsOp : program_id_inst<"get_num_programs">; +def TT_GetNumProgramsOp : TT_Op<"get_num_programs", [Pure]> { + let arguments = (ins TT_ProgramDim:$axis); + + let results = (outs I32:$result); + + let assemblyFormat = "$axis attr-dict `:` type($result)"; + + let extraClassDeclaration = [{ + int32_t getAxisAsInt() { + return static_cast(getAxis()); + } + }]; +} // // Dot Op From 80d6e9c71bca813b38d4644bee9751a0f99ea0b5 Mon Sep 17 00:00:00 2001 From: Yichen Yan Date: Thu, 14 Mar 2024 10:36:03 +0800 Subject: [PATCH 5/5] Resolve comments Co-authored-by: Justin Lebar --- .github/workflows/integration-tests.yml | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml index ba52359e5d48..5daaff08a3c4 100644 --- a/.github/workflows/integration-tests.yml +++ b/.github/workflows/integration-tests.yml @@ -110,12 +110,11 @@ jobs: python3 -m pytest -vs hopper/test_flashattention.py - name: Run python examples on CUDA - # Execute on H100 due to its ~30s speed advantage over A100; - # example duration matches this gain and is device-agnostic. + # Use H100 as it's ~2min faster than A100; + # additional example takes 30s and doesn't depend on the device. if: ${{(matrix.runner[0] == 'self-hosted') && (matrix.runner[1] == 'H100')}} run: | cd python/examples - # Test example find . -type f -name '*.py' -exec echo {} \; -exec python3 {} \; - name: Clear cache