[mlir][linalg] Add a pattern to split linalg.pad_tensor ops

[antiagainst]

Summary

This commit adds a pattern to wrap a linalg.pad_tensor op with
an scf.if op to separate the cases where we don't need padding
(all pad sizes are actually zeros) and where we indeed need
padding.

This pattern is meant to handle padding inside tiled loops.
Under such cases the padding sizes typically depend on the
loop induction variables. Splitting them would allow treating
perfect tiles and edge tiles separately.

Depends On D117017

[antiagainst]

Rebase

[nicolasvasilache]

mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td

19

This is very limited in terms of control and heuristics.
Should this just be a test pass?

20

Can we improve the terminology?

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

18

It is now called tensor.pad

mlir/lib/Dialect/Tensor/Transforms/SplitPadding.cpp

77

trivial braces

[antiagainst]

Address comments

mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td

19

Good point. Converted into an uber pass for testing various tensor patterns.

[antiagainst]

Fix BUILD.bazel

mlir/include/mlir/Dialect/Tensor/Transforms/Passes.h

 #define MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
 /// Creates an instance of `tensor` dialect bufferization pass.
 std::unique_ptr<Pass> createTensorBufferizePass();
 
+/// Creates a pass to wrap tensor.pad ops with scf.if ops to allow handle
+/// padding-elided and padding-needed cases separately.
+std::unique_ptr<Pass> createTensorSplitPaddingPass();
+
 //===----------------------------------------------------------------------===//
 // Registration
 //===----------------------------------------------------------------------===//
 
 namespace tensor {
 /// Generate the code for registering passes.
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Tensor/Transforms/Passes.h.inc"
 } // namespace tensor
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES_H_

mlir/include/mlir/Dialect/Tensor/Transforms/Passes.td

 
 include "mlir/Pass/PassBase.td"
 
 def TensorBufferize : Pass<"tensor-bufferize", "FuncOp"> {
   let summary = "Bufferize the `tensor` dialect";
   let constructor = "mlir::createTensorBufferizePass()";
 }
 
+def TensorSplitPadding : Pass<"tensor-split-padding", "FuncOp"> {
+  let summary = "Split `tensor.pad` op into padding-unnecessary and "
+                "padding-needed cases";
+  let description = [{
+     This pass creates scf.if ops to wrap tensor.pad ops to allow handle
+     padding-elided and padding-needed cases separately.
+  }];
+  let constructor = "mlir::createTensorSplitPaddingPass()";
+}
+
 #endif // MLIR_DIALECT_TENSOR_TRANSFORMS_PASSES

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

+//===- Transforms.h - Tensor Transformation Patterns ------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMS_H
+#define MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMS_H
+
+#include "mlir/IR/PatternMatch.h"
+
+namespace mlir {
+namespace tensor {
+
+/// Populates `patterns` with patterns that split linalg.pad_tensor ops by
+/// creating scf.if ops to wrap linalg.pad_tensor ops and handle
+/// padding-unncessary and padding-needed cases separately.
+void populateSplitPaddingPatterns(RewritePatternSet &patterns,
+                                  PatternBenefit baseBenefit = 1);
+
+} // namespace tensor
+} // namespace mlir
+
+#endif // MLIR_DIALECT_TENSOR_TRANSFORMS_TRANSFORMS_H

mlir/lib/Dialect/Tensor/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   Bufferize.cpp
+  SplitPadding.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Tensor/Transforms
 
   DEPENDS
   MLIRTensorTransformsIncGen
 
   LINK_LIBS PUBLIC

mlir/lib/Dialect/Tensor/Transforms/SplitPadding.cpp

