[BACKEND] Fix wrong check in tmem_alloc canonicalization pattern (#7719)

Author: ThomasRaoux

lib/Dialect/TritonGPU/Transforms/HoistTMEMAlloc.cpp

@@ -177,7 +177,7 @@ class CombineTMEMStoreAndAlloc : public OpRewritePattern<ttng::TMEMStoreOp> {
     auto alloc = store.getDep().getDefiningOp<TMEMTokenAllocOp>();
     if (!alloc)
       return failure();
-    if (store.getSrc() != alloc.getResult())
+    if (store.getDst() != alloc.getResult())
       return failure();
     if (alloc->getBlock() != store->getBlock())
       return failure();

test/TritonGPU/hoist-tmem-alloc.mlir

@@ -341,13 +341,15 @@ module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.targ
 #tmem = #ttng.tensor_memory_encoding<blockM = 128, blockN = 128, unpacked = true>
 module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 4 : i32, ttg.target = "cuda:100", "ttg.threads-per-warp" = 32 : i32} {
   tt.func public @forward_tmem_load(%m: !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>, %t: !ttg.async.token) -> (!ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>, !ttg.async.token) {
+    %true = arith.constant true
     %result, %token0 = ttng.tmem_load %m[%t] : !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable> -> tensor<128x128xf32, #blocked>
     // HOIST-IF-LABEL: @forward_tmem_load
     // HOIST-IF-SAME:    %[[ARG0:.+]]: !ttg.memdesc<128x128xf32,
     // HOIST-IF-SAME:    %[[ARG1:.+]]: !ttg.async.token
     // HOIST-IF-NEXT:    tt.return %[[ARG0]], %[[ARG1]]
-    %result1, %token1 = ttng.tmem_alloc %result : (tensor<128x128xf32, #blocked>) -> (!ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>, !ttg.async.token)
-    tt.return %result1, %token1 : !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>, !ttg.async.token
+    %result1, %token1 = ttng.tmem_alloc : () -> (!ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>, !ttg.async.token)
+    %token2 = ttng.tmem_store %result, %result1[%token1], %true : tensor<128x128xf32, #blocked> -> !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>
+    tt.return %result1, %token2 : !ttg.memdesc<128x128xf32, #tmem, #ttng.tensor_memory, mutable>, !ttg.async.token
   }
 }