[mlir][linalg][bufferize] Fix buffer equivalence around scf.if ops

[springerm]

Summary

Also extend the comments for aliasInfo and equivalenceInfo.

mlir/include/mlir/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.h

private:
   using EquivalenceClassRangeType = llvm::iterator_range<
       llvm::EquivalenceClasses<Value, ValueComparator>::member_iterator>;
   /// Check that aliasInfo for `v` exists and return a reference to it.
   EquivalenceClassRangeType getAliases(Value v) const;
 
   /// Set of tensors that are known to bufferize to writable memory.
   llvm::DenseSet<Value> bufferizeToWritableMemory;
 
-  /// Auxiliary structure to store all the values a given value aliases with.
-  /// These are the conservative cases that can further decompose into
-  /// "equivalent" buffer relationships.
+  /// Auxiliary structure to store all the values a given value may alias with.
+  /// Alias information is "may be" conservative: In the presence of branches, a
+  /// value may alias with one of multiple other values. The concrete aliasing
+  /// value may not even be known at compile time. All such values are
+  /// considered to be aliases.
   llvm::EquivalenceClasses<Value, ValueComparator> aliasInfo;
 
-  /// Auxiliary structure to store all the equivalent buffer classes.
+  /// Auxiliary structure to store all the equivalent buffer classes. Equivalent
+  /// buffer information is "must be" conservative: Only if two values are
+  /// guaranteed to be equivalent at runtime, they said to be equivalent. It is
+  /// possible that, in the presence of branches, it cannot be determined
+  /// statically if two values are equivalent. In that case, the values are
+  /// considered to be not equivalent.
   llvm::EquivalenceClasses<Value, ValueComparator> equivalentInfo;
 };
 
 /// Analyze the `ops` to determine which OpResults are inplaceable.
 LogicalResult inPlaceAnalysis(SmallVector<Operation *> &ops,
                               BufferizationAliasInfo &aliasInfo,
                               const DominanceInfo &domInfo,
                               unsigned analysisFuzzerSeed = 0);

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ComprehensiveBufferize.cpp

BufferizationAliasInfo::BufferizationAliasInfo(Operation *rootOp) {
 
   // The return value of an scf::IfOp aliases with both yield values.
   rootOp->walk([&](scf::IfOp ifOp) {
     if (ifOp->getNumResults() > 0) {
       for (auto it : llvm::zip(ifOp.thenYield().results(),
                                ifOp.elseYield().results(), ifOp.results())) {
         aliasInfo.unionSets(std::get<0>(it), std::get<1>(it));
         aliasInfo.unionSets(std::get<0>(it), std::get<2>(it));
-        equivalentInfo.unionSets(std::get<0>(it), std::get<1>(it));
-        equivalentInfo.unionSets(std::get<0>(it), std::get<2>(it));
       }
     }
   });
 }
 
 /// Add a new entry for `v` in the `aliasInfo` and `equivalentInfo`. In the
 /// beginning the alias and equivalence sets only contain `v` itself.
 void BufferizationAliasInfo::createAliasInfoEntry(Value v) {
static Value getResultBuffer(OpBuilder &b, OpResult result,
   Operation *op = result.getOwner();
   SmallVector<OpOperand *> aliasingOperands = getAliasingOpOperand(result);
   assert(!aliasingOperands.empty() && "could not get aliasing OpOperand");
   Value operand = aliasingOperands.front()->get();
   Value operandBuffer = lookup(bvm, operand);
   assert(operandBuffer && "operand buffer not found");
   // Make sure that all OpOperands are the same buffer. If this is not the case,
   // we would have to materialize a memref value.
+  // TODO: Should be looking for checking for "equivalent buffers" instead of
+  // operator== here, but equivalent buffers for scf.if yield values are not
+  // set up yet.
   if (!llvm::all_of(aliasingOperands, [&](OpOperand *o) {
         return lookup(bvm, o->get()) == operandBuffer;
       })) {
     op->emitError("result buffer is ambiguous");
     return Value();
   }
 
   // If bufferizing out-of-place, allocate a new buffer.

mlir/test/Dialect/Linalg/comprehensive-module-bufferize-invalid.mlir

^bb1:
     }
     return %T#0, %T#1 : tensor<f32>, tensor<f32>
   ^bb2:
     return %t2, %t1 : tensor<f32>, tensor<f32>
 }
 
 // -----
 
+func @scf_if_not_equivalent(
+    %cond: i1, %t1: tensor<?xf32> {linalg.inplaceable = true},
+    %idx: index) -> tensor<?xf32> {
+  // expected-error @+1 {{result buffer is ambiguous}}
+  %r = scf.if %cond -> (tensor<?xf32>) {
+    scf.yield %t1 : tensor<?xf32>
+  } else {
+    // This buffer aliases, but is not equivalent.
+    %t2 = tensor.extract_slice %t1 [%idx] [%idx] [1] : tensor<?xf32> to tensor<?xf32>
+    scf.yield %t2 : tensor<?xf32>
+  }
+  return %r : tensor<?xf32>
+}
+
+// -----
+
 // expected-error @-3 {{expected callgraph to be free of circular dependencies}}
 
 func @foo() {
   call @bar() : () -> ()
   return
 }
 
 func @bar() {

mlir/test/Dialect/Linalg/comprehensive-module-bufferize.mlir

func @scf_if_inplace(%cond: i1,
   %r = scf.if %cond -> (tensor<?xf32>) {
     scf.yield %t1 : tensor<?xf32>
   } else {
     %t2 = vector.transfer_write %v, %t1[%idx] : vector<5xf32>, tensor<?xf32>
     scf.yield %t2 : tensor<?xf32>
   }
   return %r : tensor<?xf32>
 }
-