+//===- SplitPadding.cpp - Splitting tensor.pad Op -------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements patterns and passes for creating scf.if ops to wrap
+// tensor.pad ops to allow handle padding-elided and padding-needed cases
+// separately.
+//
+//===----------------------------------------------------------------------===//
+
+#include "PassDetail.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tensor/Transforms/Passes.h"
+#include "mlir/Dialect/Tensor/Transforms/Transforms.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "mlir-tensor-split-padding"
+
+using namespace mlir;
+
+/// Returns true if the the given `attrOrValue` is a constant zero.
+static bool isZero(OpFoldResult attrOrValue) {
+  if (Optional<int64_t> val = getConstantIntValue(attrOrValue))
+    return val.getValue() == 0;
+  return false;
+}
+
+/// Gets the given `attrOrValue` as a Value by creating constant ops for
+/// attributes.
+static Value getAsValue(OpFoldResult attrOrValue, OpBuilder &builder,
+                        Location loc) {
+  if (Value val = attrOrValue.dyn_cast<Value>())
+    return val;
+  auto attr = attrOrValue.get<Attribute>().cast<IntegerAttr>();
+  return builder.create<arith::ConstantIndexOp>(loc, attr.getInt());
+}
+
+namespace {
+
+/// Splits a tensor.pad op by wrapping it in a scf.if op to handle
+/// padding-unnecessary and padding-needed cases.
+struct SplitPadding final : public OpRewritePattern<tensor::PadOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::PadOp padOp,
+                                PatternRewriter &rewriter) const override {
+    // Avoid infinitely applying this pattern.
+    if (padOp->getParentOfType<scf::IfOp>())
+      return failure();
+
+    // If all padding sizes are zero, we don't need to do anything.
+    SmallVector<OpFoldResult> lowPads = padOp.getMixedLowPad();
+    SmallVector<OpFoldResult> highPads = padOp.getMixedHighPad();
+    if (llvm::all_of(lowPads, isZero) && llvm::all_of(highPads, isZero))
+      return failure();
+
+    // Build the condition for the scf.if op: all pad sizes are zero.
+    Location loc = padOp.getLoc();
+    Value cstZero = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+    SmallVector<Value> eqZeroCmpVals;
+    for (OpFoldResult pad : llvm::concat<OpFoldResult>(lowPads, highPads)) {
+      if (!isZero(pad)) {
+        eqZeroCmpVals.push_back(rewriter.create<arith::CmpIOp>(
+            loc, arith::CmpIPredicate::eq, getAsValue(pad, rewriter, loc),
+            cstZero));
+      }
+    }
+    Value ifCond = eqZeroCmpVals.front();
+    for (Value cmp : llvm::makeArrayRef(eqZeroCmpVals).drop_front()) {
+      ifCond = rewriter.create<arith::AndIOp>(loc, ifCond, cmp);
+    }
+
+    // Build the scf.if op itself. For the "then" branch, we can elide the
+    // padding. For the "else" branch, we retain the clone op.
+    auto thenBuilder = [&padOp](OpBuilder &builder, Location loc) {
+      builder.create<scf::YieldOp>(loc, padOp.source());
+    };
+    auto elseBuilder = [&padOp](OpBuilder &builder, Location loc) {
+      Operation *newOp = builder.clone(*padOp);
+      builder.create<scf::YieldOp>(loc, newOp->getResults());
+    };
+    rewriter.replaceOpWithNewOp<scf::IfOp>(padOp, padOp.getType(), ifCond,
+                                           thenBuilder, elseBuilder);
+    return success();
+  }
+};
+
+struct TensorSplitPaddingPass final
+    : public TensorSplitPaddingBase<TensorSplitPaddingPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<arith::ArithmeticDialect, scf::SCFDialect>();
+  }
+
+  void runOnOperation() override {
+    FuncOp fn = getOperation();
+    RewritePatternSet patterns(&getContext());
+    tensor::populateSplitPaddingPatterns(patterns);
+    if (failed(applyPatternsAndFoldGreedily(fn, std::move(patterns))))
+      return signalPassFailure();
+  }
+};
+
+} // namespace
+
+void tensor::populateSplitPaddingPatterns(RewritePatternSet &patterns,
+                                          PatternBenefit baseBenefit) {
+  patterns.add<SplitPadding>(patterns.getContext(), baseBenefit);
+}
+
+std::unique_ptr<Pass> mlir::createTensorSplitPaddingPass() {
+  return std::make_unique<TensorSplitPaddingPass>();
+}

