Add test and address comments

Author: hsjts0u

mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp

@@ -1063,7 +1063,7 @@ tileAndFuseFirstExtractUse(RewriterBase &rewriter, Diagnostic &diag,
       resultNumber, offsets, sizes);
 
   // Cleanup clone.
-  if (dyn_cast<LoopLikeOpInterface>(containingOp))
+  if (isa<LoopLikeOpInterface>(containingOp))
     rewriter.eraseOp(tileableProducer);
 
   return std::make_tuple(tileAndFuseResult->tiledOps, newContainingOp);

mlir/test/Dialect/Linalg/transform-op-fuse-into-containing.mlir

@@ -253,6 +253,40 @@ module {
 
 // -----
 
+#map = affine_map<(d0) -> (d0 * 2)>
+#map1 = affine_map<(d0) -> (d0 * 4)>
+module {
+  // CHECK-LABEL: func.func @fuse_tileable_op_no_dps
+  func.func @fuse_tileable_op_no_dps(%arg0: tensor<4x4x4xf32>, %arg1: tensor<4x4x4xf32>) -> tensor<4x4x4xf32> {
+    %0 = "test.tiling_no_dps_op"(%arg0, %arg1) : (tensor<4x4x4xf32>, tensor<4x4x4xf32>) -> tensor<4x4x4xf32>
+    %1 = tensor.empty() : tensor<4x4x4xf32>
+    // CHECK: scf.forall
+    %2 = scf.forall (%arg2, %arg3, %arg4) in (4, 2, 1) shared_outs(%arg5 = %1) -> (tensor<4x4x4xf32>) {
+      %3 = affine.apply #map(%arg3)
+      %4 = affine.apply #map1(%arg4)
+      // CHECK: "test.tiling_no_dps_op"
+      // CHECK: "test.unregistered_op"
+      %extracted_slice = tensor.extract_slice %0[%arg2, %3, %4] [1, 2, 4] [1, 1, 1] : tensor<4x4x4xf32> to tensor<1x2x4xf32>
+      %5 = "test.unregistered_op"(%extracted_slice, %extracted_slice) : (tensor<1x2x4xf32>, tensor<1x2x4xf32>) -> tensor<1x2x4xf32>
+      scf.forall.in_parallel {
+        tensor.parallel_insert_slice %5 into %arg5[%arg2, %3, %4] [1, 2, 4] [1, 1, 1] : tensor<1x2x4xf32> into tensor<4x4x4xf32>
+      }
+    }
+    return %2 : tensor<4x4x4xf32>
+  }
+
+  module attributes {transform.with_named_sequence} {
+    transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
+      %op = transform.structured.match ops{["test.tiling_no_dps_op"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+      %forall = transform.structured.match ops{["scf.forall"]} in %arg0 : (!transform.any_op) -> !transform.any_op
+      %fused, %new_containing = transform.structured.fuse_into_containing_op %op into %forall : (!transform.any_op, !transform.any_op) -> (!transform.any_op, !transform.any_op)
+      transform.yield
+    }
+  }
+}
+
+// -----
+
 module {
   // CHECK-LABEL: func.func @fuse_tileable_op_through_bbarg_inout_nested
   //  CHECK-SAME:   %[[ARG0:[0-9a-z]+]]: tensor<?x?x?xf32>

mlir/test/lib/Dialect/Test/TestOpDefs.cpp

@@ -1051,6 +1051,84 @@ LogicalResult OpWithRefineTypeInterfaceOp::refineReturnTypes(
   return success();
 }
 
+//===----------------------------------------------------------------------===//
+// TilingNoDpsOp
+//===----------------------------------------------------------------------===//
+
+static Value getSlice(OpBuilder &builder, Location loc, Value source,
+                      ArrayRef<OpFoldResult> offsets,
+                      ArrayRef<OpFoldResult> sizes,
+                      ArrayRef<OpFoldResult> strides) {
+  auto staticOffsets = getConstantIntValues(offsets);
+  auto staticSizes = getConstantIntValues(sizes);
+  auto staticStrides = getConstantIntValues(strides);
+
+  auto sourceShape = cast<ShapedType>(source.getType()).getShape();
+  if (staticSizes && ArrayRef(*staticSizes) == sourceShape)
+    return source;
+
+  return {mlir::tensor::ExtractSliceOp::create(builder, loc, source, offsets,
+                                               sizes, strides)};
+}
+
+static ShapedType getSliceType(ShapedType type, ArrayRef<OpFoldResult> sizes) {
+  auto staticSizes = getConstantIntValues(sizes);
+  if (staticSizes.has_value())
+    return type.cloneWith(*staticSizes, type.getElementType());
+  return nullptr;
+}
+
+SmallVector<Range> TilingNoDpsOp::getIterationDomain(OpBuilder &builder) {
+  auto shape = cast<ShapedType>(getResult().getType()).getShape();
+  auto zero = getAsIndexOpFoldResult(getContext(), 0);
+  auto one = getAsIndexOpFoldResult(getContext(), 1);
+  return llvm::map_to_vector(shape, [&](int64_t size) {
+    return Range{.offset = zero,
+                 .size = getAsIndexOpFoldResult(getContext(), size),
+                 .stride = one};
+  });
+}
+
+SmallVector<utils::IteratorType> TilingNoDpsOp::getLoopIteratorTypes() {
+  auto tensorType = cast<ShapedType>(getResult().getType());
+  SmallVector<utils::IteratorType> types(
+      static_cast<size_t>(tensorType.getRank()), utils::IteratorType::parallel);
+  return types;
+}
+
+FailureOr<TilingResult>
+TilingNoDpsOp::getTiledImplementation(OpBuilder &builder,
+                                      ArrayRef<OpFoldResult> offsets,
+                                      ArrayRef<OpFoldResult> sizes) {
+  auto loc = getLoc();
+  auto strides = SmallVector<OpFoldResult>(
+      static_cast<size_t>(cast<ShapedType>(getOperand(0).getType()).getRank()),
+      getAsIndexOpFoldResult(getContext(), 1));
+  auto inputSlices = llvm::map_to_vector(getOperands(), [&](Value operand) {
+    return getSlice(builder, loc, operand, offsets, sizes, strides);
+  });
+  auto resultType =
+      getSliceType(cast<ShapedType>(getResult().getType()), sizes);
+  auto tiledOp = TilingNoDpsOp::create(builder, loc, TypeRange{resultType},
+                                       ValueRange(inputSlices));
+  return TilingResult{.tiledOps = {tiledOp},
+                      .tiledValues = SmallVector<Value>{tiledOp.getResult()},
+                      .generatedSlices =
+                          map_to_vector(inputSlices, [](Value val) {
+                            return val.getDefiningOp();
+                          })};
+}
+
+LogicalResult TilingNoDpsOp::getResultTilePosition(
+    OpBuilder & builder, unsigned resultNumber,
+    ArrayRef<OpFoldResult> offsets, ArrayRef<OpFoldResult> sizes,
+    SmallVector<OpFoldResult> &resultOffsets,
+    SmallVector<OpFoldResult> &resultSizes) {
+  resultOffsets.assign(offsets.begin(), offsets.end());
+  resultSizes.assign(sizes.begin(), sizes.end());
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // OpWithShapedTypeInferTypeAdaptorInterfaceOp
 //===----------------------------------------------------------------------===//

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

@@ -30,6 +30,7 @@ include "mlir/Interfaces/InferTypeOpInterface.td"
 include "mlir/Interfaces/LoopLikeInterface.td"
 include "mlir/Interfaces/MemorySlotInterfaces.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
+include "mlir/Interfaces/TilingInterface.td"
 include "mlir/Interfaces/ValueBoundsOpInterface.td"
 include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.td"
 include "mlir/Dialect/Bufferization/IR/BufferizationTypeInterfaces.td"
@@ -2887,6 +2888,20 @@ def TestLinalgFillOp :
   }];
 }
 
+//===----------------------------------------------------------------------===//
+// Test TilingInterface.
+//===----------------------------------------------------------------------===//
+
+def Test_TilingNoDpsOp : TEST_Op<"tiling_no_dps_op",
+    [Pure, DeclareOpInterfaceMethods<TilingInterface,
+      ["getIterationDomain",
+       "getLoopIteratorTypes",
+       "getResultTilePosition",
+       "getTiledImplementation"]>]> {
+  let arguments = (ins AnyRankedTensor:$lhs, AnyRankedTensor:$rhs);
+  let results = (outs AnyRankedTensor:$result);
+}
+
 //===----------------------------------------------------------------------===//
 // Test NVVM RequiresSM trait.
 //===----------------------------------------------------------------------===//