[VectorDistribution] Add vector distribution support multi-dim reduction with scalars (#18800)
Splitting https://github.com/iree-org/iree/pull/18519 into four patches.
Depends #18784
This is the second one, adding the corresponding layout analysis and
especially supporting the case where reduction is performed inside
scf.for operation.
Also, the relevant tests are added.
Since patch 2 includes changes from patch #18784, the necessary updates
from the first patch have also been included here.
---------
Signed-off-by: Bangtian Liu <liubangtian@gmail.com>
diff --git a/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel b/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel
index d9d23b2..d6cc75d 100644
--- a/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel
+++ b/compiler/src/iree/compiler/Codegen/Common/BUILD.bazel
@@ -78,6 +78,7 @@
"@llvm-project//mlir:Analysis",
"@llvm-project//mlir:DialectUtils",
"@llvm-project//mlir:IR",
+ "@llvm-project//mlir:SCFDialect",
"@llvm-project//mlir:VectorDialect",
],
)
diff --git a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt
index ee7c406..8f729de 100644
--- a/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt
+++ b/compiler/src/iree/compiler/Codegen/Common/CMakeLists.txt
@@ -58,6 +58,7 @@
LLVMSupport
MLIRAnalysis
MLIRIR
+ MLIRSCFDialect
MLIRVectorDialect
iree::compiler::Codegen::Dialect::VectorExt::IR::IREEVectorExtDialect
PUBLIC
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUDistributionPatterns.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUDistributionPatterns.cpp
index 5133d9d..0ef6e64 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUDistributionPatterns.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUDistributionPatterns.cpp
@@ -107,11 +107,11 @@
Type elementType = constant.getType().getElementType();
auto vectorType =
VectorType::get(layout.getDistributedShape(), elementType);
- Operation *distirbutedOp = rewriter.create<arith::ConstantOp>(
+ auto distributedOp = rewriter.create<arith::ConstantOp>(
constantOp.getLoc(), vectorType,
SplatElementsAttr::get(vectorType, attr.getSplatValue<Attribute>()));
replaceOpWithDistributedValues(rewriter, constantOp,
- distirbutedOp->getResult(0));
+ distributedOp->getResult(0));
return success();
}
};
@@ -536,8 +536,10 @@
SmallVector<Value> newInitArgs;
for (Value initArg : forOp.getInitArgs()) {
if (auto vectorInitArg = dyn_cast<VectorValue>(initArg)) {
- initArg =
- getDistributed(rewriter, vectorInitArg, signature[vectorInitArg]);
+ if (isNonZeroRank(vectorInitArg)) {
+ initArg =
+ getDistributed(rewriter, vectorInitArg, signature[vectorInitArg]);
+ }
}
newInitArgs.push_back(initArg);
}
@@ -582,8 +584,14 @@
SmallVector<Value> operands;
for (Value operand : yieldOp->getOperands()) {
if (auto vectorOperand = dyn_cast<VectorValue>(operand)) {
- operand = DistributionPattern::getDistributed(rewriter, vectorOperand,
- signature[vectorOperand]);
+ // Distributing the operand requires it to have a non-zero rank, meaning
+ // it must have at least one dimension. If the vector has a non-zero
+ // rank, the operand is distributed according to the provided layout
+ // signature.
+ if (isNonZeroRank(vectorOperand)) {
+ operand = DistributionPattern::getDistributed(
+ rewriter, vectorOperand, signature[vectorOperand]);
+ }
}
operands.push_back(operand);
}
@@ -606,8 +614,10 @@
for (auto [bbArg, oldInit] : llvm::zip_equal(bbArgs, oldInits)) {
Value val = bbArg;
if (auto oldVectorInit = dyn_cast<VectorValue>(oldInit)) {
- val = rewriter.create<IREE::VectorExt::ToSIMDOp>(
- oldVectorInit.getLoc(), oldVectorInit.getType(), val);
+ if (isNonZeroRank(oldVectorInit)) {
+ val = rewriter.create<IREE::VectorExt::ToSIMDOp>(
+ oldVectorInit.getLoc(), oldVectorInit.getType(), val);
+ }
}
replacements.push_back(val);
}
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUNestedLayoutDistributionPatterns.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUNestedLayoutDistributionPatterns.cpp
index e36ad99..c8b2ede 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUNestedLayoutDistributionPatterns.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUNestedLayoutDistributionPatterns.cpp
@@ -305,7 +305,9 @@
auto vectorType = VectorType::get(distShape, elementType);
VectorValue srcVector = dyn_cast<VectorValue>(broadcastOp.getSource());
- if (!srcVector) {
+ // If the srcVector is a scalar (like f32) or a rank-0 vector (like
+ // vector<f32>), we proceed with the scalar distribution branch.
+ if (!srcVector || !isNonZeroRank(srcVector)) {
// The way distribution currently works, there is no partial thread
// distribution, so a scalar is available to all threads. Scalar
// distribution is simply a broadcast from scalar to the distributed
@@ -413,16 +415,10 @@
DistributionSignature &signature,
PatternRewriter &rewriter) const override {
VectorValue srcVector = multiReduceOp.getSource();
- auto accVector = dyn_cast<VectorValue>(multiReduceOp.getAcc());
- if (!accVector) {
- return rewriter.notifyMatchFailure(
- multiReduceOp, "unimplemented: scalar accumulator distribution");
- }
- auto resVector = dyn_cast<VectorValue>(multiReduceOp.getResult());
- if (!resVector) {
- return rewriter.notifyMatchFailure(
- multiReduceOp, "unimplemented: scalar result distribution");
- }
+ Value acc = multiReduceOp.getAcc();
+ Value res = multiReduceOp.getResult();
+ auto accVector = dyn_cast<VectorValue>(acc);
+ auto resVector = dyn_cast<VectorValue>(res);
auto srcLayout = dyn_cast_or_null<NestedLayoutAttr>(signature[srcVector]);
if (!srcLayout) {
@@ -440,8 +436,14 @@
VectorValue disSrc =
getDistributed(rewriter, srcVector, signature[srcVector]);
- VectorValue disAcc =
- getDistributed(rewriter, accVector, signature[accVector]);
+
+ Value disAcc;
+ if (accVector) {
+ disAcc = getDistributed(rewriter, accVector, signature[accVector]);
+ } else {
+ // Scalars are always distributed to all threads already.
+ disAcc = multiReduceOp.getAcc();
+ }
Location loc = multiReduceOp.getLoc();
@@ -462,7 +464,16 @@
auto localReduction = rewriter.create<vector::MultiDimReductionOp>(
loc, disSrc, localInit, distributedReductionMask,
multiReduceOp.getKind());
- auto locallyReduced = dyn_cast<VectorValue>(localReduction.getResult());
+
+ VectorValue locallyReduced;
+ if (accVector) {
+ locallyReduced = dyn_cast<VectorValue>(localReduction.getResult());
+ } else {
+ // Broadcast scalar accumulator to vector.
+ VectorType vecType = VectorType::get(ArrayRef{int64_t(1)}, elemTy);
+ locallyReduced = rewriter.create<vector::BroadcastOp>(
+ loc, vecType, localReduction.getResult());
+ }
assert(locallyReduced && "result should have been a vector");
@@ -485,15 +496,30 @@
// reduction.
VectorValue unflattened = rewriter.create<vector::ShapeCastOp>(
loc, shaped, threadReduced.value());
+
+ if (!accVector) {
+ // Broadcast the scalar (e.g., f32) to a vector type (e.g., vector<f32>)
+ // because the following implementation requires the operand to be a
+ // vector.
+ disAcc = rewriter.create<vector::BroadcastOp>(loc, shaped, disAcc);
+ }
+
Value accReduction = vector::makeArithReduction(
rewriter, loc, multiReduceOp.getKind(), unflattened, disAcc);
auto accReduced = dyn_cast<VectorValue>(accReduction);
if (!accReduced) {
return failure();
}
- replaceOpWithDistributedValues(rewriter, multiReduceOp, accReduced);
- return failure();
+ if (resVector) {
+ replaceOpWithDistributedValues(rewriter, multiReduceOp, accReduced);
+ } else {
+ Value accReducedVal = rewriter.create<vector::ExtractOp>(
+ loc, accReduction, ArrayRef{int64_t(0)});
+ replaceOpWithDistributedValues(rewriter, multiReduceOp, accReducedVal);
+ }
+
+ return success();
}
FailureOr<VectorValue> doThreadReduction(RewriterBase &rewriter,
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.cpp b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.cpp
index a883180..7e927b4 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/GPUVectorDistribution.cpp
@@ -132,14 +132,16 @@
for (auto [opResult, replacement] :
llvm::zip_equal(op->getOpResults(), values)) {
// If this value is a vector type, it must be converted back to simd.
- if (isa<VectorType>(replacement.getType())) {
- auto oldResult = cast<VectorValue>(opResult);
- // Create a toSIMD op to convert the value back to the simd.
- rewriter.setInsertionPointAfterValue(oldResult);
- Value toSIMD = rewriter.create<IREE::VectorExt::ToSIMDOp>(
- oldResult.getLoc(), oldResult.getType(), replacement);
- // Add to replacements.
- replacement = toSIMD;
+ if (auto replacementType = dyn_cast<VectorType>(replacement.getType())) {
+ if (replacementType.getRank() != 0) {
+ auto oldResult = cast<VectorValue>(opResult);
+ // Create a toSIMD op to convert the value back to the simd.
+ rewriter.setInsertionPointAfterValue(oldResult);
+ Value toSIMD = rewriter.create<IREE::VectorExt::ToSIMDOp>(
+ oldResult.getLoc(), oldResult.getType(), replacement);
+ // Add to replacements.
+ replacement = toSIMD;
+ }
}
replacements.push_back(replacement);
}
diff --git a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_nested_layout_vector_distribution.mlir b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_nested_layout_vector_distribution.mlir
index f05b992..1fd7682 100644
--- a/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_nested_layout_vector_distribution.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/GPU/test/gpu_nested_layout_vector_distribution.mlir
@@ -1047,3 +1047,95 @@
// CHECK: gpu.subgroup_reduce maximumf %{{.*}} cluster(size = 2, stride = 32) : (f32) -> f32
// Accumulator reduction
// CHECK: arith.maximumf %{{.*}}, %{{.*}} : vector<1x1x1xf32>
+
+// -----
+
+#nested = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1],
+ batch_tile = [2, 2],
+ outer_tile = [1, 1],
+ thread_tile = [16, 4],
+ element_tile = [1, 4],
+
+ subgroup_strides = [1, 1],
+ thread_strides = [1, 16]
+>
+
+func.func @mfma_16x16x16_out_reduced_alldims(%arg0: vector<32x32xf32>, %arg1: f32) -> f32 {
+ %arg0l = iree_vector_ext.to_layout %arg0 to layout(#nested) : vector<32x32xf32>
+ %0 = vector.multi_reduction <maximumf>, %arg0l, %arg1 [0, 1] : vector<32x32xf32> to f32
+ return %0 : f32
+}
+
+builtin.module attributes { transform.with_named_sequence } {
+ transform.named_sequence @__transform_main(%variant_op: !transform.any_op {transform.readonly}) {
+ %top_level_func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op
+ transform.iree.test_gpu_vector_distribution %top_level_func : !transform.any_op
+ transform.yield
+ }
+}
+
+// CHECK-LABEL: func @mfma_16x16x16_out_reduced_alldims
+// Local reduction
+// CHECK: vector.multi_reduction <maximumf>, %{{.*}}, %{{.*}} [0, 1, 2, 3, 4, 5] : vector<2x2x1x1x1x4xf32> to f32
+// Global reduction
+// CHECK: gpu.subgroup_reduce maximumf %{{.*}} cluster(size = 16) : (f32) -> f32
+// CHECK-NEXT: gpu.subgroup_reduce maximumf %{{.*}} cluster(size = 4, stride = 16) : (f32) -> f32
+// Accumulator reduction
+// CHECK: arith.maximumf %{{.*}}, %{{.*}} : vector<1xf32>
+
+// -----
+
+#layout = #iree_vector_ext.nested_layout<
+ subgroup_tile = [1, 1],
+ batch_tile = [2, 2],
+ outer_tile = [1, 1],
+ thread_tile = [16, 4],
+ element_tile = [1, 4],
+
+ subgroup_strides = [1, 1],
+ thread_strides = [1, 16]
+>
+
+func.func @distribute_scf_for(%arr: memref<32x32xf16>, %a: vector<32x32xf16>) -> vector<f32> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c128 = arith.constant 128 : index
+ %cst = arith.constant dense<0.000000e+00> : vector<f32>
+ %cst_0 = arith.constant 0.0 : f16
+ %out = scf.for %i = %c0 to %c128 step %c1 iter_args(%arg0 = %cst) -> (vector<f32>) {
+ %root = vector.transfer_read %arr[%c0, %c0], %cst_0 {in_bounds = [true, true]} : memref<32x32xf16>, vector<32x32xf16>
+ %rootl = iree_vector_ext.to_layout %root to layout(#layout) : vector<32x32xf16>
+ %b = arith.addf %rootl, %a : vector<32x32xf16>
+ %c = arith.extf %b : vector<32x32xf16> to vector<32x32xf32>
+ %init = vector.extractelement %arg0[] : vector<f32>
+ %root_red = vector.multi_reduction<add>, %c, %init [0, 1] : vector<32x32xf32> to f32
+ %d = vector.broadcast %root_red : f32 to vector<f32>
+ scf.yield %d : vector<f32>
+ }
+ return %out : vector<f32>
+}
+
+builtin.module attributes { transform.with_named_sequence } {
+ transform.named_sequence @__transform_main(%variant_op: !transform.any_op {transform.readonly}) {
+ %top_level_func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op
+ transform.iree.test_gpu_vector_distribution %top_level_func : !transform.any_op
+ transform.yield
+ }
+}
+
+// CHECK-LABEL: func @distribute_scf_for
+// CHECK: %[[ROOT:.*]] = arith.constant dense<0.000000e+00> : vector<f32>
+// CHECK: iter_args(%[[ARG0:.*]] = %[[ROOT]]) -> (vector<f32>)
+// CHECK: %[[A:.*]] = iree_vector_ext.to_simt %{{.*}} : vector<32x32xf16> -> vector<2x2x1x1x1x4xf16>
+// CHECK: %[[B:.*]] = arith.addf %{{.*}}, %[[A]]
+// CHECK: %[[C:.*]] = arith.extf %[[B]]
+// CHECK-NEXT: %[[D:.*]] = vector.extractelement %[[ARG0]][] : vector<f32>
+// Local reduction
+// CHECK: vector.multi_reduction <add>, %[[C]], %{{.*}} [0, 1, 2, 3, 4, 5] : vector<2x2x1x1x1x4xf32> to f32
+// Global reduction
+// CHECK: gpu.subgroup_reduce add %{{.*}} cluster(size = 16) : (f32) -> f32
+// CHECK-NEXT: gpu.subgroup_reduce add %{{.*}} cluster(size = 4, stride = 16) : (f32) -> f32
+// Accumulator reduction
+// CHECK: vector.broadcast %[[D]] : f32 to vector<1xf32>
+// CHECK: arith.addf %{{.*}}, %{{.*}} : vector<1xf32>
diff --git a/compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp b/compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp
index deabc58..28b75d1 100644
--- a/compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp
+++ b/compiler/src/iree/compiler/Codegen/Common/VectorLayoutAnalysis.cpp
@@ -13,6 +13,7 @@
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/Diagnostics.h"
@@ -135,6 +136,9 @@
RegionBranchPoint branchPoint,
MutableArrayRef<OpOperand> operands);
+ void visitRegionBranchTerminatorOpInterface(RegionBranchOpInterface branch,
+ RegionBranchPoint branchPoint);
+
DistributionLayout *getLatticeElement(Value val);
MLIRContext *ctx;
@@ -662,6 +666,9 @@
ArrayRef<DistributionLayout *> operandLattices,
ArrayRef<const DistributionLayout *> resultLattices,
std::function<void(DistributionLayout *, ChangeResult)> update) {
+ if (resultLattices.empty()) {
+ return;
+ }
// Reductions should always propagate value layout to result. Result can
// enforce it's layout on init.
const DistributionLayout *result = resultLattices[0];
@@ -727,9 +734,12 @@
auto resultShape = broadcast.getResultVectorType().getShape();
auto inputType = broadcast.getSourceType();
- assert(isa<VectorType>(inputType) &&
- "Scalar broadcast not supported for now.");
- auto inputShape = cast<VectorType>(inputType).getShape();
+
+ VectorType inputVectorType = dyn_cast<VectorType>(inputType);
+ if (!inputVectorType)
+ return;
+
+ auto inputShape = inputVectorType.getShape();
SmallVector<bool> reductionMask(resultShape.size(), false);
// Set the trailing dimensions to be reduced.
@@ -994,6 +1004,9 @@
if (auto branch = dyn_cast<RegionBranchOpInterface>(op)) {
visitRegionSuccessors(branch, RegionBranchPoint::parent(),
branch->getOpOperands());
+
+ // Handle the propagation from scf.for to yield op.
+ visitRegionBranchTerminatorOpInterface(branch, RegionBranchPoint::parent());
return;
}
@@ -1086,6 +1099,43 @@
}
}
+void EnforceLayout::visitRegionBranchTerminatorOpInterface(
+ RegionBranchOpInterface branch, RegionBranchPoint branchPoint) {
+ SmallVector<RegionSuccessor> successors;
+ branch.getSuccessorRegions(branchPoint, successors);
+ if (!branch.hasLoop())
+ return;
+ SmallVector<DistributionLayout *> resultLattices;
+ for (Value result : branch->getResults()) {
+ DistributionLayout *resultLattice = getLatticeElement(result);
+ if (resultLattice->isUninitialized())
+ continue;
+ resultLattices.push_back(resultLattice);
+ }
+
+ // We do not support multiple results yet.
+ if (resultLattices.size() != 1)
+ return;
+
+ for (RegionSuccessor successor : successors) {
+ if (Region *succ = successor.getSuccessor()) {
+ Operation *terminator = succ->back().getTerminator();
+ if (scf::YieldOp yieldOp = dyn_cast<scf::YieldOp>(terminator)) {
+ for (Value operand : yieldOp.getOperands()) {
+ if (!isa<VectorType>(operand.getType())) {
+ continue;
+ }
+ DistributionLayout *forwardLattice = getLatticeElement(operand);
+ ChangeResult changed = forwardLattice->resolve(resultLattices[0]);
+ propagateIfChanged(forwardLattice, changed);
+ }
+ }
+ }
+ }
+
+ return;
+}
+
DistributionLayout *EnforceLayout::getLatticeElement(Value val) {
// Add dependency of operation on the analysis state.
assert(isa<VectorType>(val.getType()) && "Lattice value should be a vector");
diff --git a/compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir b/compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir
index 28e8ab0..6533a09 100644
--- a/compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir
+++ b/compiler/src/iree/compiler/Codegen/Common/test/vector_layout_analysis.mlir
@@ -562,3 +562,36 @@
transform.yield
}
}
+
+// -----
+
+#layout = #iree_vector_ext.layout<<[VECTORY], [16]>, <[BATCHY, VECTORX], [2, 8]>>
+
+// Propagate and enforce through scf.for
+builtin.module attributes { transform.with_named_sequence } {
+ func.func @scffor(%arr: memref<16x16xf16>, %arr2: memref<16xf16>, %a: vector<16xf16>, %b: vector<16xf16>) -> vector<f16> {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c1024 = arith.constant 1024 : index
+ %cst = arith.constant dense<0.000000e+00> : vector<f16>
+ %cst_0 = arith.constant 0.0 : f16
+
+ %out = scf.for %iv = %c0 to %c1024 step %c1 iter_args(%arg1 = %cst) -> (vector<f16>) {
+ %root = vector.transfer_read %arr[%c0, %c0], %cst_0 {in_bounds = [true, true]} : memref<16x16xf16>, vector<16x16xf16>
+ // expected-remark @above {{layout of result #0 is #iree_vector_ext.layout<<[ VECTORY], [16]>, <[ BATCHY, VECTORX], [2, 8]>>}}
+ %rootl = iree_vector_ext.to_layout %root to layout(#layout) : vector<16x16xf16>
+ %init = vector.extractelement %arg1[] : vector<f16>
+ %root_red = vector.multi_reduction<add>, %rootl, %init [0, 1] : vector<16x16xf16> to f16
+ %c = vector.broadcast %root_red : f16 to vector<f16>
+ scf.yield %c : vector<f16>
+ }
+
+ func.return %out : vector<f16>
+ }
+
+ transform.named_sequence @__transform_main(%variant_op: !transform.any_op {transform.readonly}) {
+ %top_level_func = transform.structured.match ops{["func.func"]} in %variant_op : (!transform.any_op) -> !transform.any_op
+ transform.iree.test_vector_layout_analysis %top_level_func : !transform.any_op
+ transform.yield
+ }
+}
diff --git a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp
index 5eb4519..e996aba 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp
+++ b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.cpp
@@ -212,6 +212,10 @@
return computeFn;
}
+bool isNonZeroRank(TypedValue<VectorType> val) {
+ return val.getType().getRank() != 0;
+}
+
//===----------------------------------------------------------------------===//
// GPU workgroup memory
//===----------------------------------------------------------------------===//
diff --git a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h
index cdbc297..4e7c108 100644
--- a/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h
+++ b/compiler/src/iree/compiler/Codegen/Utils/GPUUtils.h
@@ -105,6 +105,9 @@
FailureOr<scf::SCFTileSizeComputationFunction>
getGPUScfTileSizeComputeFn(mlir::FunctionOpInterface funcOp, int tilingLevel);
+/// Returns true iff the rank of the input value 'val' is non-zero.
+bool isNonZeroRank(TypedValue<VectorType> val);
+
//===----------------------------------------------------------------------===//
// GPU workgroup memory
//===----------------------------------------------------------------------===//