mlir/test/Dialect/Tensor/split-padding.mlir

+// RUN: mlir-opt -split-input-file -tensor-split-padding %s | FileCheck %s
+
+// CHECK-LABEL: func @pad_all_zero_sizes
+func @pad_all_zero_sizes(%input: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {
+  %f0 = arith.constant 0.0 : f32
+  %c0 = arith.constant 0 : index
+  %0 = tensor.pad %input low[0, %c0, 0] high[%c0, 0, 0] {
+  ^bb0(%dim0: index, %dim1: index, %dim2: index):
+    tensor.yield %f0 : f32
+  } : tensor<?x?x?xf32> to tensor<?x?x?xf32>
+  return %0 : tensor<?x?x?xf32>
+}
+
+// CHECK-NOT: scf.if
+//     CHECK: tensor.pad
+
+// -----
+
+// CHECK-LABEL: func @pad_non_zero_sizes
+//  CHECK-SAME: (%[[INPUT:.+]]: tensor<?x?x8xf32>, %[[LOW0:.+]]: index, %[[HIGH1:.+]]: index)
+func @pad_non_zero_sizes(%input: tensor<?x?x8xf32>, %low0: index, %high1: index) -> tensor<?x?x8xf32> {
+  %f0 = arith.constant 0.0 : f32
+  %0 = tensor.pad %input low[%low0, 0, 0] high[0, %high1, 0] {
+  ^bb0(%dim0: index, %dim1: index, %dim2: index):
+    tensor.yield %f0 : f32
+  } : tensor<?x?x8xf32> to tensor<?x?x8xf32>
+  return %0 : tensor<?x?x8xf32>
+}
+
+// CHECK: %[[F0:.+]] = arith.constant 0.000000e+00 : f32
+// CHECK: %[[C0:.+]] = arith.constant 0 : index
+// CHECK: %[[EQ0:.+]] = arith.cmpi eq, %[[LOW0]], %[[C0]] : index
+// CHECK: %[[EQ1:.+]] = arith.cmpi eq, %[[HIGH1]], %[[C0]] : index
+// CHECK: %[[AND:.+]] = arith.andi %[[EQ0]], %[[EQ1]] : i1
+// CHECK: %[[IF:.+]] = scf.if %[[AND]] -> (tensor<?x?x8xf32>) {
+// CHECK:   scf.yield %[[INPUT]] : tensor<?x?x8xf32>
+// CHECK: } else {
+// CHECK:   %[[PAD:.+]] = tensor.pad %[[INPUT]] low[%[[LOW0]], 0, 0] high[0, %[[HIGH1]], 0]  {
+// CHECK:   ^bb0(%{{.+}}: index, %{{.+}}: index, %{{.+}}: index):
+// CHECK:     tensor.yield %[[F0]] : f32
+// CHECK:   } : tensor<?x?x8xf32> to tensor<?x?x8xf32>
+// CHECK:   scf.yield %[[PAD]] : tensor<?x?x8xf32>
+// CHECK: }
+// CHECK: return %[[IF]] : tensor<?x?x8xf32>

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

srcs = glob(
         [
             "lib/Dialect/Tensor/Transforms/*.cpp",
             "lib/Dialect/Tensor/Transforms/*.h",
         ],
     ),
     hdrs = [
         "include/mlir/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.h",
         "include/mlir/Dialect/Tensor/Transforms/Passes.h",
+        "include/mlir/Dialect/Tensor/Transforms/Transforms.h",
     ],
     includes = ["include"],
     deps = [
         ":ArithmeticDialect",
         ":Async",
         ":BufferizationDialect",
         ":BufferizationTransforms",
+        ":DialectUtils",
         ":IR",
         ":MemRefDialect",
         ":ParallelLoopMapperAttrGen",
         ":Pass",
         ":SCFDialect",
         ":StandardOps",
         ":Support",
         ":TensorDialect",