From 4dac289e65a38774a9b72ae09669890151122bff Mon Sep 17 00:00:00 2001
From: Goran Flegar <gflegar@google.com>
Date: Thu, 11 Jan 2024 13:48:22 +0100
Subject: [PATCH] Fix use-after-free detected by address sanitizer (#2914)

`DistributedEncodingTrait::getCTAOrder()` returns a SmallVector by
value, which is deleted as soon as it is assigned to `ref`. `ref` then
becomes a dangling reference.

To prevent that, we now use a vector instead of an array reference.
---
 lib/Dialect/TritonGPU/IR/Dialect.cpp | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/lib/Dialect/TritonGPU/IR/Dialect.cpp b/lib/Dialect/TritonGPU/IR/Dialect.cpp
index 025b3265f9..47cd0712e1 100644
--- a/lib/Dialect/TritonGPU/IR/Dialect.cpp
+++ b/lib/Dialect/TritonGPU/IR/Dialect.cpp
@@ -300,17 +300,17 @@ SmallVector<unsigned> getCTASplitNum(Attribute layout) {
 }
 
 SmallVector<unsigned> getCTAOrder(Attribute layout) {
-  ArrayRef<unsigned> ref;
+  SmallVector<unsigned> res;
   if (auto distributedLayout = layout.dyn_cast<DistributedEncodingTrait>()) {
-    ref = distributedLayout.getCTAOrder();
+    res = distributedLayout.getCTAOrder();
   } else if (auto mfmaLayout = layout.dyn_cast<MfmaEncodingAttr>()) {
     return {0, 1};
   } else if (auto sharedLayout = layout.dyn_cast<SharedEncodingAttr>()) {
-    ref = sharedLayout.getCTALayout().getCTAOrder();
+    res = SmallVector<unsigned>(sharedLayout.getCTALayout().getCTAOrder());
   } else {
     llvm::report_fatal_error("Unimplemented usage of getCTAOrder");
   }
-  return SmallVector<unsigned>(ref.begin(), ref.end());
+  return res;
 }
 
 SmallVector<int64_t> getShapePerCTA(ArrayRef<unsigned> CTASplitNum,