[mlir][bufferization] Privatize buffers for parallel regions

[springerm]

Summary

One-Shot Bufferize correctly handles RaW conflicts around repetitive regions (loops). Specical handling is needed for parallel regions. These are a special kind of repetitive regions that can have additional RaW conflicts that would not be present if the regions would be executed sequentially.

Example:

%0 = bufferization.alloc_tensor()
scf.forall ... {
  %1 = linalg.fill ins(...) outs(%0)
  ...
  scf.forall.in_parallel {
    tensor.parallel_insert_slice %1 into ...
  }
}

A separate (private) buffer must be allocated for each iteration of the scf.forall loop.

This change adds a new interface method to BufferizableOpInterface to detect parallel regions. By default, regions are assumed to be sequential.

A buffer is privatized if an OpOperand bufferizes to a memory read inside a parallel region that is different from the parallel region where operand's value is defined.

[maerhart]

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

687

region?

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

127

Nit: maybe assert that it is also a repetitive region since that is requested in the documentation?

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

551

Does this include reads of the result of the region-containing op? I.e., there are writes inside the parallel region, but they are basically discarded (like doing nothing in the in_parallel), is it guaranteed that they won't show up?

[springerm]

address comments

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

551

If an alias of the scf.forall result (which is also an alias of the init_arg) is written in the loop body and read afterwards, something (usually the op inside the loop) must bufferize out-of-place. This is irrespective of parallelism. But I added a test case where there is no read that bufferizes in-place.

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h

 /// owner of the block. In case of an OpResult that is the defining op.
 Operation *getOwnerOfValue(Value value);
 
 /// Assuming that the given region is repetitive, find the next enclosing
 /// repetitive region.
 Region *getNextEnclosingRepetitiveRegion(Region *region,
                                          const BufferizationOptions &options);
 
