Fixed bug in buffer deallocation pass using unranked memref types.

[dfki-jugr]

Summary

In the buffer deallocation pass, unranked memref types are not properly supported.
After investigating this issue, it turns out that the Clone and Dealloc operation
does not support unranked memref types in the current implementation.
This patch adds the missing feature and enables the transformation of any memref
type.

This patch solves this bug: https://bugs.llvm.org/show_bug.cgi?id=48385

[herhut]

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

355

Add let verifier = ? here and below to avoid specifying empty verifier.

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

478

This is already checked by the operand type definitions. You can drop the entire function.

551

Drop this function, as well.

[dfki-jugr]

Addressed comments.

[herhut]

Thanks!

mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td

let description = [{
     ```mlir
     %arg1 = memref.clone %arg0 : memref<?xf32> to memref<?xf32>
     ```
 
     Note, that mutating the source or result of the clone operation leads to
     undefined behavior.
   }];
 
-  let arguments = (ins Arg<AnyMemRef, "", []>:$input);
-  let results = (outs Arg<AnyMemRef, "", []>:$output);
+  let arguments = (ins Arg<AnyRankedOrUnrankedMemRef, "", []>:$input);
+  let results = (outs Arg<AnyRankedOrUnrankedMemRef, "", []>:$output);
 
   let extraClassDeclaration = [{
     Value getSource() { return input();}
     Value getTarget() { return output(); }
   }];
 
   let assemblyFormat = "$input attr-dict `:` type($input) `to` type($output)";
 
   let hasFolder = 1;
   let hasCanonicalizer = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // DeallocOp
 //===----------------------------------------------------------------------===//
 
 def MemRef_DeallocOp : MemRef_Op<"dealloc", [MemRefsNormalizable]> {
   let summary = "memory deallocation operation";
   let description = [{
     The `dealloc` operation frees the region of memory referenced by a memref
     which was originally created by the `alloc` operation.
     The `dealloc` operation should not be called on memrefs which alias an
     alloc'd memref (e.g. memrefs returned by `view` operations).
 
     Example:
 
     ```mlir
     %0 = memref.alloc() : memref<8x64xf32, (d0, d1) -> (d0, d1), 1>
     memref.dealloc %0 : memref<8x64xf32, (d0, d1) -> (d0, d1), 1>
     ```
   }];
 
-  let arguments = (ins Arg<AnyMemRef, "", [MemFree]>:$memref);
+  let arguments = (ins Arg<AnyRankedOrUnrankedMemRef, "", [MemFree]>:$memref);
 
   let hasFolder = 1;
   let assemblyFormat = "$memref attr-dict `:` type($memref)";
 }
 
 //===----------------------------------------------------------------------===//
 // DimOp
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/MemRef/IR/MemRefOps.cpp

 OpFoldResult CastOp::fold(ArrayRef<Attribute> operands) {
   return succeeded(foldMemRefCast(*this)) ? getResult() : Value();
 }
 
 //===----------------------------------------------------------------------===//
 // CloneOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(CloneOp op) { return success(); }
+static LogicalResult verify(CloneOp op) {
+  if (!op.input().getType().isa<MemRefType, UnrankedMemRefType>())
+    return op.emitOpError("operand must be a memref");
+  return success();
+}
 
 void CloneOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
   effects.emplace_back(MemoryEffects::Read::get(), input(),
                        SideEffects::DefaultResource::get());
   effects.emplace_back(MemoryEffects::Write::get(), output(),
                        SideEffects::DefaultResource::get());

 OpFoldResult CloneOp::fold(ArrayRef<Attribute> operands) {
   return succeeded(foldMemRefCast(*this)) ? getResult() : Value();
 }
 
 //===----------------------------------------------------------------------===//
 // DeallocOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(DeallocOp op) {
-  if (!op.memref().getType().isa<MemRefType>())
+  if (!op.memref().getType().isa<MemRefType, UnrankedMemRefType>())
     return op.emitOpError("operand must be a memref");
   return success();
 }
 
 LogicalResult DeallocOp::fold(ArrayRef<Attribute> cstOperands,
                               SmallVectorImpl<OpFoldResult> &results) {
   /// dealloc(memrefcast) -> dealloc
   return foldMemRefCast(*this);

mlir/test/Conversion/StandardToSPIRV/alloc.mlir

 
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {}>
   }
 {
   func @alloc_dealloc_dynamic_workgroup_mem(%arg0 : memref<4x?xf32, 3>) {
     // expected-error @+2 {{unhandled deallocation type}}
-    // expected-error @+1 {{'memref.dealloc' op operand #0 must be memref of any type values}}
+    // expected-error @+1 {{'memref.dealloc' op operand #0 must be unranked.memref of any type values or memref of any type values}}
     memref.dealloc %arg0 : memref<4x?xf32, 3>
     return
   }
 }
 
 // -----
 
 module attributes {
   spv.target_env = #spv.target_env<
     #spv.vce<v1.0, [Shader], [SPV_KHR_storage_buffer_storage_class]>, {}>
   }
 {
   func @alloc_dealloc_mem(%arg0 : memref<4x5xf32>) {
     // expected-error @+2 {{unhandled deallocation type}}
-    // expected-error @+1 {{op operand #0 must be memref of any type values}}
+    // expected-error @+1 {{op operand #0 must be unranked.memref of any type values or memref of any type values}}
     memref.dealloc %arg0 : memref<4x5xf32>
     return
   }
 }

