From 4db0de4fdccc0681917a7447a1aa7f85acd2ff7e Mon Sep 17 00:00:00 2001
From: Mogball <jeffniu22@gmail.com>
Date: Thu, 10 Jul 2025 14:10:41 -0700
Subject: [PATCH 1/2] [Backend] Fix alias analysis with `ub.poison`

Optimistically assume poison values alias nothing.
---
 lib/Analysis/Alias.cpp        |  4 ++++
 test/Analysis/test-alias.mlir | 14 ++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/lib/Analysis/Alias.cpp b/lib/Analysis/Alias.cpp
index 1817a9000658..bceb49ee842d 100644
--- a/lib/Analysis/Alias.cpp
+++ b/lib/Analysis/Alias.cpp
@@ -2,6 +2,7 @@
 
 #include "mlir/Support/LLVM.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
+#include "mlir/Dialect/UB/IR/UBOps.h"
 
 namespace mlir {
 
@@ -38,6 +39,9 @@ LogicalResult SharedMemoryAliasAnalysis::visitOperation(
   } else if (op->hasTrait<OpTrait::MemDescViewTrait>()) {
     aliasInfo = AliasInfo(operands[0]->getValue());
     pessimistic = false;
+  } else if (isa<ub::PoisonOp>(op)) {
+    aliasInfo = AliasInfo();
+    pessimistic = false;
   } else {
     assert(!isa<triton::gpu::MemDescType>(result.getType()) &&
            "unknown operation creating memory descriptor");
diff --git a/test/Analysis/test-alias.mlir b/test/Analysis/test-alias.mlir
index 2b60fc6df4e8..380ad08a3c1b 100644
--- a/test/Analysis/test-alias.mlir
+++ b/test/Analysis/test-alias.mlir
@@ -191,4 +191,18 @@ tt.func @cf_for(%arg0: index, %arg1: index, %arg2: index, %arg3: !tt.ptr<f16>, %
   tt.return
 }
 
+tt.func @poison_memdesc(%arg0: i1) {
+  // expected-remark @below {{%0 -> %0}}
+  %0 = ttg.local_alloc : () -> !ttg.memdesc<128x32xf16, #A_SHARED, #ttg.shared_memory, mutable>
+  cf.cond_br %arg0, ^bb1, ^bb2(%0 : !ttg.memdesc<128x32xf16, #A_SHARED, #ttg.shared_memory, mutable>)
+^bb1:
+  %1 = ub.poison : !ttg.memdesc<128x32xf16, #A_SHARED, #ttg.shared_memory, mutable>
+  cf.br ^bb2(%1 : !ttg.memdesc<128x32xf16, #A_SHARED, #ttg.shared_memory, mutable>)
+^bb2(%2: !ttg.memdesc<128x32xf16, #A_SHARED, #ttg.shared_memory, mutable>):
+  %c0_i32 = arith.constant 0 : i32
+  // expected-remark @below {{%3 -> %0}}
+  %3 = ttg.memdesc_subview %2[%c0_i32, %c0_i32] : !ttg.memdesc<128x32xf16, #A_SHARED, #ttg.shared_memory, mutable> -> !ttg.memdesc<128x32xf16, #A_SHARED, #ttg.shared_memory, mutable>
+  tt.return
+}
+
 }  // module

From 29b1510a113264dc4da7102c33a1e8c072c9d394 Mon Sep 17 00:00:00 2001
From: Mogball <jeffniu22@gmail.com>
Date: Thu, 10 Jul 2025 14:21:20 -0700
Subject: [PATCH 2/2] fmt

---
 lib/Analysis/Alias.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lib/Analysis/Alias.cpp b/lib/Analysis/Alias.cpp
index bceb49ee842d..1a997b2234e1 100644
--- a/lib/Analysis/Alias.cpp
+++ b/lib/Analysis/Alias.cpp
@@ -1,8 +1,8 @@
 #include "triton/Analysis/Alias.h"
 
+#include "mlir/Dialect/UB/IR/UBOps.h"
 #include "mlir/Support/LLVM.h"
 #include "triton/Dialect/TritonGPU/IR/Dialect.h"
-#include "mlir/Dialect/UB/IR/UBOps.h"
 
 namespace mlir {