[BACKEND] Preserve tt attrs in AddPtr combine and canonicalize (#7113)

Combine and canonicalize remove the attributes that relate to alignment
so we don't get vectorization of stores in some cases, this allows them
to keep it.

Author: nzaghen

include/triton/Dialect/Triton/IR/DiscardableAttributes.h

@@ -0,0 +1,15 @@
+#ifndef TRITON_DIALECT_TRITON_IR_DISCARDABLE_ATTRIBUTES_H_
+#define TRITON_DIALECT_TRITON_IR_DISCARDABLE_ATTRIBUTES_H_
+
+#include "mlir/Support/LLVM.h"
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+namespace mlir::triton {
+
+// Filter out attributes from the given operation that are not present in
+// the allowList.
+[[nodiscard]] SmallVector<NamedAttribute>
+filterDiscardableAttrs(Operation *op, ArrayRef<StringRef> allowList);
+
+} // namespace mlir::triton
+#endif // TRITON_DIALECT_TRITON_IR_DISCARDABLE_ATTRIBUTES_H_

lib/Dialect/Triton/IR/CMakeLists.txt

@@ -4,6 +4,7 @@ add_public_tablegen_target(TritonCanonicalizeIncGen)
 
 add_triton_library(TritonIR
   Dialect.cpp
+  DiscardableAttributes.cpp
   Ops.cpp
   Traits.cpp
   Types.cpp

lib/Dialect/Triton/IR/DiscardableAttributes.cpp

@@ -0,0 +1,17 @@
+#include "mlir/Support/LLVM.h"
+#include "triton/Dialect/Triton/IR/Dialect.h"
+
+namespace mlir::triton {
+
+SmallVector<NamedAttribute>
+filterDiscardableAttrs(Operation *op, ArrayRef<StringRef> allowList) {
+  SmallVector<NamedAttribute> propagatedAttrs;
+  for (auto attrName : allowList) {
+    Attribute attr = op->getDiscardableAttr(attrName);
+    if (attr)
+      propagatedAttrs.emplace_back(attrName, attr);
+  }
+  return propagatedAttrs;
+}
+
+} // namespace mlir::triton

lib/Dialect/Triton/Transforms/Combine.cpp

@@ -6,6 +6,7 @@
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
+#include "triton/Dialect/Triton/IR/DiscardableAttributes.h"
 #include "triton/Dialect/Triton/Transforms/Passes.h"
 
 namespace mlir::triton {

lib/Dialect/Triton/Transforms/Combine.td

@@ -39,9 +39,15 @@ def CombineDotAddFRevPattern : Pat<
 //   Note: leave (sub %c0, %c0) canceling to ArithDialect
 //         (ref: ArithCanonicalization.td)
 defvar DefOverflow = ConstantEnumCase<Arith_IntegerOverflowAttr, "none">;
+
+def CopyDiscardableAttrs: NativeCodeCallVoid<
+        "$1.getOwner()->setDiscardableAttrs(triton::filterDiscardableAttrs($0.getOwner(), "
+        "{\"tt.divisibility\", \"tt.contiguity\", \"tt.constancy\"}))">;
+
 def CombineAddPtrPattern : Pat<
-        (TT_AddPtrOp (TT_AddPtrOp $ptr, $idx0), $idx1),
-        (TT_AddPtrOp $ptr, (Arith_AddIOp $idx0, $idx1, DefOverflow)),
-        [(Constraint<CPred<"isAddPtrOffsetCombinable($0, $1)">> $idx0, $idx1)]>;
+        (TT_AddPtrOp:$src (TT_AddPtrOp $ptr, $idx0), $idx1),
+        (TT_AddPtrOp:$dest $ptr, (Arith_AddIOp $idx0, $idx1, DefOverflow)),
+        [(Constraint<CPred<"isAddPtrOffsetCombinable($0, $1)">> $idx0, $idx1)],
+        [(CopyDiscardableAttrs $src, $dest)]>;
 
 #endif

test/Triton/combine.mlir

@@ -84,6 +84,29 @@ tt.func @test_combine_addptr_pattern(%base: !tt.ptr<f32>) -> tensor<8x!tt.ptr<f3
     tt.return %ptr1 : tensor<8x!tt.ptr<f32>>
 }
 
+// CHECK-LABEL: @test_combine_addptr_pattern_discardableattrs
+tt.func @test_combine_addptr_pattern_discardableattrs(%base: !tt.ptr<f32>) -> !tt.ptr<f32> {
+    %off0 = arith.constant 8 : i32
+    %off1 = arith.constant 4 : i32
+    // CHECK-NEXT: %[[cst:.*]] = arith.constant 12 : i32
+    // CHECK-NEXT: %0 = tt.addptr %{{.*}}, %[[cst]] {tt.constancy = 8 : i32, tt.contiguity = 512 : i32, tt.divisibility = 16 : i32} : !tt.ptr<f32>, i32
+    %ptr0 = tt.addptr %base, %off0 : !tt.ptr<f32>, i32
+    %ptr1 = tt.addptr %ptr0, %off1 {tt.divisibility = 16 : i32, tt.constancy = 8 : i32, tt.contiguity = 512 : i32} : !tt.ptr<f32>, i32
+
+    tt.return %ptr1 : !tt.ptr<f32>
+}
+
+// CHECK-LABEL: @test_combine_addptr_pattern_discardableattrs_disallowed
+tt.func @test_combine_addptr_pattern_discardableattrs_disallowed(%base: !tt.ptr<f32>) -> !tt.ptr<f32> {
+    %off0 = arith.constant 8 : i32
+    %off1 = arith.constant 4 : i32
+    // CHECK-NEXT: %[[cst:.*]] = arith.constant 12 : i32
+    // CHECK-NEXT: %0 = tt.addptr %{{.*}}, %[[cst]] {tt.divisibility = 16 : i32} : !tt.ptr<f32>, i32
+    %ptr0 = tt.addptr %base, %off0 : !tt.ptr<f32>, i32
+    %ptr1 = tt.addptr %ptr0, %off1 {tt.divisibility = 16 : i32, tt.disallowed = 8 : i32} : !tt.ptr<f32>, i32
+
+    tt.return %ptr1 : !tt.ptr<f32>
+}
 // CHECK-LABEL: @test_combine_addptr_pattern_i64
 tt.func @test_combine_addptr_pattern_i64(%base: !tt.ptr<f32>) -> tensor<8x!tt.ptr<f32>> {
     %off0 = arith.constant 10 : i64

third_party/amd/lib/TritonAMDGPUTransforms/CanonicalizePointers.cpp

@@ -17,6 +17,7 @@
 #include "mlir/Transforms/DialectConversion.h"
 #include "triton/Analysis/Utility.h"
 #include "triton/Dialect/Triton/IR/Dialect.h"
+#include "triton/Dialect/Triton/IR/DiscardableAttributes.h"
 #include "triton/Dialect/Triton/IR/Types.h"
 #include "triton/Dialect/TritonGPU/Transforms/Utility.h"
 #include "llvm/ADT/STLExtras.h"
@@ -561,12 +562,20 @@ class ConvertAddPtrOp : public PointerCanonicalizationPattern<tt::AddPtrOp> {
     RewriterBase::InsertionGuard guard(rewriter);
     rewriter.setInsertionPoint(addPtrOp);
 
+    // Query all discardable attributes that we want to preserve
+    std::array<StringRef, 3> propagateList{"tt.divisibility", "tt.contiguity",
+                                           "tt.constancy"};
+    SmallVector<NamedAttribute> propagatedAttrs =
+        tt::filterDiscardableAttrs(addPtrOp.getOperation(), propagateList);
+
     // If it is a scalar pointer update, simply bump the base pointer
     if (llvm::isa<tt::PointerType>(addPtrOp.getPtr().getType())) {
       assert(llvm::isa<IntegerType>(origOffset.getType()) &&
              "expected offset to be integer type");
       auto newAddPtrOp = rewriter.create<tt::AddPtrOp>(
           curLoc, fatPtrBase.getType(), fatPtrBase, origOffset);
+      newAddPtrOp->setDiscardableAttrs(propagatedAttrs);
+
       rewriter.replaceOpWithMultiple(addPtrOp, {{newAddPtrOp, fatPtrOffset}});
       fatPtrs[{newAddPtrOp, fatPtrOffset}] =
           fatPtrs.at({fatPtrBase, fatPtrOffset});
@@ -581,6 +590,8 @@ class ConvertAddPtrOp : public PointerCanonicalizationPattern<tt::AddPtrOp> {
             maybeGetOrCreateScalarConstant(rewriter, curLoc, origOffset)) {
       tt::AddPtrOp newAddPtrOp = rewriter.create<tt::AddPtrOp>(
           curLoc, fatPtrBase.getType(), fatPtrBase, *scalarConst);
+      newAddPtrOp->setDiscardableAttrs(propagatedAttrs);
+
       rewriter.replaceOpWithMultiple(addPtrOp, {{newAddPtrOp, fatPtrOffset}});
       // If we are updating the tensor pointer with a constant value, we can
       // propagate the attributes of the tensor pointer to the fat pointer.
@@ -596,6 +607,7 @@ class ConvertAddPtrOp : public PointerCanonicalizationPattern<tt::AddPtrOp> {
 
     auto newAddPtrOp = rewriter.create<tt::AddPtrOp>(
         curLoc, fatPtrBase.getType(), fatPtrBase, uniformOffset);
+    newAddPtrOp->setDiscardableAttrs(propagatedAttrs);
 
     // Vector offset update (if any): bump the tensor offset
     bool canNarrow = fatPtrs.at({fatPtrBase, fatPtrOffset}).canNarrow;