df71000d7d
This commit adds a pattern to lower linalg.generic for reduction to spv.GroupNonUniform* ops. Right now this only supports integer reduction on 1-D input memref. Shader entry point ABI is queried to make sure that the input memref's shape matches the local workgroup's invocation configuration. This makes sure that the workload fits in one local workgroup so that we can leverage SPIR-V group non-uniform operations. linglg.generic is a structured op that preserves the right level of information. It is easier to recognize reduction at this level than performing analysis on loops. This commit also exposes `getElementPtr` in SPIRVLowering.h given that it's a generally useful utility function. Differential Revision: https://reviews.llvm.org/D73437
107 lines
4.0 KiB
C++
107 lines
4.0 KiB
C++
//===- Utils.cpp - Utilities to support the Linalg dialect ----------------===//
|
|
//
|
|
// 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 utilities for the Linalg dialect.
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Linalg/Utils/Utils.h"
|
|
#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
|
|
#include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
|
|
#include "mlir/Dialect/LoopOps/LoopOps.h"
|
|
#include "mlir/Dialect/StandardOps/Ops.h"
|
|
#include "mlir/EDSC/Helpers.h"
|
|
#include "mlir/IR/AffineExpr.h"
|
|
#include "mlir/IR/AffineMap.h"
|
|
#include "mlir/IR/Matchers.h"
|
|
#include "mlir/IR/OpImplementation.h"
|
|
#include "mlir/Pass/Pass.h"
|
|
#include "mlir/Support/STLExtras.h"
|
|
#include "mlir/Transforms/FoldUtils.h"
|
|
|
|
using namespace mlir;
|
|
using namespace mlir::edsc;
|
|
using namespace mlir::edsc::intrinsics;
|
|
using namespace mlir::linalg;
|
|
using namespace mlir::loop;
|
|
|
|
Optional<RegionMatcher::BinaryOpKind>
|
|
RegionMatcher::matchAsScalarBinaryOp(GenericOp op) {
|
|
auto ®ion = op.region();
|
|
if (!has_single_element(region))
|
|
return llvm::None;
|
|
|
|
Block &block = region.front();
|
|
if (block.getNumArguments() != 2 ||
|
|
!block.getArgument(0).getType().isIntOrFloat() ||
|
|
!block.getArgument(1).getType().isIntOrFloat())
|
|
return llvm::None;
|
|
|
|
auto &ops = block.getOperations();
|
|
if (!has_single_element(block.without_terminator()))
|
|
return llvm::None;
|
|
|
|
using mlir::matchers::m_Val;
|
|
auto a = m_Val(block.getArgument(0));
|
|
auto b = m_Val(block.getArgument(1));
|
|
|
|
auto addPattern = m_Op<linalg::YieldOp>(m_Op<AddIOp>(a, b));
|
|
if (addPattern.match(&ops.back()))
|
|
return BinaryOpKind::IAdd;
|
|
|
|
return llvm::None;
|
|
}
|
|
|
|
static Value emitOrFoldComposedAffineApply(OpBuilder &b, Location loc,
|
|
AffineMap map,
|
|
ArrayRef<Value> operandsRef,
|
|
OperationFolder *folder) {
|
|
SmallVector<Value, 4> operands(operandsRef.begin(), operandsRef.end());
|
|
fullyComposeAffineMapAndOperands(&map, &operands);
|
|
canonicalizeMapAndOperands(&map, &operands);
|
|
return folder ? folder->create<AffineApplyOp>(b, loc, map, operands)
|
|
: b.create<AffineApplyOp>(loc, map, operands);
|
|
}
|
|
|
|
SmallVector<Value, 4> mlir::linalg::applyMapToValues(OpBuilder &b, Location loc,
|
|
AffineMap map,
|
|
ArrayRef<Value> values,
|
|
OperationFolder *folder) {
|
|
SmallVector<Value, 4> res;
|
|
res.reserve(map.getNumResults());
|
|
unsigned numDims = map.getNumDims();
|
|
// For each `expr` in `map`, applies the `expr` to the values extracted from
|
|
// ranges. If the resulting application can be folded into a Value, the
|
|
// folding occurs eagerly. Otherwise, an affine.apply operation is emitted.
|
|
for (auto expr : map.getResults()) {
|
|
AffineMap map = AffineMap::get(numDims, 0, expr);
|
|
res.push_back(emitOrFoldComposedAffineApply(b, loc, map, values, folder));
|
|
}
|
|
return res;
|
|
}
|
|
|
|
/// Returns all the operands of `linalgOp` that are not views.
|
|
/// Asserts that these operands are value types to allow transformations like
|
|
/// tiling to just use the values when cloning `linalgOp`.
|
|
SmallVector<Value, 4>
|
|
mlir::linalg::getAssumedNonViewOperands(LinalgOp linalgOp) {
|
|
auto *op = linalgOp.getOperation();
|
|
unsigned numViews = linalgOp.getNumInputsAndOutputs();
|
|
unsigned nOperands = op->getNumOperands() - numViews;
|
|
SmallVector<Value, 4> res;
|
|
res.reserve(nOperands);
|
|
for (unsigned i = 0; i < nOperands; ++i) {
|
|
res.push_back(op->getOperand(numViews + i));
|
|
auto t = res.back().getType();
|
|
(void)t;
|
|
assert((t.isIntOrIndexOrFloat() || t.isa<VectorType>()) &&
|
|
"expected scalar or vector type");
|
|
}
|
|
return res;
|
|
}
|