228 lines
8.8 KiB
C++
228 lines
8.8 KiB
C++
//===- StructuredOpsUtils.cpp - Utilities used by structured ops ----------===//
|
|
//
|
|
// 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
|
|
//
|
|
//===----------------------------------------------------------------------===//
|
|
|
|
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
|
|
#include "mlir/IR/AffineMap.h"
|
|
#include "mlir/IR/Builders.h"
|
|
#include "mlir/IR/BuiltinAttributes.h"
|
|
#include "mlir/IR/IRMapping.h"
|
|
#include "llvm/ADT/StringSet.h"
|
|
|
|
#include "mlir/Dialect/Utils/DialectUtilsEnums.cpp.inc"
|
|
|
|
using namespace mlir;
|
|
|
|
bool mlir::isRowMajorMatmul(ArrayAttr indexingMaps) {
|
|
if (indexingMaps.size() != 3)
|
|
return false;
|
|
|
|
AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
|
|
AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
|
|
AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
|
|
|
|
if (map0.getNumResults() != 2 || map1.getNumResults() != 2 ||
|
|
map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
|
|
map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
|
|
return false;
|
|
}
|
|
|
|
// Extract dimensions for MxK * KxN -> MxN
|
|
AffineExpr m = map2.getResult(0);
|
|
AffineExpr n = map2.getResult(1);
|
|
AffineExpr k = map0.getResult(1);
|
|
auto *context = indexingMaps.getContext();
|
|
auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {m, k}, context));
|
|
auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {k, n}, context));
|
|
auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {m, n}, context));
|
|
auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
|
|
return indexingMaps == maps;
|
|
}
|
|
|
|
bool mlir::isColumnMajorMatmul(ArrayAttr indexingMaps) {
|
|
if (indexingMaps.size() != 3)
|
|
return false;
|
|
|
|
AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
|
|
AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
|
|
AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
|
|
|
|
if (map0.getNumResults() != 2 || map1.getNumResults() != 2 ||
|
|
map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
|
|
map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
|
|
return false;
|
|
}
|
|
|
|
// Extract dimensions for KxM * NxK -> NxM
|
|
AffineExpr n = map2.getResult(0);
|
|
AffineExpr m = map2.getResult(1);
|
|
AffineExpr k = map0.getResult(0);
|
|
auto *context = indexingMaps.getContext();
|
|
auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {k, m}, context));
|
|
auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {n, k}, context));
|
|
auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {n, m}, context));
|
|
auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
|
|
return indexingMaps == maps;
|
|
}
|
|
|
|
bool mlir::isRowMajorBatchMatmul(ArrayAttr indexingMaps) {
|
|
if (indexingMaps.size() != 3)
|
|
return false;
|
|
|
|
AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
|
|
AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
|
|
AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
|
|
|
|
if (map0.getNumResults() != 3 || map1.getNumResults() != 3 ||
|
|
map2.getNumResults() != 3 || map0.getNumInputs() != 4 ||
|
|
map1.getNumInputs() != 4 || map2.getNumInputs() != 4) {
|
|
return false;
|
|
}
|
|
|
|
// Extract dimensions for BxMxK * BxKxN -> BxMxN
|
|
AffineExpr b = map2.getResult(0);
|
|
AffineExpr m = map2.getResult(1);
|
|
AffineExpr n = map2.getResult(2);
|
|
AffineExpr k = map0.getResult(2);
|
|
auto *context = indexingMaps.getContext();
|
|
auto mapA = AffineMapAttr::get(AffineMap::get(4, 0, {b, m, k}, context));
|
|
auto mapB = AffineMapAttr::get(AffineMap::get(4, 0, {b, k, n}, context));
|
|
auto mapC = AffineMapAttr::get(AffineMap::get(4, 0, {b, m, n}, context));
|
|
auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
|
|
return indexingMaps == maps;
|
|
}
|
|
|
|
bool mlir::isVecmat(ArrayAttr indexingMaps) {
|
|
if (indexingMaps.size() != 3)
|
|
return false;
|
|
AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
|
|
AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
|
|
AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
|
|
|
|
if (map0.getNumResults() != 1 || map1.getNumResults() != 2 ||
|
|
map2.getNumResults() != 1 || map0.getNumInputs() != 2 ||
|
|
map1.getNumInputs() != 2 || map2.getNumInputs() != 2) {
|
|
return false;
|
|
}
|
|
|
|
// Extract dimensions for K * KxN -> N
|
|
AffineExpr k = map0.getResult(0);
|
|
AffineExpr n = map2.getResult(0);
|
|
auto *context = indexingMaps.getContext();
|
|
auto mapA = AffineMapAttr::get(AffineMap::get(2, 0, {k}, context));
|
|
auto mapB = AffineMapAttr::get(AffineMap::get(2, 0, {k, n}, context));
|
|
auto mapC = AffineMapAttr::get(AffineMap::get(2, 0, {n}, context));
|
|
auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
|
|
return indexingMaps == maps;
|
|
}
|
|
|
|
bool mlir::isBatchVecmat(ArrayAttr indexingMaps) {
|
|
if (indexingMaps.size() != 3)
|
|
return false;
|
|
AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
|
|
AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
|
|
AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
|
|
|
|
if (map0.getNumResults() != 2 || map1.getNumResults() != 3 ||
|
|
map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
|
|
map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
|
|
return false;
|
|
}
|
|
|
|
// Extract dimensions for B*K * B*K*N -> B*N
|
|
AffineExpr b = map0.getResult(0);
|
|
AffineExpr k = map0.getResult(1);
|
|
AffineExpr n = map2.getResult(1);
|
|
auto *context = indexingMaps.getContext();
|
|
auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {b, k}, context));
|
|
auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {b, k, n}, context));
|
|
auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {b, n}, context));
|
|
auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
|
|
return indexingMaps == maps;
|
|
}
|
|
|
|
bool mlir::isMatvec(ArrayAttr indexingMaps) {
|
|
if (indexingMaps.size() != 3)
|
|
return false;
|
|
AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
|
|
AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
|
|
AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
|
|
|
|
if (map0.getNumResults() != 2 || map1.getNumResults() != 1 ||
|
|
map2.getNumResults() != 1 || map0.getNumInputs() != 2 ||
|
|
map1.getNumInputs() != 2 || map2.getNumInputs() != 2) {
|
|
return false;
|
|
}
|
|
|
|
// Extract dimensions for N*K * K -> N
|
|
AffineExpr k = map1.getResult(0);
|
|
AffineExpr n = map2.getResult(0);
|
|
auto *context = indexingMaps.getContext();
|
|
auto mapA = AffineMapAttr::get(AffineMap::get(2, 0, {n, k}, context));
|
|
auto mapB = AffineMapAttr::get(AffineMap::get(2, 0, {k}, context));
|
|
auto mapC = AffineMapAttr::get(AffineMap::get(2, 0, {n}, context));
|
|
auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
|
|
return indexingMaps == maps;
|
|
}
|
|
|
|
bool mlir::isBatchMatvec(ArrayAttr indexingMaps) {
|
|
if (indexingMaps.size() != 3)
|
|
return false;
|
|
AffineMap map0 = cast<AffineMapAttr>(indexingMaps[0]).getValue();
|
|
AffineMap map1 = cast<AffineMapAttr>(indexingMaps[1]).getValue();
|
|
AffineMap map2 = cast<AffineMapAttr>(indexingMaps[2]).getValue();
|
|
|
|
if (map0.getNumResults() != 3 || map1.getNumResults() != 2 ||
|
|
map2.getNumResults() != 2 || map0.getNumInputs() != 3 ||
|
|
map1.getNumInputs() != 3 || map2.getNumInputs() != 3) {
|
|
return false;
|
|
}
|
|
|
|
// Extract dimensions for B*N*K * B*K -> B*N
|
|
AffineExpr b = map0.getResult(0);
|
|
AffineExpr k = map1.getResult(1);
|
|
AffineExpr n = map2.getResult(1);
|
|
auto *context = indexingMaps.getContext();
|
|
auto mapA = AffineMapAttr::get(AffineMap::get(3, 0, {b, n, k}, context));
|
|
auto mapB = AffineMapAttr::get(AffineMap::get(3, 0, {b, k}, context));
|
|
auto mapC = AffineMapAttr::get(AffineMap::get(3, 0, {b, n}, context));
|
|
auto maps = ArrayAttr::get(context, {mapA, mapB, mapC});
|
|
return indexingMaps == maps;
|
|
}
|
|
|
|
Operation *mlir::clone(OpBuilder &b, Operation *op, TypeRange newResultTypes,
|
|
ValueRange newOperands) {
|
|
IRMapping bvm;
|
|
OperationState state(op->getLoc(), op->getName(), newOperands, newResultTypes,
|
|
op->getAttrs());
|
|
for (Region &r : op->getRegions())
|
|
r.cloneInto(state.addRegion(), bvm);
|
|
return b.create(state);
|
|
}
|
|
|
|
Operation *mlir::cloneWithoutRegions(OpBuilder &b, Operation *op,
|
|
TypeRange newResultTypes,
|
|
ValueRange newOperands) {
|
|
OperationState state(op->getLoc(), op->getName(), newOperands, newResultTypes,
|
|
op->getAttrs());
|
|
for (size_t cnt = 0, e = op->getNumRegions(); cnt < e; ++cnt)
|
|
state.addRegion();
|
|
return b.create(state);
|
|
}
|
|
|
|
SmallVector<NamedAttribute>
|
|
mlir::getPrunedAttributeList(Operation *op, ArrayRef<StringRef> elidedAttrs) {
|
|
llvm::StringSet<> elidedAttrsSet;
|
|
elidedAttrsSet.insert(elidedAttrs.begin(), elidedAttrs.end());
|
|
SmallVector<NamedAttribute> attrs;
|
|
for (auto attr : op->getAttrs()) {
|
|
if (elidedAttrsSet.count(attr.getName()))
|
|
continue;
|
|
attrs.push_back(attr);
|
|
}
|
|
return attrs;
|
|
}
|