[mlir] Remove over specified memory effects

[youngar]

Summary

The standard and gpu dialect both have alloc operations which use the
memory effect MemAlloc. In both cases, it is specified on both the
operation itself and on the result. This results in two memory effects
being created for these operations. When MemAlloc is defined on an
operation, it represents some background effect which the compiler
cannot reason about, and inhibits the ability of the compiler to
remove dead std.alloc operations. This change removes the uneeded
MemAlloc effect from these operations and leaves the effect on the
result, which allows dead allocs to be erased.

There is the same problem, but to a lesser extent, with MemFree, MemRead
and MemWrite. Over-specifying these traits is not currently inhibiting
any optimization.

[aartbik]

LGTM for the sparse changes, it removes the index array when the sparse kernel does not need it

[rriddle]

Thanks!

mlir/include/mlir/Dialect/GPU/GPUOps.td

834

Can you remove the <DefaultResource> here too? It isn't necessary.

[youngar]

[mlir][gpu] Remove unneeded DefaultResource

DefaultResource is implied.

[youngar]

Should I get more reviews, or should I arc land this now?

[rriddle]

In D94662#2499581, @youngar wrote:

Should I get more reviews, or should I arc land this now?

If you have commit access, feel free to go ahead and land. If not, let me know and I can land it for you.

[youngar]

rebase

[youngar]

I manually rebased+pushed the commit to the github main branch. Will phabricator automatically close this? Did I do the something wrong?

[mehdi_amini]

Phab linked the commit as "added a commit" to this revision, I'm not sure why it didn't close it...

[youngar]

Thanks for the help and reviews everyone!

[herhut]

Thanks for cleaning this up!

mlir/docs/Interfaces.md

 ### Attribute/Operation/Type Interfaces
 
 Attribute/Operation/Type interfaces, as the names suggest, are those registered
 at the level of a specific attribute/operation/type. These interfaces provide
 access to derived objects by providing a virtual interface that must be
 implemented. As an example, many analyses and transformations want to reason
 about the side effects of an operation to improve performance and correctness.
 The side effects of an operation are generally tied to the semantics of a
-specific operation, for example an `affine.load` operation has a `write` effect
+specific operation, for example an `affine.load` operation has a `read` effect
 (as the name may suggest).
 
 These interfaces are defined by overriding the
 [CRTP](https://en.wikipedia.org/wiki/Curiously_recurring_template_pattern) class
 for the specific IR entity; `AttrInterface`, `OpInterface`, or `TypeInterface`
 respectively. These classes take, as a template parameter, a `Traits` class that
 defines a `Concept` and a `Model` class. These classes provide an implementation
 of concept-based polymorphism, where the `Concept` defines a set of virtual

mlir/include/mlir/Dialect/GPU/GPUOps.td

def GPU_WaitOp : GPU_Op<"wait", [GPU_AsyncOpInterface]> {
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) attr-dict
   }];
 }
 
 def GPU_AllocOp : GPU_Op<"alloc", [
     GPU_AsyncOpInterface,
-    AttrSizedOperandSegments,
-    MemoryEffects<[MemAlloc<DefaultResource>]>
+    AttrSizedOperandSegments
   ]> {
 
   let summary = "GPU memory allocation operation.";
   let description = [{
     The `gpu.alloc` operation allocates a region of memory on the GPU. It is
     similar to the `std.alloc` op, but supports asynchronous GPU execution.
 
     The op does not execute before all async dependencies have finished
     executing.
 
     If the `async` keyword is present, the op is executed asynchronously (i.e.
     it does not block until the execution has finished on the device). In
     that case, it also returns a !gpu.async.token.
 
     Example:
 
     ```mlir
     %memref, %token = gpu.alloc async [%dep] (%width) : memref<64x?xf32, 1>
     ```
   }];
 
   let arguments = (ins Variadic<GPU_AsyncToken>:$asyncDependencies,
                    Variadic<Index>:$dynamicSizes, Variadic<Index>:$symbolOperands);
-  let results = (outs Res<AnyMemRef, "", [MemAlloc<DefaultResource>]>:$memref,
+  let results = (outs Res<AnyMemRef, "", [MemAlloc]>:$memref,
                  Optional<GPU_AsyncToken>:$asyncToken);
 
   let extraClassDeclaration = [{
     MemRefType getType() { return memref().getType().cast<MemRefType>(); }
   }];
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies) ` `
     `(` $dynamicSizes `)` (`` `[` $symbolOperands^ `]`)? attr-dict `:` type($memref)
   }];
 }
 
