[llvm] [mlir][bufferization] Empty tensor elimination for materialize_in_destination (PR #65468)

Matthias Springer via llvm-commits llvm-commits at lists.llvm.org
Tue Sep 12 00:46:02 PDT 2023


================
@@ -1147,13 +1023,117 @@ struct SplatOpInterface
   }
 };
 
+namespace {
+/// Return "true" if `insertSliceOp` inserts into a subset that is equivalent
+/// to the subset defined by `candidate`. `equivalenceFn` is used to determine
+/// equivalence of tensors.
+template <typename OpTy>
+bool isSubsetEquivalentToInsertSliceLikeOp(
+    OpTy insertSliceOp, Value candidate,
+    function_ref<bool(Value, Value)> equivalenceFn) {
+  // Look for a matching tensor.extract_slice op.
+  auto extractSliceOp = candidate.getDefiningOp<tensor::ExtractSliceOp>();
+  if (!extractSliceOp)
+    return false;
+  if (!equivalenceFn(extractSliceOp.getSource(), insertSliceOp.getDest()))
+    return false;
+  if (!sameOffsetsSizesAndStrides(extractSliceOp, insertSliceOp,
+                                  isEqualConstantIntOrValue))
+    return false;
+  return true;
+}
+
+template <typename OpTy>
+Value getSubsetOfInsertSliceLikeOp(OpBuilder &b, Location loc,
+                                   OpTy insertSliceOp) {
+  auto extractOp = b.create<tensor::ExtractSliceOp>(
+      loc, insertSliceOp.getSourceType(), insertSliceOp.getDest(),
+      insertSliceOp.getMixedOffsets(), insertSliceOp.getMixedSizes(),
+      insertSliceOp.getMixedStrides());
+  return extractOp.getResult();
+}
+
+template <typename OpTy>
+SmallVector<Value>
+getValuesNeededToBuildSubsetOfInsertSliceLikeOp(OpTy insertSliceOp) {
+  SmallVector<Value> neededValues;
+  // Collect all values that are needed to construct the replacement op.
+  neededValues.append(insertSliceOp.getOffsets().begin(),
+                      insertSliceOp.getOffsets().end());
+  neededValues.append(insertSliceOp.getSizes().begin(),
+                      insertSliceOp.getSizes().end());
+  neededValues.append(insertSliceOp.getStrides().begin(),
+                      insertSliceOp.getStrides().end());
+  neededValues.push_back(insertSliceOp.getDest());
+  return neededValues;
+}
+} // namespace
+
+struct InsertSliceOpSubsetOpInterface
----------------
matthias-springer wrote:

This is so that the `MLIRTensorDialect` build target does not have to depend on `MLIRBufferizationDialect`. That's also why the `BufferizableOpInterface` implementations in this file are implemented as external models. It also keeps `TensorOps.cpp` and `TensorOps.td` a bit simpler.

https://github.com/llvm/llvm-project/pull/65468


More information about the llvm-commits mailing list