[Mlir-commits] [mlir] [mlir][memref] Fix alloca lowering with 0 dimensions (PR #111119)

[Matthias Springer]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/111119

The `memref.alloca` lowering computed the allocation size incorrectly when there were 0 dimensions.

Previously:
```
memref.alloca() : memref<10x0x2xf32>
--> llvm.alloca 20xf32
```

Now:
```
memref.alloca() : memref<10x0x2xf32>
--> llvm.alloca 0xf32
```

>From the `llvm.alloca` documentation:
```
Allocating zero bytes is legal, but the returned pointer may not be unique.
```

>From 3e3b23f6b4bcfa1a30a2a5a2acc4b73868829720 Mon Sep 17 00:00:00 2001
From: Matthias Springer <mspringer at nvidia.com>
Date: Fri, 4 Oct 2024 10:32:13 +0200
Subject: [PATCH] [mlir][memref] Fix alloca lowering with 0 dimensions

The `memref.alloca` lowering computed the allocation size incorrectly when there were 0 dimensions.

Previously:
```
memref.alloca() : memref<10x0x2xf32>
--> llvm.alloca 20xf32
```

Now:
```
memref.alloca() : memref<10x0x2xf32>
--> llvm.alloca 0xf32
```

>From the `llvm.alloca` documentation:
```
Allocating zero bytes is legal, but the returned pointer may not be unique.
```
---
 mlir/lib/Conversion/LLVMCommon/Pattern.cpp       |  2 --
 .../MemRefToLLVM/convert-static-memref-ops.mlir  | 16 ++++++++++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Conversion/LLVMCommon/Pattern.cpp b/mlir/lib/Conversion/LLVMCommon/Pattern.cpp
index 1886dfa870961a..d551506485a454 100644
--- a/mlir/lib/Conversion/LLVMCommon/Pattern.cpp
+++ b/mlir/lib/Conversion/LLVMCommon/Pattern.cpp
@@ -139,8 +139,6 @@ void ConvertToLLVMPattern::getMemRefDescriptorSizes(
     strides[i] = runningStride;
 
     int64_t staticSize = memRefType.getShape()[i];
-    if (staticSize == 0)
-      continue;
     bool useSizeAsStride = stride == 1;
     if (staticSize == ShapedType::kDynamic)
       stride = ShapedType::kDynamic;
diff --git a/mlir/test/Conversion/MemRefToLLVM/convert-static-memref-ops.mlir b/mlir/test/Conversion/MemRefToLLVM/convert-static-memref-ops.mlir
index 96cf2264e9c2f3..4f4ce0d0ed978e 100644
--- a/mlir/test/Conversion/MemRefToLLVM/convert-static-memref-ops.mlir
+++ b/mlir/test/Conversion/MemRefToLLVM/convert-static-memref-ops.mlir
@@ -95,6 +95,22 @@ func.func @static_alloca() -> memref<32x18xf32> {
 
 // -----
 
+// CHECK-LABEL: func @static_alloca_zero()
+func.func @static_alloca_zero() -> memref<32x0x18xf32> {
+// CHECK: %[[sz1:.*]] = llvm.mlir.constant(32 : index) : i64
+// CHECK: %[[sz2:.*]] = llvm.mlir.constant(0 : index) : i64
+// CHECK: %[[sz3:.*]] = llvm.mlir.constant(18 : index) : i64
+// CHECK: %[[st1:.*]] = llvm.mlir.constant(1 : index) : i64
+// CHECK: %[[st2:.*]] = llvm.mlir.constant(0 : index) : i64
+// CHECK: %[[num_elems:.*]] = llvm.mlir.constant(0 : index) : i64
+// CHECK: %[[allocated:.*]] = llvm.alloca %[[num_elems]] x f32 : (i64) -> !llvm.ptr
+ %0 = memref.alloca() : memref<32x0x18xf32>
+
+ return %0 : memref<32x0x18xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @static_dealloc
 func.func @static_dealloc(%static: memref<10x8xf32>) {
 // CHECK: %[[ptr:.*]] = llvm.extractvalue %{{.*}}[0] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>