+/// If `region` is a parallel region, return `region`. Otherwise, find the first
+/// enclosing parallel region of `region`. If there is no such region, return
+/// "nullptr".
+///
+/// Note: Whether a region is parallel or sequential is queried from the
+/// `BufferizableOpInterface`.
+Region *getParallelRegion(Region *region, const BufferizationOptions &options);
+
 namespace detail {
 /// This is the default implementation of
 /// BufferizableOpInterface::getAliasingOpOperands. Should not be called from
 /// other places.
 AliasingOpOperandList defaultGetAliasingOpOperands(Value value,
                                                    const AnalysisState &state);
 
 /// This is the default implementation of

mlir/include/mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td

let methods = [
         /*methodName=*/"isRepetitiveRegion",
         /*args=*/(ins "unsigned":$index),
         /*methodBody=*/"",
         /*defaultImplementation=*/[{
           return ::mlir::bufferization::detail::defaultIsRepetitiveRegion(
               ::llvm::cast<BufferizableOpInterface>($_op.getOperation()), index);
         }]
       >,
+      InterfaceMethod<
+        /*desc=*/[{
+          Return `true` if the given region of this op is parallel, i.e.,
+          multiple instances of the region may be executing at the same time.
+          If a region is parallel, it must also be marked as "repetitive".
+
+          The RaW conflict detection of One-Shot Analysis is more strict inside
+          parallel regions: Buffer may have to be privatized.
+
+          By default, regions are assumed to be sequential.
+        }],
+        /*retType=*/"bool",
+        /*methodName=*/"isParallelRegion",
+        /*args=*/(ins "unsigned":$index),
+        /*methodBody=*/"",
+        /*defaultImplementation=*/[{
+          return false;
+        }]
+      >,
       StaticInterfaceMethod<
         /*desc=*/[{
           Return `true` if the op and this interface implementation supports
           unstructured control flow. I.e., regions with multiple blocks. This is
           not supported in most ops, so the default answer is `false`.
         }],
         /*retType=*/"bool",
         /*methodName=*/"supportsUnstructuredControlFlow",

mlir/lib/Dialect/Bufferization/IR/BufferizableOpInterface.cpp

Region *bufferization::getNextEnclosingRepetitiveRegion(
   assert(isRepetitiveRegion(region, options) && "expected repetitive region");
   while ((region = region->getParentRegion())) {
     if (isRepetitiveRegion(region, options))
       break;
   }
   return region;
 }
 
+Region *bufferization::getParallelRegion(Region *region,
+                                         const BufferizationOptions &options) {
+  while (region) {
+    auto bufferizableOp = options.dynCastBufferizableOp(region->getParentOp());
+    if (bufferizableOp &&
+        bufferizableOp.isParallelRegion(region->getRegionNumber())) {
+      assert(isRepetitiveRegion(region, options) &&
+             "expected that all parallel regions are also repetitive regions");
+      return region;
+    }
+    region = region->getParentRegion();
+  }
+  return nullptr;
+}
+
 Operation *bufferization::getOwnerOfValue(Value value) {
   if (auto opResult = llvm::dyn_cast<OpResult>(value))
     return opResult.getDefiningOp();
   return llvm::cast<BlockArgument>(value).getOwner()->getParentOp();
 }
 
 bool bufferization::allocationDoesNotEscape(OpResult opResult) {
 #ifndef NDEBUG

mlir/lib/Dialect/Bufferization/Transforms/OneShotAnalysis.cpp

 /// actually reads another definition W2.
 static bool
 hasReadAfterWriteInterference(const DenseSet<OpOperand *> &usesRead,
                               const DenseSet<OpOperand *> &usesWrite,
                               const DominanceInfo &domInfo,
                               OneShotAnalysisState &state) {
   const BufferizationOptions &options = state.getOptions();
 
+  // Before going through the main RaW analysis, find cases where a buffer must
+  // be privatized due to parallelism. If the result of a write is never read,
+  // privatization is not necessary (and large parts of the IR are likely dead).
+  if (!usesRead.empty()) {
+    for (OpOperand *uConflictingWrite : usesWrite) {
+      // Find the allocation point or last write (definition) of the buffer.
+      // Note: In contrast to `findDefinitions`, this also returns results of
+      // ops that do not bufferize to memory write when no other definition
+      // could be found. E.g., "bufferization.alloc_tensor" would be included,
+      // even though that op just bufferizes to an allocation but does define
+      // the contents of the buffer.
+      SetVector<Value> definitionsOrLeaves =
+          state.findValueInReverseUseDefChain(
+              uConflictingWrite->get(),
+              [&](Value v) { return state.bufferizesToMemoryWrite(v); });
+      assert(!definitionsOrLeaves.empty() &&
+             "expected at least one definition or leaf");
+
+      // The writing op must bufferize out-of-place if the definition is in a
+      // different parallel region than this write.
+      for (Value def : definitionsOrLeaves) {
+        if (getParallelRegion(def.getParentRegion(), options) !=
+            getParallelRegion(uConflictingWrite->getOwner()->getParentRegion(),
+                              options)) {
+          LLVM_DEBUG(
+              llvm::dbgs()
+              << "\n- bufferizes out-of-place due to parallel region:\n");
+          LLVM_DEBUG(llvm::dbgs()
+                     << "  unConflictingWrite = operand "
+                     << uConflictingWrite->getOperandNumber() << " of "
+                     << *uConflictingWrite->getOwner() << "\n");
+          return true;
+        }
+      }
+    }
+  }
+
   for (OpOperand *uRead : usesRead) {
     Operation *readingOp = uRead->getOwner();
     LLVM_DEBUG(llvm::dbgs() << "\n- check conflict:\n");
     LLVM_DEBUG(llvm::dbgs() << "  uRead = operand " << uRead->getOperandNumber()
                             << " of " << *readingOp << "\n");
 
     // Find the definition of uRead by following the SSA use-def chain.
     // E.g.:

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

for (auto [lb, ub, step] :
       if (!stepConstant)
         return true;
 
       if (*lbConstant + *stepConstant < *ubConstant)
         return true;
     }
     return false;
   }
+
+  bool isParallelRegion(Operation *op, unsigned index) const {
+    return isRepetitiveRegion(op, index);
+  }
 };
 
 /// Nothing to do for InParallelOp.
 struct InParallelOpInterface
     : public BufferizableOpInterface::ExternalModel<InParallelOpInterface,
                                                     InParallelOp> {
   LogicalResult bufferize(Operation *op, RewriterBase &b,
                           const BufferizationOptions &options) const {

mlir/test/Dialect/SCF/one-shot-bufferize-analysis.mlir

scf.for %iv2 = %a to %b step %c {
         %2 = tensor.extract_slice %t[0][4][1] : tensor<10xf32> to tensor<4xf32>
         %3 = tensor.extract %2[%a] : tensor<4xf32>
         vector.print %3 : f32
       }
     }
   }
   return
 }
+
+// -----
+
+// CHECK-LABEL: func @parallel_region()
+func.func @parallel_region() -> tensor<320xf32>
+{
+  %alloc0 = bufferization.alloc_tensor() : tensor<320xf32>
+  %alloc1 = bufferization.alloc_tensor() : tensor<1xf32>
+  %c320 = arith.constant 320 : index
+  // CHECK: scf.forall
+  %0 = scf.forall (%arg0) in (%c320) shared_outs(%arg1 = %alloc0) -> (tensor<320xf32>) {
+    %val = "test.foo"() : () -> (f32)
+    // linalg.fill must bufferize out-of-place because every thread needs a
+    // private copy of %alloc1.
+    // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "false"]}
+    %fill = linalg.fill ins(%val : f32) outs(%alloc1 : tensor<1xf32>) -> tensor<1xf32>
+    scf.forall.in_parallel {
+      // CHECK: tensor.parallel_insert_slice {{.*}} {__inplace_operands_attr__ = ["true", "true", "none"]}
+      tensor.parallel_insert_slice %fill into %arg1[%arg0] [1] [1] : tensor<1xf32> into tensor<320xf32>
+    }
+  }
+  // CHECK: } {__inplace_operands_attr__ = ["none", "true"]}
+  return %0 : tensor<320xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @parallel_region_mixed_def(
+func.func @parallel_region_mixed_def(%c: i1) -> tensor<320xf32>
+{
+  %alloc0 = bufferization.alloc_tensor() : tensor<320xf32>
+  %alloc1 = bufferization.alloc_tensor() : tensor<1xf32>
+  %c320 = arith.constant 320 : index
+  // CHECK: scf.forall
+  %0 = scf.forall (%arg0) in (%c320) shared_outs(%arg1 = %alloc0) -> (tensor<320xf32>) {
+    %alloc2 = bufferization.alloc_tensor() : tensor<1xf32>
+    %selected = scf.if %c -> tensor<1xf32> {
+      scf.yield %alloc1 : tensor<1xf32>
+    } else {
+      scf.yield %alloc2 : tensor<1xf32>
+    }
+    %val = "test.foo"() : () -> (f32)
+    // linalg.fill must bufferize out-of-place because every thread needs a
+    // private copy of %alloc1.
+    // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "false"]}
+    %fill = linalg.fill ins(%val : f32) outs(%selected : tensor<1xf32>) -> tensor<1xf32>
+    scf.forall.in_parallel {
+      // CHECK: tensor.parallel_insert_slice {{.*}} {__inplace_operands_attr__ = ["true", "true", "none"]}
+      tensor.parallel_insert_slice %fill into %arg1[%arg0] [1] [1] : tensor<1xf32> into tensor<320xf32>
+    }
+  }
+  // CHECK: } {__inplace_operands_attr__ = ["none", "true"]}
+  return %0 : tensor<320xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @parallel_region_two_writes(
+func.func @parallel_region_two_writes(%f: f32) -> tensor<320xf32>
+{
+  %alloc0 = bufferization.alloc_tensor() : tensor<320xf32>
+  %alloc1 = bufferization.alloc_tensor() : tensor<1xf32>
+  %c320 = arith.constant 320 : index
+  %c0 = arith.constant 0 : index
+  // CHECK: scf.forall
+  %0 = scf.forall (%arg0) in (%c320) shared_outs(%arg1 = %alloc0) -> (tensor<320xf32>) {
+    %val = "test.foo"() : () -> (f32)
+    // linalg.fill must bufferize out-of-place because every thread needs a
+    // private copy of %alloc1.
+    // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "false"]}
+    %fill = linalg.fill ins(%val : f32) outs(%alloc1 : tensor<1xf32>) -> tensor<1xf32>
+    // CHECK: tensor.insert
+    // CHECK-SAME: __inplace_operands_attr__ = ["none", "true", "none"]
+    %inserted = tensor.insert %f into %fill[%c0] : tensor<1xf32>
+
+    scf.forall.in_parallel {
+      // CHECK: tensor.parallel_insert_slice {{.*}} {__inplace_operands_attr__ = ["true", "true", "none"]}
+      tensor.parallel_insert_slice %inserted into %arg1[%arg0] [1] [1] : tensor<1xf32> into tensor<320xf32>
+    }
+  }
+  // CHECK: } {__inplace_operands_attr__ = ["none", "true"]}
+  return %0 : tensor<320xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @parallel_region_no_read()
+func.func @parallel_region_no_read()
+{
+  %alloc0 = bufferization.alloc_tensor() : tensor<320xf32>
+  %alloc1 = bufferization.alloc_tensor() : tensor<1xf32>
+  %c320 = arith.constant 320 : index
+  // CHECK: scf.forall
+  scf.forall (%arg0) in (%c320) {
+    %val = "test.foo"() : () -> (f32)
+    // linalg.fill can bufferize in-place because no alias of %alloc1 is read.
+    // CHECK: linalg.fill {__inplace_operands_attr__ = ["none", "true"]}
+    %fill = linalg.fill ins(%val : f32) outs(%alloc1 : tensor<1xf32>) -> tensor<1xf32>
+    scf.forall.in_parallel {
+    }
+  }
+  return
+}