From a5c5aae16af41382e1a9441864b85ac2d614ff93 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar@google.com>
Date: Wed, 20 Nov 2024 11:45:49 -0800
Subject: [PATCH] [mlir] Make ArrayType shape return smallvector

Previously the ArrayRef returned was flagged as returning address of local temporary object. This is a little bit surprising that this works given different return type difference etc and may get reworked (beyond this seemingly satisfying ubsan).

PiperOrigin-RevId: 698466497
---
 xls/contrib/mlir/IR/xls_ops.td | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/xls/contrib/mlir/IR/xls_ops.td b/xls/contrib/mlir/IR/xls_ops.td
index ed481b35c8..5d9e4f8c7d 100644
--- a/xls/contrib/mlir/IR/xls_ops.td
+++ b/xls/contrib/mlir/IR/xls_ops.td
@@ -55,7 +55,9 @@ def Xls_ArrayType : Xls_Type<"Array", [ShapedTypeInterface]> {
   ];
   let extraClassDeclaration = [{
     constexpr bool hasRank() { return true; }
-    ::llvm::ArrayRef<int64_t> getShape() const { return getNumElements(); }
+    ::llvm::SmallVector<int64_t> getShape() const {
+      return ::llvm::SmallVector<int64_t>(1, getNumElements());
+    }
     ::mlir::ShapedType cloneWith(std::optional<llvm::ArrayRef<int64_t>> shape,
                                  Type elementType) {
       assert(0 && "ArrayType::cloneWith unimplemented!");