From 901031a8b367edd5a4228865473c3babd0d9aca7 Mon Sep 17 00:00:00 2001 From: Micah Weston Date: Wed, 27 Aug 2025 15:18:45 +0000 Subject: [PATCH 1/3] Update triton to e44bd1c83c1c3e8deac7c4f02683cfb3cc395c8b --- .gitmodules | 3 -- backend/driver.py | 2 - .../ConversionPatterns.hpp | 52 +++++++++---------- .../Dialect/TPtr/IR/TPtrDialect.td | 3 -- .../StructuredToMemref/StructuredToMemref.cpp | 15 +++--- .../TritonArithToLinalg.cpp | 1 + .../TritonToLinalg/TritonToLinalg.cpp | 1 + .../UnstructuredToMemrefPass.cpp | 4 +- python/examples/conftest.py | 26 +++++++++- ...duce_unsplat.mlir => convert_unsplat.mlir} | 5 +- tools/triton-shared-opt/CMakeLists.txt | 1 + .../RegisterTritonSharedDialects.h | 6 +-- 12 files changed, 63 insertions(+), 56 deletions(-) rename test/Conversion/TritonToLinalgExperimental/{reduce_unsplat.mlir => convert_unsplat.mlir} (95%) diff --git a/.gitmodules b/.gitmodules index 711388b9a..e69de29bb 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,3 +0,0 @@ -[submodule "triton"] - path = triton - url = https://github.com/triton-lang/triton.git diff --git a/backend/driver.py b/backend/driver.py index e8053cc4f..a115eca14 100644 --- a/backend/driver.py +++ b/backend/driver.py @@ -54,8 +54,6 @@ def _ty_to_cpp(ty): "u16": "uint16_t", "u32": "uint32_t", "u64": "uint64_t", - "fp16": "float", - "bf16": "float", "fp32": "float", "f32": "float", "fp64": "double", diff --git a/include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp b/include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp index f9c102bdf..b1ebf1f51 100644 --- a/include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp +++ b/include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp @@ -360,7 +360,7 @@ struct LoadConverter : public OpConversionPattern { loc, rewriter); auto zeroMap = AffineMap::getConstantMap(0, rewriter.getContext()); auto loadOp = rewriter.create( - op.getLoc(), sMemRef, zeroMap, std::nullopt); + op.getLoc(), sMemRef, zeroMap, ValueRange{}); rewriter.replaceOp(op, loadOp.getResult()); return success(); } @@ -520,7 +520,7 @@ struct StoreConverter : public OpConversionPattern { PtrAnalysis::getScalarMemRef(op.getPtr(), ptr, loc, rewriter); auto zeroMap = AffineMap::getConstantMap(0, rewriter.getContext()); rewriter.create(loc, val, sMemRef, zeroMap, - std::nullopt); + ValueRange{}); rewriter.eraseOp(op); return success(); } @@ -649,6 +649,28 @@ struct SplatConverter : public OpConversionPattern { } }; +struct UnsplatConverter : public OpConversionPattern { + using OpConversionPattern::OpConversionPattern; + + LogicalResult + matchAndRewrite(triton::UnsplatOp op, OpAdaptor adaptor, + ConversionPatternRewriter &rewriter) const override { + auto tensorType = op.getSrc().getType(); + + // Only generate indices for non-zero rank tensors. + SmallVector indices(tensorType.getRank()); + if (indices.size() > 0) { + auto zeroIdx = + rewriter.createOrFold(op.getLoc(), 0); + llvm::fill(indices, zeroIdx); + } + + rewriter.replaceOpWithNewOp(op, adaptor.getSrc(), + indices); + return success(); + } +}; + struct BroadcastConverter : public OpConversionPattern { private: using OpConversionPattern::OpConversionPattern; @@ -1397,24 +1419,6 @@ struct ReduceConverter : public OpConversionPattern { return success(); } - LogicalResult - convertToTensorExtract(triton::ReduceOp op, - typename triton::ReduceOp::Adaptor adaptor, - ConversionPatternRewriter &rewriter) const { - assert(llvm::hasSingleElement(op.getSrcs())); - - auto returnOp = cast(*op.getOps().begin()); - assert(llvm::hasSingleElement(returnOp.getResult())); - assert(cast(returnOp.getResult().front()).getArgNumber() == - 0); - - auto source = op.getSrcs().front(); - auto zeroIdx = - rewriter.createOrFold(op.getLoc(), 0); - rewriter.replaceOpWithNewOp(op, source, zeroIdx); - return success(); - } - public: LogicalResult matchAndRewrite(triton::ReduceOp op, @@ -1431,14 +1435,6 @@ struct ReduceConverter : public OpConversionPattern { "axis is within " "operand's rank"); - // Unsplat is implemented as a single element, rank 1 reduction where - // single element is yielded immediately. This can be simplified into - // a single element extract. - if (llvm::hasSingleElement(op.getOps()) && sourceType.getRank() == 1 && - sourceType.getShape()[0] == 1) { - return convertToTensorExtract(op, adaptor, rewriter); - } - return convertToLinalgReduce(op, adaptor, rewriter); } }; diff --git a/include/triton-shared/Dialect/TPtr/IR/TPtrDialect.td b/include/triton-shared/Dialect/TPtr/IR/TPtrDialect.td index 8333aa380..04d894ff4 100644 --- a/include/triton-shared/Dialect/TPtr/IR/TPtrDialect.td +++ b/include/triton-shared/Dialect/TPtr/IR/TPtrDialect.td @@ -109,9 +109,6 @@ def TPTR_TypeOffsetOp : TPTR_Op<"type_offset", [ConstantLike, Pure]> { let arguments = (ins TypeAttr:$baseType); let results = (outs AnySignlessIntegerOrIndex:$result); - let builders = [ - OpBuilder<(ins "TypeAttr":$baseType, CArg<"Type", "nullptr">:$resultTy)> - ]; let assemblyFormat = [{ attr-dict $baseType custom(type($result)) }]; diff --git a/lib/Conversion/StructuredToMemref/StructuredToMemref.cpp b/lib/Conversion/StructuredToMemref/StructuredToMemref.cpp index 4ee575a88..30fd93edd 100644 --- a/lib/Conversion/StructuredToMemref/StructuredToMemref.cpp +++ b/lib/Conversion/StructuredToMemref/StructuredToMemref.cpp @@ -577,21 +577,18 @@ struct MakeTensorPtrConverter struct MakeGatherScatterTensorPtrConverter : public OpConversionPattern { -private: - using OpConversionPattern::OpConversionPattern; - -public: - MakeGatherScatterTensorPtrConverter(const TypeConverter &typeConverter, - MLIRContext *context) - : OpConversionPattern(typeConverter, context) {} + using OpConversionPattern::OpConversionPattern; LogicalResult matchAndRewrite(tts::MakeGatherScatterTensorPtrOp op, OpAdaptor adaptor, ConversionPatternRewriter &rewriter) const override { // The gatherScatterPtr is rewritten as separate rows during load/store // operations. Therefore, no action is needed here except saving - // adaptor.getBase(). - rewriter.replaceOp(op, adaptor.getBase()); + // adaptor.getBase(). DialectConversion will ignore pure type conversion if + // we were to simply replace the op with adaptor.getBase(). To circumvent + // this we create an identity cast. + rewriter.replaceOpWithNewOp( + op, adaptor.getBase().getType(), adaptor.getBase()); return success(); } }; diff --git a/lib/Conversion/TritonArithToLinalg/TritonArithToLinalg.cpp b/lib/Conversion/TritonArithToLinalg/TritonArithToLinalg.cpp index 56b1b541d..44ecd7211 100644 --- a/lib/Conversion/TritonArithToLinalg/TritonArithToLinalg.cpp +++ b/lib/Conversion/TritonArithToLinalg/TritonArithToLinalg.cpp @@ -78,6 +78,7 @@ void mlir::triton::populateTritonArithToLinalgConversionPatterns( patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); diff --git a/lib/Conversion/TritonToLinalg/TritonToLinalg.cpp b/lib/Conversion/TritonToLinalg/TritonToLinalg.cpp index d85cab7ac..1c8ed9cf2 100644 --- a/lib/Conversion/TritonToLinalg/TritonToLinalg.cpp +++ b/lib/Conversion/TritonToLinalg/TritonToLinalg.cpp @@ -63,6 +63,7 @@ void mlir::triton::populateTritonToLinalgConversionPatterns( patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); + patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); patterns.add(patterns.getContext()); diff --git a/lib/Conversion/UnstructuredToMemref/UnstructuredToMemrefPass.cpp b/lib/Conversion/UnstructuredToMemref/UnstructuredToMemrefPass.cpp index 22e6ffc88..316537bfc 100644 --- a/lib/Conversion/UnstructuredToMemref/UnstructuredToMemrefPass.cpp +++ b/lib/Conversion/UnstructuredToMemref/UnstructuredToMemrefPass.cpp @@ -104,7 +104,7 @@ struct ScalarLoadConverter : public OpConversionPattern { auto zeroMap = AffineMap::getConstantMap(0, rewriter.getContext()); auto scalarLoadOp = rewriter.create( - loc, memref, zeroMap, std::nullopt); + loc, memref, zeroMap, ValueRange{}); rewriter.replaceOp(gatherOp, scalarLoadOp.getResult()); @@ -150,7 +150,7 @@ struct ScalarStoreConverter : public OpConversionPattern { auto zeroMap = AffineMap::getConstantMap(0, rewriter.getContext()); rewriter.create(loc, storeVal, memref, zeroMap, - std::nullopt); + ValueRange{}); rewriter.eraseOp(scatterOp); return success(); diff --git a/python/examples/conftest.py b/python/examples/conftest.py index 749761d3f..4bac989cc 100644 --- a/python/examples/conftest.py +++ b/python/examples/conftest.py @@ -18,6 +18,19 @@ def empty_decorator(func): def device(request): return "cpu" + +# this fixture is used for test_enable_fp_fusion +@pytest.fixture +def fresh_knobs(): + from triton._internal_testing import _fresh_knobs_impl + + fresh_function, reset_function = _fresh_knobs_impl() + try: + yield fresh_function() + finally: + reset_function() + + # this fixture is used for test_trans_4d && test_trans_reshape @pytest.fixture def with_allocator(): @@ -32,7 +45,7 @@ def with_allocator(): triton.set_allocator(NullAllocator()) -tests_supported = { +core_tests_supported = { "test_store_eviction_policy", "test_unary_op", "test_umulhi", @@ -77,6 +90,11 @@ def with_allocator(): "test_arange", } +annotations_tests_supported = { + "test_int_annotation", + "test_unknown_annotation", +} + def pytest_collection_modifyitems(config, items): skip_marker = pytest.mark.skip(reason="CPU backend does not support it yet") @@ -89,7 +107,11 @@ def pytest_collection_modifyitems(config, items): test_func_name = item.originalname if item.originalname else item.name test_file = str(item.fspath) - if test_file.endswith("test_core.py") and test_func_name not in tests_supported: + if test_file.endswith("test_core.py") and test_func_name not in core_tests_supported: + item.add_marker(skip_marker) + continue + + if test_file.endswith("test_annotations.py") and test_func_name not in annotations_tests_supported: item.add_marker(skip_marker) continue diff --git a/test/Conversion/TritonToLinalgExperimental/reduce_unsplat.mlir b/test/Conversion/TritonToLinalgExperimental/convert_unsplat.mlir similarity index 95% rename from test/Conversion/TritonToLinalgExperimental/reduce_unsplat.mlir rename to test/Conversion/TritonToLinalgExperimental/convert_unsplat.mlir index d97813532..9fd9cc52b 100644 --- a/test/Conversion/TritonToLinalgExperimental/reduce_unsplat.mlir +++ b/test/Conversion/TritonToLinalgExperimental/convert_unsplat.mlir @@ -6,10 +6,7 @@ module { %0 = tt.splat %arg0 : !tt.ptr -> tensor<1x!tt.ptr> %1 = tt.load %0 : tensor<1x!tt.ptr> %2 = arith.cmpi sgt, %1, %cst : tensor<1xi32> - %3 = "tt.reduce"(%2) <{axis = 0 : i32}> ({ - ^bb0(%arg1: i1, %arg2: i1): - tt.reduce.return %arg1 : i1 - }) : (tensor<1xi1>) -> i1 + %3 = tt.unsplat %2 : tensor<1xi1> scf.if %3 { tt.store %arg0, %c42_i32 : !tt.ptr } diff --git a/tools/triton-shared-opt/CMakeLists.txt b/tools/triton-shared-opt/CMakeLists.txt index 53a305c41..260c3e5d7 100644 --- a/tools/triton-shared-opt/CMakeLists.txt +++ b/tools/triton-shared-opt/CMakeLists.txt @@ -15,6 +15,7 @@ target_link_libraries(triton-shared-opt PRIVATE # MLIR core MLIROptLib MLIRPass + MLIRRegisterAllPasses MLIRTransforms ) diff --git a/tools/triton-shared-opt/RegisterTritonSharedDialects.h b/tools/triton-shared-opt/RegisterTritonSharedDialects.h index 5f30a8c36..6d92953fa 100644 --- a/tools/triton-shared-opt/RegisterTritonSharedDialects.h +++ b/tools/triton-shared-opt/RegisterTritonSharedDialects.h @@ -45,7 +45,7 @@ inline void registerTritonSharedDialects(mlir::DialectRegistry ®istry) { mlir::ttx::TritonTilingExtDialect, mlir::tts::TritonStructuredDialect, mlir::triton::TritonDialect, mlir::cf::ControlFlowDialect, mlir::math::MathDialect, mlir::arith::ArithDialect, mlir::scf::SCFDialect, - mlir::gpu::GPUDialect, mlir::linalg::LinalgDialect, - mlir::func::FuncDialect, mlir::tensor::TensorDialect, - mlir::memref::MemRefDialect, mlir::bufferization::BufferizationDialect>(); + mlir::linalg::LinalgDialect, mlir::func::FuncDialect, + mlir::tensor::TensorDialect, mlir::memref::MemRefDialect, + mlir::bufferization::BufferizationDialect>(); } From 9df315b99b95f77e51280643d3e4f44627f7af9e Mon Sep 17 00:00:00 2001 From: Micah Weston Date: Tue, 2 Sep 2025 15:12:07 +0000 Subject: [PATCH 2/3] Updates triton hash. --- triton-hash.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/triton-hash.txt b/triton-hash.txt index 7b9861abc..80f8308f4 100644 --- a/triton-hash.txt +++ b/triton-hash.txt @@ -1 +1 @@ -ec8cb09329cf25ac241a7dee1eea5a5d94daef8a +e44bd1c83c1c3e8deac7c4f02683cfb3cc395c8b From 71179e68188b795c1ea8f2ab05e81602cbf4273a Mon Sep 17 00:00:00 2001 From: Micah Weston Date: Tue, 2 Sep 2025 15:54:56 +0000 Subject: [PATCH 3/3] Adds comment about lack of bfloat16 and float16 support in the CPU backend. --- backend/driver.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/backend/driver.py b/backend/driver.py index a115eca14..dcfe9b743 100644 --- a/backend/driver.py +++ b/backend/driver.py @@ -54,6 +54,10 @@ def _ty_to_cpp(ty): "u16": "uint16_t", "u32": "uint32_t", "u64": "uint64_t", + # Proper support for bfloat16 and float16 is not yet handled. + # https://github.com/microsoft/triton-shared/issues/348 + # "fp16": "TODO", + # "bf16": "TODO", "fp32": "float", "f32": "float", "fp64": "double",