replica
/
hanchenye-scalehls
2
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[HLSKernelToAffine] impl of GemmOp and SymmOp lowering

This commit is contained in:
Hanchen Ye 2020-11-30 22:57:42 -06:00
parent 29aff5e6fb
commit ac1c9966a8
4 changed files with 146 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
include/Dialect/HLSKernel/BLASOps.td
View File

@ -40,7 +40,7 @@ def SymmOp : HLSKernelOp<"symm", [HLSKernelOpInterface]> {
C = alpha * B * A + beta * C,
A: N x N, symmetric,
UPLO (false / true): A is upper / lower triangular,
UPLO (false / true): A is lower / upper triangular,
B: M x N,
C: M x N
@ -68,7 +68,7 @@ def SyrkOp : HLSKernelOp<"syrk", [HLSKernelOpInterface]> {
C = alpha * A^T * A + beta * C,
A: K x N,
UPLO (false / true): C is upper / lower triangular,
UPLO (false / true): C is lower / upper triangular,
C: N x N, symmetric
}];
@ -96,7 +96,7 @@ def Syr2kOp : HLSKernelOp<"syr2k", [HLSKernelOpInterface]> {
A: K x N,
B: K x N,
UPLO (false / true): C is upper / lower triangular,
UPLO (false / true): C is lower / upper triangular,
C: N x N, symmetric
}];
@ -123,7 +123,7 @@ def TrmmOp : HLSKernelOp<"trmm", [HLSKernelOpInterface]> {
B = alpha * B * op(A),
A: N x N, triangular,
UPLO (false / true): A is upper / lower triangular,
UPLO (false / true): A is lower / upper triangular,
TRANSA (false / true): op(A) = A / op(A) = A^T,
DIAG (false / true): A is non-unit / unit triangular,

128
lib/Conversion/HLSKernelToAffine/HLSKernelToAffine.cpp
View File

@ -44,13 +44,42 @@ private:
OpBuilder &builder;
Location loc;
// Helpers for creating loops, loads, stores and binary operations.
Value createLoop(unsigned upper, unsigned step = 1, unsigned lower = 0) {
// Helpers for creating loops.
// Constant upper and lower bound.
Value createLoop(int64_t upper, int64_t lower = 0, int64_t step = 1) {
auto loop = builder.create<mlir::AffineForOp>(loc, lower, upper, step);
builder.setInsertionPointToStart(&loop.getLoopBody().front());
return loop.getInductionVar();
}
// General case.
Value createLoop(std::initializer_list<Value> upper, AffineMap upperMap,
std::initializer_list<Value> lower, AffineMap lowerMap,
int64_t step = 1) {
auto loop = builder.create<mlir::AffineForOp>(loc, lower, lowerMap, upper,
upperMap, step);
builder.setInsertionPointToStart(&loop.getLoopBody().front());
return loop.getInductionVar();
}
Value createLoop(Value upper, Value lower, int64_t step = 1) {
auto indexMap = AffineMap::get(1, 0, getDim(0), builder.getContext());
return createLoop({upper}, indexMap, {lower}, indexMap);
}
Value createLoop(int64_t upper, Value lower, int64_t step = 1) {
auto lowerMap = AffineMap::get(1, 0, getDim(0), builder.getContext());
auto upperMap = AffineMap::get(0, 0, getConst(upper), builder.getContext());
return createLoop({}, upperMap, {lower}, lowerMap);
}
Value createLoop(Value upper, int64_t lower, int64_t step = 1) {
auto lowerMap = AffineMap::get(0, 0, getConst(lower), builder.getContext());
auto upperMap = AffineMap::get(1, 0, getDim(0), builder.getContext());
return createLoop({upper}, upperMap, {}, lowerMap);
}
// Helpers for creating loads, stores and binary operations.
Value createLoad(Value array, std::initializer_list<Value> index) {
return builder.create<mlir::AffineLoadOp>(loc, array,
ArrayRef<Value>(index));
@ -319,9 +348,100 @@ bool HLSKernelVisitor::visitOp(MergeOp op) {
//===----------------------------------------------------------------------===//
// Only default attributes configuration are supported.
bool HLSKernelVisitor::visitOp(GemmOp op) { return true; }
bool HLSKernelVisitor::visitOp(GemmOp op) {
auto alpha = op.getOperand(0);
auto beta = op.getOperand(1);
bool HLSKernelVisitor::visitOp(SymmOp op) { return true; }
auto A = op.getOperand(2);
auto B = op.getOperand(3);
auto C = op.getOperand(4);
auto AShape = A.getType().cast<MemRefType>().getShape();
auto CShape = C.getType().cast<MemRefType>().getShape();
// Set insertion point of builder.
builder.setInsertionPoint(op);
// Create M dimension loop.
auto m = createLoop(CShape[0]);
// Create N dimension loop.
auto n = createLoop(CShape[1]);
// Update C with beta * C.
auto initC = createLoad(C, {m, n});
auto betaC = createBinaryOp<mlir::MulFOp>(beta, initC);
createStore(betaC, C, {m, n});
// Create K dimension loop.
auto k = createLoop(AShape[1]);
// Accumulate C with alpha * A * B.
auto valA = createLoad(A, {m, k});
auto valB = createLoad(B, {k, n});
auto valC = createLoad(C, {m, n});
auto alphaA = createBinaryOp<mlir::MulFOp>(alpha, valA);
auto alphaAB = createBinaryOp<mlir::MulFOp>(alphaA, valB);
auto accumC = createBinaryOp<mlir::AddFOp>(alphaAB, valC);
createStore(accumC, C, {m, n});
return true;
}
bool HLSKernelVisitor::visitOp(SymmOp op) {
auto alpha = op.getOperand(0);
auto beta = op.getOperand(1);
auto A = op.getOperand(2);
auto B = op.getOperand(3);
auto C = op.getOperand(4);
auto CShape = C.getType().cast<MemRefType>().getShape();
// Set insertion point of builder.
builder.setInsertionPoint(op);
// Create M dimension loop.
auto m = createLoop(CShape[0]);
// Create N dimension loop.
auto n = createLoop(CShape[1]);
// Update C with beta * C.
auto initC = createLoad(C, {m, n});
auto betaC = createBinaryOp<mlir::MulFOp>(beta, initC);
createStore(betaC, C, {m, n});
// Create K dimension loop for lower triangle.
auto lk = createLoop(m, 0);
// Accumulate C with alpha * A * B.
auto valA = createLoad(A, {m, lk});
auto valB = createLoad(B, {lk, n});
auto valC = createLoad(C, {m, n});
auto alphaA = createBinaryOp<mlir::MulFOp>(alpha, valA);
auto alphaAB = createBinaryOp<mlir::MulFOp>(alphaA, valB);
auto accumC = createBinaryOp<mlir::AddFOp>(alphaAB, valC);
createStore(accumC, C, {m, n});
// Create K dimension loop for upper triangle.
builder.setInsertionPoint(n.getParentBlock()->getTerminator());
auto hk = createLoop(CShape[0], m);
// Accumulate C with alpha * A * B.
valA = createLoad(A, {hk, m});
valB = createLoad(B, {hk, n});
valC = createLoad(C, {m, n});
alphaA = createBinaryOp<mlir::MulFOp>(alpha, valA);
alphaAB = createBinaryOp<mlir::MulFOp>(alphaA, valB);
accumC = createBinaryOp<mlir::AddFOp>(alphaAB, valC);
createStore(accumC, C, {m, n});
return true;
}
bool HLSKernelVisitor::visitOp(SyrkOp op) { return true; }

9
test/Conversion/HLSKernelToAffine/test_gemm.mlir Normal file
View File

@ -0,0 +1,9 @@
// RUN: scalehls-opt -hlskernel-to-affine %s | FileCheck %s
// CHECK: module {
func @test_gemm(%A: memref<32x16xf32>, %B: memref<16x8xf32>, %C: memref<32x8xf32>) -> () {
%alpha = constant 11.0 : f32
%beta = constant 42.0 : f32
"hlskernel.gemm" (%alpha, %beta, %A, %B, %C) {} : (f32, f32, memref<32x16xf32>, memref<16x8xf32>, memref<32x8xf32>) -> ()
return
}

9
test/Conversion/HLSKernelToAffine/test_symm.mlir Normal file
View File

@ -0,0 +1,9 @@
// RUN: scalehls-opt -hlskernel-to-affine %s | FileCheck %s
// CHECK: module {
func @test_symm(%A: memref<32x32xf32>, %B: memref<32x8xf32>, %C: memref<32x8xf32>) -> () {
%alpha = constant 11.0 : f32
%beta = constant 42.0 : f32
"hlskernel.symm" (%alpha, %beta, %A, %B, %C) {} : (f32, f32, memref<32x32xf32>, memref<32x8xf32>, memref<32x8xf32>) -> ()
return
}