[mlir][linalg] Add a new helper hook: `hasVectorizationImpl` #110708

banach-space · 2024-10-01T17:37:21Z

The newly added hook simply returns false for Ops for which there's no
"vectorization logic" in the Linalg Vectorizer (i.e. the vectorize()
method). It's added so that the following two TD ops expose identical
level of functionality (that's not the case ATM):

transform.structured.vectorize_children_and_apply_patterns
transform.structured.vectorize

Specifically, ATM, the former works only for Linalg Ops, while the
latter works for all Ops that the vectorizer supports (*). With this change,
I am making sure that both TD will behave consistently.

Note, this shouldn't affect any of the current uses of the vectorizer.

(*) This is implemented via the vectorize() method in Vectorization.cpp.

llvmbot · 2024-10-01T17:37:58Z

@llvm/pr-subscribers-mlir-linalg

Author: Andrzej Warzyński (banach-space)

Changes

The newly added hook simply returns false for Ops for which there's no
"vectorization logic" in the Linalg Vectorizer (i.e. the vectorize
method). It's added so that the following two TD ops expose identical
level of functionality (that's not the case ATM):

transform.structured.vectorize_children_and_apply_patterns
transform.structured.vectorize

Full diff: https://github.com/llvm/llvm-project/pull/110708.diff

3 Files Affected:

(modified) mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h (+7)
(modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+6-7)
(modified) mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp (+5)

diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 48e657cca96e39..facda13eeaa7ed 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -762,6 +762,13 @@ LogicalResult copyToGPUPrivateMemory(OpBuilder &b, Value src, Value dst);
 /// memory is freed when going outside of the scope.
 LogicalResult deallocateGPUPrivateMemory(OpBuilder &, Value /*buffer*/);
 
+/// Check if this Op is vectorizable. All Linalg Op are vectorizable, as well
+/// as selected Tensor Ops. Note that this is merely a high level check and
+/// that the vectorizer also requires various additional pre-conditions to be
+/// met for it to work. These are only checked for Ops that are supported,
+/// other Ops should be rejected early.
+bool isVectorizable(Operation *);
+
 /// Emit a suitable vector form for an operation. If provided,
 /// `inputVectorSizes` are used to vectorize this operation. `inputVectorSizes`
 /// must match the rank of the iteration space of the operation and the sizes
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 46c8510f4ed514..8632e7db76353c 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -3411,11 +3411,11 @@ struct VectorizationPattern : public RewritePattern {
         flatten1DDepthwiseConv(flattenConv) {}
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
-    LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
-    if (!linalgOp)
-      return rewriter.notifyMatchFailure(op, "expected Linalg Op");
-    return vectorize(rewriter, linalgOp, /*inputVectorSizes=*/{},
-                     /*scalableVecDims=*/{}, vectorizeNDExtract,
+    if (!linalg::isVectorizable(op))
+      return rewriter.notifyMatchFailure(op,
+                                         "Unsupported Op, cannot vectorize");
+    return vectorize(rewriter, op, /*inputVectorSizes=*/{},
+                     /*inputScalableVecDims=*/{}, vectorizeNDExtract,
                      flatten1DDepthwiseConv);
   }
 
@@ -3496,8 +3496,7 @@ DiagnosedSilenceableFailure transform::VectorizeOp::apply(
 
   // TODO: Check that the correct number of vectorSizes was provided.
   for (Operation *target : targets) {
-    if (!isa<linalg::LinalgOp, tensor::PadOp, tensor::PackOp, tensor::UnPackOp>(
-            target)) {
+    if (!linalg::isVectorizable(target)) {
       return mlir::emitSilenceableFailure(target->getLoc())
              << "Unsupported Op, cannot vectorize";
     }
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index ca85f4b9b9c156..d873ca41ee86f8 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -2129,6 +2129,11 @@ static void convertAffineApply(RewriterBase &rewriter, LinalgOp linalgOp) {
   }
 }
 
+bool mlir::linalg::isVectorizable(Operation *op) {
+  return isa<linalg::LinalgOp, tensor::PadOp, tensor::PackOp, tensor::UnPackOp>(
+      op);
+}
+
 /// Emit a suitable vector form for an operation. If provided,
 /// `inputVectorSizes` are used to vectorize this operation.
 /// `inputVectorSizes` must match the rank of the iteration space of the

llvmbot · 2024-10-01T17:37:58Z

@llvm/pr-subscribers-mlir

Author: Andrzej Warzyński (banach-space)

Changes

The newly added hook simply returns false for Ops for which there's no
"vectorization logic" in the Linalg Vectorizer (i.e. the vectorize
method). It's added so that the following two TD ops expose identical
level of functionality (that's not the case ATM):

transform.structured.vectorize_children_and_apply_patterns
transform.structured.vectorize

Full diff: https://github.com/llvm/llvm-project/pull/110708.diff

3 Files Affected:

(modified) mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h (+7)
(modified) mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp (+6-7)
(modified) mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp (+5)

diff --git a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
index 48e657cca96e39..facda13eeaa7ed 100644
--- a/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
+++ b/mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
@@ -762,6 +762,13 @@ LogicalResult copyToGPUPrivateMemory(OpBuilder &b, Value src, Value dst);
 /// memory is freed when going outside of the scope.
 LogicalResult deallocateGPUPrivateMemory(OpBuilder &, Value /*buffer*/);
 
+/// Check if this Op is vectorizable. All Linalg Op are vectorizable, as well
+/// as selected Tensor Ops. Note that this is merely a high level check and
+/// that the vectorizer also requires various additional pre-conditions to be
+/// met for it to work. These are only checked for Ops that are supported,
+/// other Ops should be rejected early.
+bool isVectorizable(Operation *);
+
 /// Emit a suitable vector form for an operation. If provided,
 /// `inputVectorSizes` are used to vectorize this operation. `inputVectorSizes`
 /// must match the rank of the iteration space of the operation and the sizes
diff --git a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
index 46c8510f4ed514..8632e7db76353c 100644
--- a/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
+++ b/mlir/lib/Dialect/Linalg/TransformOps/LinalgTransformOps.cpp
@@ -3411,11 +3411,11 @@ struct VectorizationPattern : public RewritePattern {
         flatten1DDepthwiseConv(flattenConv) {}
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
-    LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
-    if (!linalgOp)
-      return rewriter.notifyMatchFailure(op, "expected Linalg Op");
-    return vectorize(rewriter, linalgOp, /*inputVectorSizes=*/{},
-                     /*scalableVecDims=*/{}, vectorizeNDExtract,
+    if (!linalg::isVectorizable(op))
+      return rewriter.notifyMatchFailure(op,
+                                         "Unsupported Op, cannot vectorize");
+    return vectorize(rewriter, op, /*inputVectorSizes=*/{},
+                     /*inputScalableVecDims=*/{}, vectorizeNDExtract,
                      flatten1DDepthwiseConv);
   }
 
@@ -3496,8 +3496,7 @@ DiagnosedSilenceableFailure transform::VectorizeOp::apply(
 
   // TODO: Check that the correct number of vectorSizes was provided.
   for (Operation *target : targets) {
-    if (!isa<linalg::LinalgOp, tensor::PadOp, tensor::PackOp, tensor::UnPackOp>(
-            target)) {
+    if (!linalg::isVectorizable(target)) {
       return mlir::emitSilenceableFailure(target->getLoc())
              << "Unsupported Op, cannot vectorize";
     }
diff --git a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
index ca85f4b9b9c156..d873ca41ee86f8 100644
--- a/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
+++ b/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
@@ -2129,6 +2129,11 @@ static void convertAffineApply(RewriterBase &rewriter, LinalgOp linalgOp) {
   }
 }
 
+bool mlir::linalg::isVectorizable(Operation *op) {
+  return isa<linalg::LinalgOp, tensor::PadOp, tensor::PackOp, tensor::UnPackOp>(
+      op);
+}
+
 /// Emit a suitable vector form for an operation. If provided,
 /// `inputVectorSizes` are used to vectorize this operation.
 /// `inputVectorSizes` must match the rank of the iteration space of the

rengolin

I was expecting this would actually check if the op has the right method, perhaps driven by implementing an interface or something. The code here is mostly moving things round, NFC almost.

banach-space · 2024-10-01T21:30:55Z

Yeah, I realise that the method name suggests something grand, but in reality is just a wrapper for an if stmt that I want to make sure is shared and re-used.

I hope that this is OK. Not strictly needed today, but I' working on a small refactor that this is blocking (as one of the TD ops bails out too early).

rengolin · 2024-10-01T21:37:22Z

I don't mind the little snippet. @ftynse ?

hanhanW · 2024-10-01T21:56:48Z

It is confusing because we have vectorizeOpPrecondition method; the real checks happen there.

llvm-project/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Lines 2091 to 2115 in ce72c76

    
           LogicalResult mlir::linalg::vectorizeOpPrecondition( 
        
               Operation *op, ArrayRef<int64_t> inputVectorSizes, 
        
               ArrayRef<bool> inputScalableVecDims, bool vectorizeNDExtract, 
        
               bool flatten1DDepthwiseConv) { 
        
             if (failed(vectorizeScalableVectorPrecondition(op, inputVectorSizes, 
        
                                                            inputScalableVecDims))) 
        
               return failure(); 
        
             return TypeSwitch<Operation *, LogicalResult>(op) 
        
                 .Case<linalg::LinalgOp>([&](auto linalgOp) { 
        
                   return vectorizeLinalgOpPrecondition(linalgOp, inputVectorSizes, 
        
                                                        vectorizeNDExtract, 
        
                                                        flatten1DDepthwiseConv); 
        
                 }) 
        
                 .Case<tensor::PadOp>([&](auto padOp) { 
        
                   return vectorizePadOpPrecondition(padOp, inputVectorSizes); 
        
                 }) 
        
                 .Case<tensor::PackOp>([&](auto packOp) { 
        
                   return vectorizePackOpPrecondition(packOp, inputVectorSizes); 
        
                 }) 
        
                 .Case<tensor::UnPackOp>([&](auto unpackOp) { 
        
                   return vectorizeUnPackOpPrecondition(unpackOp, inputVectorSizes); 
        
                 }) 
        
                 .Default([](auto) { return failure(); }); 
        
           }

hanhanW · 2024-10-01T21:58:44Z

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

@@ -762,6 +762,13 @@ LogicalResult copyToGPUPrivateMemory(OpBuilder &b, Value src, Value dst);
 /// memory is freed when going outside of the scope.
 LogicalResult deallocateGPUPrivateMemory(OpBuilder &, Value /*buffer*/);

+/// Check if this Op is vectorizable. All Linalg Op are vectorizable, as well


I think this statement is not correct. Not all the Linalg op are vectorizable. Linalg ops with non-trivial indexing maps are not vectorizable atm.

llvm-project/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Lines 1916 to 1926 in ce72c76

// TODO: the common vector shape is equal to the static loop sizes only when

// all indexing maps are projected permutations. For convs and stencils the

// logic will need to evolve.

if (!allIndexingsAreProjectedPermutation(linalgOp)) {

LDBG("precondition failed: not projected permutations\n");

return failure();

}

if (failed(reductionPreconditions(linalgOp))) {

LDBG("precondition failed: reduction preconditions\n");

return failure();

}

banach-space · 2024-10-02T07:17:29Z

I don't mind the little snippet. @ftynse ?

Sure, but then we have to repeat the same snippet twice. Also, isVectorizable is really a property of the vectorizer that should be captured/documented somewhere. That's missing today.

In case it wasn't clear from the diff, this is what's happening today (pseudo-code to capture the high-level idea):

// CURRENT logic for transform.structured.vectorize
LogicalResult process_vectorize_td(Operation *target) {
  if (!isa<linalg::LinalgOp, tensor::PadOp, tensor::PackOp, tensor::UnPackOp>(target))
    return rewriter.notifyMatchFailure(op, "Unsupported Op, cannot vectorize");

  return vectorize(target);
}

// CURRENT logic for transform.structured.vectorize_children_and_apply_patterns
LogicalResult process_vectorize_children_and_apply_patterns_td(Operation *target) {
  // NOTE - this condition is unnecessarily restrictive and is being relaxed in this PR
  if (!isa<linalg::LinalgOp>(target))
    return rewriter.notifyMatchFailure(op, "expected Linalg Op");

  return vectorize(target);
}

After this change, we will have this:

bool isVectorizable(Operation *op) {
  return isa<linalg::LinalgOp, tensor::PadOp, tensor::PackOp, tensor::UnPackOp>(
      op);
}

// NEW logic for transform.structured.vectorize
LogicalResult process_vectorize_td(Operation *target) {
  if (!isVectorizable(target))
    return mlir::emitSilenceableFailure(target->getLoc())
           << "Unsupported Op, cannot vectorize";

  return vectorize(target);
}

// NEW logic for transform.structured.vectorize_children_and_apply_patterns
LogicalResult
process_vectorize_children_and_apply_patterns_td(Operation *target) {
  if (!isVectorizable(target))
    return mlir::emitSilenceableFailure(target->getLoc())
           << "Unsupported Op, cannot vectorize";

  return vectorize(target);
}

Put differently, this PR aims to achieve two things:

healthier code re-use,
relax the artificially imposed restrictions on ransform.structured.vectorize_children_and_apply_patterns.

banach-space · 2024-10-02T07:30:35Z

It is confusing because we have vectorizeOpPrecondition method; the real checks happen there.

llvm-project/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

Lines 2091 to 2115 in ce72c76

LogicalResult mlir::linalg::vectorizeOpPrecondition(

Operation *op, ArrayRef<int64_t> inputVectorSizes,

ArrayRef<bool> inputScalableVecDims, bool vectorizeNDExtract,

bool flatten1DDepthwiseConv) {

if (failed(vectorizeScalableVectorPrecondition(op, inputVectorSizes,

inputScalableVecDims)))

return failure();

return TypeSwitch<Operation *, LogicalResult>(op)

.Case<linalg::LinalgOp>([&](auto linalgOp) {

return vectorizeLinalgOpPrecondition(linalgOp, inputVectorSizes,

vectorizeNDExtract,

flatten1DDepthwiseConv);

})

.Case<tensor::PadOp>([&](auto padOp) {

return vectorizePadOpPrecondition(padOp, inputVectorSizes);

})

.Case<tensor::PackOp>([&](auto packOp) {

return vectorizePackOpPrecondition(packOp, inputVectorSizes);

})

.Case<tensor::UnPackOp>([&](auto unpackOp) {

return vectorizeUnPackOpPrecondition(unpackOp, inputVectorSizes);

})

.Default([](auto) { return failure(); });

}

Yes, there are effectively two levels of verification (note that this PR doesn't change that). Here's pseudo-code to explain this:

// Verification 1: At the level of transform.structured.vectorize (very crude, initial filtering)
LogicalResult process_vectorize_td(Operation *target) {
  if (!isVectorizable)
    return rewriter.notifyMatchFailure(op, "Unsupported Op, cannot vectorize");

  return vectorize(target);
}

// Verification 2: At the level of linalg::vectorize (very detailed)
LogicalResult vectorize(Operation *op) {
    if (failed(vectorizeOpPrecondition(op) {
      LDBG("Vectorization pre-conditions failed\n");
      return failure();
    }
    // ...
}

Basically:

first, before doing anything expensive (and entering vectorize()), we check that the candidate Op is supported (*),
second, for every Op that's supported, we run more involved checks inside vectorize().

I can re-wire things, but would rather start with this small change :) I have a few other changes in flight that I'd like to prioritise (related to vectorising tensor.pad/tensor.pack).

(*) You are right that even some Linalg Ops don't qualify for vectorization (at least today). The idea behind isVectorizable (I should probably rename it) is to reject Ops that are known not to be supported at all, i.e. Ops outside the Linalg Dialect and selected Ops from Tensor. And for the supported Ops, as you pointed out, we still need to run the pre-condition checks and reject some of them.

Hope I didn't miss anything 😅

banach-space · 2024-10-02T12:32:44Z

Note - I've added two tests for tensor.pack to demonstrate how transform.structured.vectorize_children_and_apply_patterns is being extended.

The newly added hook simply returns `false` for Ops for which there's no "vectorization logic" in the Linalg Vectorizer (i.e. the `vectorize` method). It's added so that the following two TD ops expose identical level of functionality (that's not the case ATM): * `transform.structured.vectorize_children_and_apply_patterns` * `transform.structured.vectorize`

Add tests to demonstrate the new functionality added to `transform.structured.vectorize_children_and_apply_patterns`.

hanhanW · 2024-10-02T15:30:01Z

I see your point, thanks for the details. It's very helpful! I'm generally +1 on removing the dup code between transform ops and pattern rewriting. It reduces the probability that one is supported in transform ops but not pattern rewriting.

I think hasVectorizationImpl is a better name. What do you think?

@hanhanW

Rename according to @hanhanW's suggestion

banach-space · 2024-10-02T15:59:22Z

I think hasVectorizationImpl is a better name. What do you think?

Agreed, thank you for the suggestion, updated :) 🙏🏻

hanhanW

LGTM, thanks! Please wait for an additional approval before landing it. Others could have different comments.

dcaballe · 2024-10-02T17:09:35Z

LG. Just one comment: shouldn't we use that check within the vectorizer preconditions?

…izable` Use hasVectorizationImpl in vectorizeOpPrecondition

banach-space · 2024-10-03T18:05:41Z

@rengolin I will assume that you are OK with this and land tomorrow. Do let me know if my assumption is incorrect :)

banach-space requested review from ftynse, nicolasvasilache, dcaballe, hanhanW and rengolin as code owners October 1, 2024 17:37

llvmbot added mlir:linalg mlir labels Oct 1, 2024

rengolin reviewed Oct 1, 2024

View reviewed changes

hanhanW reviewed Oct 1, 2024

View reviewed changes

banach-space added 2 commits October 2, 2024 13:33

fixup! [mlir][linalg] Add a new helper hook - isVectorizable

778fd1c

Add tests to demonstrate the new functionality added to `transform.structured.vectorize_children_and_apply_patterns`.

banach-space force-pushed the andrzej/add_is_vectorizable branch from 5b0db64 to 778fd1c Compare October 2, 2024 12:46

fixup! fixup! [mlir][linalg] Add a new helper hook - isVectorizable

b2bae45

Rename according to @hanhanW's suggestion

hanhanW approved these changes Oct 2, 2024

View reviewed changes

banach-space changed the title ~~[mlir][linalg] Add a new helper hook - isVectorizable~~ [mlir][linalg] Add a new helper hook: hasVectorizationImpl Oct 3, 2024

fixup! fixup! fixup! [mlir][linalg] Add a new helper hook - `isVector…

3f8d9f8

…izable` Use hasVectorizationImpl in vectorizeOpPrecondition

dcaballe approved these changes Oct 4, 2024

View reviewed changes

banach-space merged commit d9d6233 into llvm:main Oct 4, 2024
8 checks passed

banach-space deleted the andrzej/add_is_vectorizable branch October 4, 2024 09:07

	// TODO: the common vector shape is equal to the static loop sizes only when
	// all indexing maps are projected permutations. For convs and stencils the
	// logic will need to evolve.
	if (!allIndexingsAreProjectedPermutation(linalgOp)) {
	LDBG("precondition failed: not projected permutations\n");
	return failure();
	}
	if (failed(reductionPreconditions(linalgOp))) {
	LDBG("precondition failed: reduction preconditions\n");
	return failure();
	}

[mlir][linalg] Add a new helper hook: hasVectorizationImpl #110708

[mlir][linalg] Add a new helper hook: hasVectorizationImpl #110708

Uh oh!

Conversation

banach-space commented Oct 1, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Oct 1, 2024

Uh oh!

llvmbot commented Oct 1, 2024

Uh oh!

rengolin left a comment

Choose a reason for hiding this comment

Uh oh!

banach-space commented Oct 1, 2024

Uh oh!

rengolin commented Oct 1, 2024

Uh oh!

hanhanW commented Oct 1, 2024

Uh oh!

hanhanW Oct 1, 2024

Choose a reason for hiding this comment

Uh oh!

banach-space commented Oct 2, 2024

Uh oh!

banach-space commented Oct 2, 2024

Uh oh!

banach-space commented Oct 2, 2024

Uh oh!

hanhanW commented Oct 2, 2024

Uh oh!

banach-space commented Oct 2, 2024

Uh oh!

hanhanW left a comment

Choose a reason for hiding this comment

Uh oh!

dcaballe commented Oct 2, 2024

Uh oh!

banach-space commented Oct 3, 2024

Uh oh!

Uh oh!

Uh oh!

[mlir][linalg] Add a new helper hook: `hasVectorizationImpl` #110708

[mlir][linalg] Add a new helper hook: `hasVectorizationImpl` #110708

banach-space commented Oct 1, 2024 •

edited

Loading