Update triton to e44bd1c83c1c3e8deac7c4f02683cfb3cc395c8b (#345)

This updates to latest Triton main which should resolve the nightly
build issues.

Primary changes:
* Change conversion of tts::MakeGatherScatterTensorPtrOp to insert
identity cast so that DialectConversion does not ignore the type
converted operand of the associated tts::LoadOp when supplying the op
adaptor to the pattern.
* Move support of unsplat to the new dedicated operation since they do
not do reduction anymore
* Add missing test figure to conftest
* Disable float annotation tests since bfloat16 and float16 are not
supported in CPU backend
* Add missing link library to registering of passes
* Update construction of ValueRange{} which was causing compilation
errors.
* Remove unused builder in TPtrOps since it caused linker errors.
* Remove bfloat16 and float16 from CPU backend since not testing is
present. Better to crash during compilation than get runtime errors.
* Remove GPUDialect that is no longer present upstream

Author: red1bluelost

.gitmodules

@@ -1,3 +0,0 @@
-[submodule "triton"]
-	path = triton
-	url = https://github.com/triton-lang/triton.git

backend/driver.py

@@ -54,8 +54,10 @@ def _ty_to_cpp(ty):
         "u16": "uint16_t",
         "u32": "uint32_t",
         "u64": "uint64_t",
-        "fp16": "float",
-        "bf16": "float",
+        # 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",

include/triton-shared/Conversion/TritonArithToLinalg/ConversionPatterns.hpp

@@ -360,7 +360,7 @@ struct LoadConverter : public OpConversionPattern<triton::LoadOp> {
                                                   loc, rewriter);
       auto zeroMap = AffineMap::getConstantMap(0, rewriter.getContext());
       auto loadOp = rewriter.create<affine::AffineLoadOp>(
-          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<triton::StoreOp> {
           PtrAnalysis::getScalarMemRef(op.getPtr(), ptr, loc, rewriter);
       auto zeroMap = AffineMap::getConstantMap(0, rewriter.getContext());
       rewriter.create<affine::AffineStoreOp>(loc, val, sMemRef, zeroMap,
-                                             std::nullopt);
+                                             ValueRange{});
       rewriter.eraseOp(op);
       return success();
     }
@@ -649,6 +649,28 @@ struct SplatConverter : public OpConversionPattern<triton::SplatOp> {
   }
 };
 
+struct UnsplatConverter : public OpConversionPattern<triton::UnsplatOp> {
+  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<Value, 1> indices(tensorType.getRank());
+    if (indices.size() > 0) {
+      auto zeroIdx =
+          rewriter.createOrFold<arith::ConstantIndexOp>(op.getLoc(), 0);
+      llvm::fill(indices, zeroIdx);
+    }
+
+    rewriter.replaceOpWithNewOp<tensor::ExtractOp>(op, adaptor.getSrc(),
+                                                   indices);
+    return success();
+  }
+};
+
 struct BroadcastConverter : public OpConversionPattern<triton::BroadcastOp> {
 private:
   using OpConversionPattern<triton::BroadcastOp>::OpConversionPattern;
@@ -1397,24 +1419,6 @@ struct ReduceConverter : public OpConversionPattern<triton::ReduceOp> {
     return success();
   }
 
-  LogicalResult
-  convertToTensorExtract(triton::ReduceOp op,
-                         typename triton::ReduceOp::Adaptor adaptor,
-                         ConversionPatternRewriter &rewriter) const {
-    assert(llvm::hasSingleElement(op.getSrcs()));
-
-    auto returnOp = cast<triton::ReduceReturnOp>(*op.getOps().begin());
-    assert(llvm::hasSingleElement(returnOp.getResult()));
-    assert(cast<BlockArgument>(returnOp.getResult().front()).getArgNumber() ==
-           0);
-
-    auto source = op.getSrcs().front();
-    auto zeroIdx =
-        rewriter.createOrFold<arith::ConstantIndexOp>(op.getLoc(), 0);
-    rewriter.replaceOpWithNewOp<tensor::ExtractOp>(op, source, zeroIdx);
-    return success();
-  }
-
 public:
   LogicalResult
   matchAndRewrite(triton::ReduceOp op,
@@ -1431,14 +1435,6 @@ struct ReduceConverter : public OpConversionPattern<triton::ReduceOp> {
            "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);
   }
 };

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<IntType>(type($result))
   }];