-def GPU_DeallocOp : GPU_Op<"dealloc", [
-    GPU_AsyncOpInterface, MemoryEffects<[MemFree]>
-  ]> {
+def GPU_DeallocOp : GPU_Op<"dealloc", [GPU_AsyncOpInterface]> {
 
   let summary = "GPU memory deallocation operation";
 
   let description = [{
     The `gpu.dealloc` operation frees the region of memory referenced by a
     memref which was originally created by the `gpu.alloc` operation. It is
     similar to the `std.dealloc` op, but supports asynchronous GPU execution.
 
def GPU_DeallocOp : GPU_Op<"dealloc", [GPU_AsyncOpInterface]> {
   let results = (outs Optional<GPU_AsyncToken>:$asyncToken);
 
   let assemblyFormat = [{
     custom<AsyncDependencies>(type($asyncToken), $asyncDependencies)
     $memref attr-dict `:` type($memref)
   }];
 }
 
-def GPU_MemcpyOp : GPU_Op<"memcpy", [
-    GPU_AsyncOpInterface, MemoryEffects<[MemRead, MemWrite]>
-  ]> {
+def GPU_MemcpyOp : GPU_Op<"memcpy", [GPU_AsyncOpInterface]> {
 
   let summary = "GPU memcpy operation";
 
   let description = [{
     The `gpu.memcpy` operation copies the content of one memref to another.
 
     The op does not execute before all async dependencies have finished
     executing.

mlir/include/mlir/Dialect/StandardOps/IR/Ops.td

 //
 //   %0 = alloclike(%m)[%s] : memref<8x?xf32, (d0, d1)[s0] -> ((d0 + s0), d1)>
 //
 class AllocLikeOp<string mnemonic,
                   Resource resource,
                   list<OpTrait> traits = []> :
     Std_Op<mnemonic,
     !listconcat([
-      MemoryEffects<[MemAlloc<resource>]>,
       AttrSizedOperandSegments
     ], traits)> {
 
   let arguments = (ins Variadic<Index>:$dynamicSizes,
                    // The symbolic operands (the ones in square brackets) bind
                    // to the symbols of the memref's layout map.
                    Variadic<Index>:$symbolOperands,
                    Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$alignment);
let description = [{
     ```
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // DeallocOp
 //===----------------------------------------------------------------------===//
 
-def DeallocOp : Std_Op<"dealloc",
-    [MemoryEffects<[MemFree]>, MemRefsNormalizable]> {
+def DeallocOp : Std_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/test/Dialect/Linalg/sparse_1d.mlir

 
 // CHECK-LABEL:   func @sum_reduction(
 // CHECK-SAME:                        %[[VAL_0:.*]]: tensor<?xf32>,
 // CHECK-SAME:                        %[[VAL_1:.*]]: tensor<f32>) -> tensor<f32> {
 // CHECK:           %[[VAL_2:.*]] = constant 999 : index
 // CHECK:           %[[VAL_3:.*]] = constant 0 : index
 // CHECK:           %[[VAL_4:.*]] = constant 1 : index
 // CHECK:           %[[VAL_5:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
-// CHECK:           %[[VAL_6:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
 // CHECK:           %[[VAL_7:.*]] = alloca(%[[VAL_2]]) : memref<?xf32>
 // CHECK:           %[[VAL_8:.*]] = alloca() : memref<f32>
 // CHECK:           %[[VAL_9:.*]] = load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
 // CHECK:           %[[VAL_10:.*]] = load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
 // CHECK:           %[[VAL_11:.*]] = load %[[VAL_8]][] : memref<f32>
 // CHECK:           %[[VAL_12:.*]] = scf.for %[[VAL_13:.*]] = %[[VAL_9]] to %[[VAL_10]] step %[[VAL_4]] iter_args(%[[VAL_14:.*]] = %[[VAL_11]]) -> (f32) {
 // CHECK:             %[[VAL_15:.*]] = load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xf32>
 // CHECK:             %[[VAL_16:.*]] = addf %[[VAL_14]], %[[VAL_15]] : f32

mlir/test/Dialect/Linalg/sparse_2d.mlir

 // CHECK-LABEL:   func @sum_reduction(
 // CHECK-SAME:                        %[[VAL_0:.*]]: tensor<10x20xf32>,
 // CHECK-SAME:                        %[[VAL_1:.*]]: tensor<f32>) -> tensor<f32> {
 // CHECK:           %[[VAL_2:.*]] = constant 999 : index
 // CHECK:           %[[VAL_3:.*]] = constant 10 : index
 // CHECK:           %[[VAL_4:.*]] = constant 0 : index
 // CHECK:           %[[VAL_5:.*]] = constant 1 : index
 // CHECK:           %[[VAL_6:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
-// CHECK:           %[[VAL_7:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
 // CHECK:           %[[VAL_8:.*]] = alloca(%[[VAL_2]]) : memref<?xf32>
 // CHECK:           %[[VAL_9:.*]] = alloca() : memref<f32>
 // CHECK:           scf.for %[[VAL_10:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
 // CHECK:             %[[VAL_11:.*]] = load %[[VAL_6]]{{\[}}%[[VAL_10]]] : memref<?xindex>
 // CHECK:             %[[VAL_12:.*]] = addi %[[VAL_10]], %[[VAL_5]] : index
 // CHECK:             %[[VAL_13:.*]] = load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex>
 // CHECK:             %[[VAL_14:.*]] = load %[[VAL_9]][] : memref<f32>
 // CHECK:             %[[VAL_15:.*]] = scf.for %[[VAL_16:.*]] = %[[VAL_11]] to %[[VAL_13]] step %[[VAL_5]] iter_args(%[[VAL_17:.*]] = %[[VAL_14]]) -> (f32) {

mlir/test/Dialect/Linalg/sparse_3d.mlir

 
 // CHECK-LABEL:   func @sum_reduction(
 // CHECK-SAME:                        %[[VAL_0:.*]]: tensor<10x20x30xf32>,
 // CHECK-SAME:                        %[[VAL_1:.*]]: tensor<f32>) -> tensor<f32> {
 // CHECK:           %[[VAL_2:.*]] = constant 999 : index
 // CHECK:           %[[VAL_3:.*]] = constant 0 : index
 // CHECK:           %[[VAL_4:.*]] = constant 1 : index
 // CHECK:           %[[VAL_5:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
-// CHECK:           %[[VAL_6:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
 // CHECK:           %[[VAL_7:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
-// CHECK:           %[[VAL_8:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
 // CHECK:           %[[VAL_9:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
-// CHECK:           %[[VAL_10:.*]] = alloca(%[[VAL_2]]) : memref<?xindex>
 // CHECK:           %[[VAL_11:.*]] = alloca(%[[VAL_2]]) : memref<?xf32>
 // CHECK:           %[[VAL_12:.*]] = alloca() : memref<f32>
 // CHECK:           %[[VAL_13:.*]] = load %[[VAL_5]]{{\[}}%[[VAL_3]]] : memref<?xindex>
 // CHECK:           %[[VAL_14:.*]] = load %[[VAL_5]]{{\[}}%[[VAL_4]]] : memref<?xindex>
 // CHECK:           scf.for %[[VAL_15:.*]] = %[[VAL_13]] to %[[VAL_14]] step %[[VAL_4]] {
 // CHECK:             %[[VAL_16:.*]] = load %[[VAL_7]]{{\[}}%[[VAL_15]]] : memref<?xindex>
 // CHECK:             %[[VAL_17:.*]] = addi %[[VAL_15]], %[[VAL_4]] : index
 // CHECK:             %[[VAL_18:.*]] = load %[[VAL_7]]{{\[}}%[[VAL_17]]] : memref<?xindex>

mlir/test/Transforms/canonicalize.mlir

func @dead_block_elim() {
   }
   return
 
 ^bb1:
   return
 }
 
 // CHECK-LABEL: func @dyn_shape_fold(%arg0: index, %arg1: index)
-func @dyn_shape_fold(%L : index, %M : index) -> (memref<? x ? x i32>, memref<? x ? x f32>) {
+func @dyn_shape_fold(%L : index, %M : index) -> (memref<? x ? x i32>, memref<? x ? x f32>, memref<4 x ? x 8 x ? x ? x f32>) {
   // CHECK: %c0 = constant 0 : index
   %zero = constant 0 : index
   // The constants below disappear after they propagate into shapes.
   %nine = constant 9 : index
   %N = constant 1024 : index
   %K = constant 512 : index
 
   // CHECK-NEXT: alloc(%arg0) : memref<?x1024xf32>
affine.for %i = 0 to %L {
     affine.for %j = 0 to 10 {
       // CHECK-NEXT: load %0[%arg2, %arg3] : memref<?x1024xf32>
       // CHECK-NEXT: store %{{.*}}, %1[%c0, %c0, %arg2, %arg3, %c0] : memref<4x1024x8x512x?xf32>
       %v = load %a[%i, %j] : memref<?x?xf32>
       store %v, %b[%zero, %zero, %i, %j, %zero] : memref<4x?x8x?x?xf32>
     }
   }
 
-  return %c, %d : memref<? x ? x i32>, memref<? x ? x f32>
+  return %c, %d, %e : memref<? x ? x i32>, memref<? x ? x f32>, memref<4 x ? x 8 x ? x ? x f32>
 }
 
 #map1 = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
 #map2 = affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0 * s2 + d1 * s1 + d2 + s0)>
 #map3 = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
 
 // CHECK-LABEL: func @dim_op_fold(
 // CHECK-SAME: %[[ARG0:[a-z0-9]*]]: index