mlir/test/Dialect/MemRef/ops.mlir

 }
 
 // CHECK-LABEL: func @read_global_memref
 func @read_global_memref() {
   %0 = memref.get_global @memref0 : memref<2xf32>
   %1 = memref.tensor_load %0 : memref<2xf32>
   return
 }
+
+// CHECK-LABEL: func @memref_clone
+func @memref_clone() {
+  %0 = memref.alloc() : memref<2xf32>
+  %1 = memref.cast %0 : memref<2xf32> to memref<*xf32>
+  %2 = memref.clone %1 : memref<*xf32> to memref<*xf32>
+  return
+}
+
+// CHECK-LABEL: func @memref_dealloc
+func @memref_dealloc() {
+  %0 = memref.alloc() : memref<2xf32>
+  %1 = memref.cast %0 : memref<2xf32> to memref<*xf32>
+  memref.dealloc %1 : memref<*xf32>
+  return
+}

mlir/test/Transforms/buffer-deallocation.mlir

 // CHECK-NEXT: br ^bb3
 // CHECK-NEXT: ^bb3(%[[ALLOC3:.*]]:{{.*}})
 //      CHECK: test.copy(%[[ALLOC3]],
 // CHECK-NEXT: memref.dealloc %[[ALLOC3]]
 // CHECK-NEXT: return
 
 // -----
 
+// Test case: See above.
+
+// CHECK-LABEL: func @condBranchUnrankedType
+func @condBranchUnrankedType(
+  %arg0: i1,
+  %arg1: memref<*xf32>,
+  %arg2: memref<*xf32>,
+  %arg3: index) {
+  cond_br %arg0, ^bb1, ^bb2(%arg3: index)
+^bb1:
+  br ^bb3(%arg1 : memref<*xf32>)
+^bb2(%0: index):
+  %1 = memref.alloc(%0) : memref<?xf32>
+  %2 = memref.cast %1 : memref<?xf32> to memref<*xf32>
+  test.buffer_based in(%arg1: memref<*xf32>) out(%2: memref<*xf32>)
+  br ^bb3(%2 : memref<*xf32>)
+^bb3(%3: memref<*xf32>):
+  test.copy(%3, %arg2) : (memref<*xf32>, memref<*xf32>)
+  return
+}
+
+// CHECK-NEXT: cond_br
+//      CHECK: %[[ALLOC0:.*]] = memref.clone
+// CHECK-NEXT: br ^bb3(%[[ALLOC0]]
+//      CHECK: ^bb2(%[[IDX:.*]]:{{.*}})
+// CHECK-NEXT: %[[ALLOC1:.*]] = memref.alloc(%[[IDX]])
+//      CHECK: test.buffer_based
+// CHECK-NEXT: %[[ALLOC2:.*]] = memref.clone
+// CHECK-NEXT: memref.dealloc %[[ALLOC1]]
+// CHECK-NEXT: br ^bb3
+// CHECK-NEXT: ^bb3(%[[ALLOC3:.*]]:{{.*}})
+//      CHECK: test.copy(%[[ALLOC3]],
+// CHECK-NEXT: memref.dealloc %[[ALLOC3]]
+// CHECK-NEXT: return
+
+// -----
+
 // Test Case:
 //      bb0
 //     /    \
 //   bb1    bb2 <- Initial position of AllocOp
 //    |     /  \
 //    |   bb3  bb4
 //    |     \  /
 //    \     bb5

mlir/test/lib/Dialect/Test/TestOps.td

 //===----------------------------------------------------------------------===//
 // Test CopyOpInterface
 //===----------------------------------------------------------------------===//
 
 def CopyOp : TEST_Op<"copy", [CopyOpInterface]> {
   let description = [{
     Represents a copy operation.
   }];
-  let arguments = (ins Res<AnyMemRef, "", [MemRead]>:$source,
-                   Res<AnyMemRef, "", [MemWrite]>:$target);
+  let arguments = (ins Res<AnyRankedOrUnrankedMemRef, "", [MemRead]>:$source,
+                   Res<AnyRankedOrUnrankedMemRef, "", [MemWrite]>:$target);
   let assemblyFormat = [{
     `(` $source `,` $target `)` `:` `(` type($source) `,` type($target) `)`
      attr-dict
   }];
   let extraClassDeclaration = [{
     Value getSource() { return source(); }
     Value getTarget() { return target(); }
   }];
def RegionYieldOp : TEST_Op<"region_yield",
   ];
 }
 
 class BufferBasedOpBase<string mnemonic, list<OpTrait> traits>
     : TEST_Op<mnemonic, traits> {
   let description = [{
     A buffer based operation, that uses memRefs as input and output.
   }];
-  let arguments = (ins AnyMemRef:$input, AnyMemRef:$output);
+  let arguments = (ins AnyRankedOrUnrankedMemRef:$input,
+                       AnyRankedOrUnrankedMemRef:$output);
 }
 
 def BufferBasedOp : BufferBasedOpBase<"buffer_based", []>{
   let assemblyFormat = [{
     `in` `(` $input`:` type($input) `)` `out` `(` $output`:` type($output) `)`
     attr-dict
   }];
 }