lib/Conversion/StructuredToMemref/StructuredToMemref.cpp

@@ -577,21 +577,18 @@ struct MakeTensorPtrConverter
 
 struct MakeGatherScatterTensorPtrConverter
     : public OpConversionPattern<tts::MakeGatherScatterTensorPtrOp> {
-private:
-  using OpConversionPattern<tts::MakeGatherScatterTensorPtrOp>::OpConversionPattern;
-
-public:
-  MakeGatherScatterTensorPtrConverter(const TypeConverter &typeConverter,
-                         MLIRContext *context)
-      : OpConversionPattern<tts::MakeGatherScatterTensorPtrOp>(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<UnrealizedConversionCastOp>(
+        op, adaptor.getBase().getType(), adaptor.getBase());
     return success();
   }
 };

lib/Conversion/TritonArithToLinalg/TritonArithToLinalg.cpp

@@ -78,6 +78,7 @@ void mlir::triton::populateTritonArithToLinalgConversionPatterns(
   patterns.add<ClampConverter>(patterns.getContext());
   patterns.add<MatmulConverter>(patterns.getContext());
   patterns.add<SplatConverter>(patterns.getContext());
+  patterns.add<UnsplatConverter>(patterns.getContext());
   patterns.add<DenseConstantConverter>(patterns.getContext());
   patterns.add<CumSumConverter>(patterns.getContext());
   patterns.add<ReshapeConverter>(patterns.getContext());

lib/Conversion/TritonToLinalg/TritonToLinalg.cpp

@@ -63,6 +63,7 @@ void mlir::triton::populateTritonToLinalgConversionPatterns(
   patterns.add<AssertConverter>(patterns.getContext());
   patterns.add<MatmulConverter>(patterns.getContext());
   patterns.add<SplatConverter>(patterns.getContext());
+  patterns.add<UnsplatConverter>(patterns.getContext());
   patterns.add<DenseConstantConverter>(patterns.getContext());
   patterns.add<UnrealizedCastConverter>(patterns.getContext());
   patterns.add<CumSumConverter>(patterns.getContext());

lib/Conversion/UnstructuredToMemref/UnstructuredToMemrefPass.cpp

@@ -104,7 +104,7 @@ struct ScalarLoadConverter : public OpConversionPattern<tts::GatherOp> {
     auto zeroMap = AffineMap::getConstantMap(0, rewriter.getContext());
 
     auto scalarLoadOp = rewriter.create<affine::AffineLoadOp>(
-        loc, memref, zeroMap, std::nullopt);
+        loc, memref, zeroMap, ValueRange{});
 
     rewriter.replaceOp(gatherOp, scalarLoadOp.getResult());
 
@@ -150,7 +150,7 @@ struct ScalarStoreConverter : public OpConversionPattern<tts::ScatterOp> {
     auto zeroMap = AffineMap::getConstantMap(0, rewriter.getContext());
 
     rewriter.create<affine::AffineStoreOp>(loc, storeVal, memref, zeroMap,
-                                           std::nullopt);
+                                           ValueRange{});
     rewriter.eraseOp(scatterOp);
 
     return success();

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
 

test/Conversion/TritonToLinalgExperimental/reduce_unsplat.mlir renamed to test/Conversion/TritonToLinalgExperimental/convert_unsplat.mlir

@@ -6,10 +6,7 @@ module {
     %0 = tt.splat %arg0 : !tt.ptr<i32> -> tensor<1x!tt.ptr<i32>>
     %1 = tt.load %0 : tensor<1x!tt.ptr<i32>>
     %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<i32>
     }