From 5907fe3f1693f865d81da4ce3ddf002060664c34 Mon Sep 17 00:00:00 2001 From: philipportner Date: Thu, 22 May 2025 11:36:50 +0200 Subject: [PATCH 1/9] Add TODOs --- doc/DaphneDSL/Builtins.md | 3 ++- src/compiler/lowering/SpecializeGenericFunctionsPass.cpp | 1 + src/ir/daphneir/DaphneOps.td | 1 + src/parser/daphnedsl/DaphneDSLVisitor.cpp | 1 + src/runtime/local/kernels/Map.h | 1 + 5 files changed, 6 insertions(+), 1 deletion(-) diff --git a/doc/DaphneDSL/Builtins.md b/doc/DaphneDSL/Builtins.md index c9214a1c3..382e2d0ac 100644 --- a/doc/DaphneDSL/Builtins.md +++ b/doc/DaphneDSL/Builtins.md @@ -296,6 +296,7 @@ The following built-in functions all follow the same scheme: Full aggregation over all elements of the matrix `arg` using aggregation function `agg` (see table below). Returns a scalar. +// TODO(#520) - **`agg`**`(arg:matrix, axis:si64)` Row or column aggregation over a *(n x m)* matrix `arg` using aggregation function `agg` (see table below). @@ -701,4 +702,4 @@ These must be provided in a separate [`.meta`-file](/doc/FileMetaDataFormat.md). - **`remove`**`(lst:list, idx:size)` Removes the element at position `idx` (counting starts at zero) from the given list `lst`. - Returns (1) the result as a new list (the argument list stays unchanged), and (2) the removed element. \ No newline at end of file + Returns (1) the result as a new list (the argument list stays unchanged), and (2) the removed element. diff --git a/src/compiler/lowering/SpecializeGenericFunctionsPass.cpp b/src/compiler/lowering/SpecializeGenericFunctionsPass.cpp index e6b98bc8a..b33f320e0 100644 --- a/src/compiler/lowering/SpecializeGenericFunctionsPass.cpp +++ b/src/compiler/lowering/SpecializeGenericFunctionsPass.cpp @@ -362,6 +362,7 @@ class SpecializeGenericFunctionsPass : public PassWrapper { let summary = "Applies a user defined function to elements of a matrix."; let arguments = (ins MatrixOrU:$arg, SymbolNameAttr:$func); diff --git a/src/parser/daphnedsl/DaphneDSLVisitor.cpp b/src/parser/daphnedsl/DaphneDSLVisitor.cpp index c3dc2e628..790a160d1 100644 --- a/src/parser/daphnedsl/DaphneDSLVisitor.cpp +++ b/src/parser/daphnedsl/DaphneDSLVisitor.cpp @@ -880,6 +880,7 @@ DaphneDSLVisitor::findMatchingUnaryUDF(mlir::Location loc, const std::string &fu return std::nullopt; } +// TODO(#520) antlrcpp::Any DaphneDSLVisitor::handleMapOpCall(DaphneDSLGrammarParser::CallExprContext *ctx) { std::string func; const auto &identifierVec = ctx->IDENTIFIER(); diff --git a/src/runtime/local/kernels/Map.h b/src/runtime/local/kernels/Map.h index 8f2382353..7b1ab32d2 100644 --- a/src/runtime/local/kernels/Map.h +++ b/src/runtime/local/kernels/Map.h @@ -64,6 +64,7 @@ template struct Map, DenseMa const VTArg *valuesArg = arg->getValues(); VTRes *valuesRes = res->getValues(); + // TODO(#520) for (size_t r = 0; r < numRows; r++) { for (size_t c = 0; c < numCols; c++) valuesRes[c] = udf(valuesArg[c]); From cea14437ff5dbef9ca3b17fb8196aac0fae26a29 Mon Sep 17 00:00:00 2001 From: Sara Lhamo Schnaterbeck Date: Sat, 7 Jun 2025 20:44:23 +0200 Subject: [PATCH 2/9] [DAPHNE-#520] Added row-/column-wise map --- src/compiler/lowering/LowerToLLVMPass.cpp | 5 +- .../SpecializeGenericFunctionsPass.cpp | 9 +- src/ir/daphneir/DaphneOps.td | 3 +- src/parser/daphnedsl/DaphneDSLBuiltins.cpp | 21 ++- src/parser/daphnedsl/DaphneDSLVisitor.cpp | 10 +- src/runtime/local/kernels/Map.h | 74 ++++++++-- src/runtime/local/kernels/kernels.json | 4 + test/api/cli/codegen/MapOpTest.cpp | 19 +++ test/api/cli/codegen/map_col.daphne | 10 ++ test/api/cli/codegen/map_row.daphne | 10 ++ test/runtime/local/kernels/MapTest.cpp | 132 +++++++++++++++++- 11 files changed, 265 insertions(+), 32 deletions(-) create mode 100644 test/api/cli/codegen/map_col.daphne create mode 100644 test/api/cli/codegen/map_row.daphne diff --git a/src/compiler/lowering/LowerToLLVMPass.cpp b/src/compiler/lowering/LowerToLLVMPass.cpp index a30dc4199..80f874e5c 100644 --- a/src/compiler/lowering/LowerToLLVMPass.cpp +++ b/src/compiler/lowering/LowerToLLVMPass.cpp @@ -508,11 +508,14 @@ class MapOpLowering : public OpConversionPattern { // Pointer to UDF callee << "__void"; + // Axis + callee << "__int64_t"; + // get pointer to UDF LLVM::LLVMFuncOp udfFuncOp = module.lookupSymbol(op.getFunc()); auto udfFnPtr = rewriter.create(loc, udfFuncOp); - std::vector kernelOperands{op.getArg(), udfFnPtr}; + std::vector kernelOperands{op.getArg(), udfFnPtr, op.getAxis()}; auto kernel = rewriter.create(loc, callee.str(), kernelOperands, op->getResultTypes()); rewriter.replaceOp(op, kernel.getResults()); diff --git a/src/compiler/lowering/SpecializeGenericFunctionsPass.cpp b/src/compiler/lowering/SpecializeGenericFunctionsPass.cpp index b33f320e0..967196c81 100644 --- a/src/compiler/lowering/SpecializeGenericFunctionsPass.cpp +++ b/src/compiler/lowering/SpecializeGenericFunctionsPass.cpp @@ -362,7 +362,6 @@ class SpecializeGenericFunctionsPass : public PassWrapper(); + int64_t axis = + CompilerUtils::constantOrThrow(mapOp.getAxis(), "map axis must be a constant."); + if (axis == 0) // row-wise map + inpMatrixTy.withShape(inpMatrixTy.getNumRows(), 1); // TODO(#520) case ncols != 1 + else if (axis == 1) // column-wise map + inpMatrixTy.withShape(1, inpMatrixTy.getNumCols()); func::FuncOp specializedFunc = createOrReuseSpecialization(inpMatrixTy.getElementType(), {}, calledFunction, mapOp.getLoc()); mapOp.setFuncAttr(specializedFunc.getSymNameAttr()); @@ -397,7 +402,7 @@ class SpecializeGenericFunctionsPass : public PassWrapper { let summary = "Applies a user defined function to elements of a matrix."; - let arguments = (ins MatrixOrU:$arg, SymbolNameAttr:$func); + let arguments = (ins MatrixOrU:$arg, SymbolNameAttr:$func, SI64:$axis); let results = (outs MatrixOrU:$res); } diff --git a/src/parser/daphnedsl/DaphneDSLBuiltins.cpp b/src/parser/daphnedsl/DaphneDSLBuiltins.cpp index a2577995a..d692d16a5 100644 --- a/src/parser/daphnedsl/DaphneDSLBuiltins.cpp +++ b/src/parser/daphnedsl/DaphneDSLBuiltins.cpp @@ -1299,14 +1299,29 @@ antlrcpp::Any DaphneDSLBuiltins::build(mlir::Location loc, const std::string &fu // **************************************************************************** if (func == "map") { - checkNumArgsExact(loc, func, numArgs, 2); + checkNumArgsBetween(loc, func, numArgs, 2, 3); + mlir::Value source = args[0]; auto co = args[1].getDefiningOp(); mlir::Attribute attr = co.getValue(); - return static_cast( - builder.create(loc, source.getType(), source, attr.dyn_cast())); + switch (numArgs) { + case 2: { // axis is not given, so it defaults to -1 + mlir::Value minusOne = builder.create(loc, int64_t(-1)); + return static_cast( + builder.create(loc, source.getType(), source, attr.dyn_cast(), minusOne)); + } + case 3: { // axis is given + int64_t axis = + CompilerUtils::constantOrThrow(args[2], "third argument of map must be a constant"); + if (axis == 0 || axis == 1) + return static_cast( + builder.create(loc, source.getType(), source, attr.dyn_cast(), args[2])); + else + throw ErrorHandler::compilerError(loc, "DSLBuiltins", "invalid axis for aggregation."); + } + } } // **************************************************************************** diff --git a/src/parser/daphnedsl/DaphneDSLVisitor.cpp b/src/parser/daphnedsl/DaphneDSLVisitor.cpp index 790a160d1..5559e2154 100644 --- a/src/parser/daphnedsl/DaphneDSLVisitor.cpp +++ b/src/parser/daphnedsl/DaphneDSLVisitor.cpp @@ -880,7 +880,6 @@ DaphneDSLVisitor::findMatchingUnaryUDF(mlir::Location loc, const std::string &fu return std::nullopt; } -// TODO(#520) antlrcpp::Any DaphneDSLVisitor::handleMapOpCall(DaphneDSLGrammarParser::CallExprContext *ctx) { std::string func; const auto &identifierVec = ctx->IDENTIFIER(); @@ -893,9 +892,9 @@ antlrcpp::Any DaphneDSLVisitor::handleMapOpCall(DaphneDSLGrammarParser::CallExpr throw ErrorHandler::compilerError(loc, "DSLVisitor", "called 'handleMapOpCall' for function " + func + " instead of 'map'"); - if (ctx->expr().size() != 2) { + if (ctx->expr().size() != 2 && ctx->expr().size() != 3) { throw ErrorHandler::compilerError(loc, "DSLVisitor", - "built-in function 'map' expects exactly 2 argument(s), but got " + + "built-in function 'map' expects 2-3 argument(s), but got " + std::to_string(ctx->expr().size())); } @@ -919,6 +918,11 @@ antlrcpp::Any DaphneDSLVisitor::handleMapOpCall(DaphneDSLGrammarParser::CallExpr args.push_back( static_cast(builder.create(loc, maybeUDF->getSymName().str()))); + if (ctx->expr().size() == 3) { + auto axis = valueOrErrorOnVisit(ctx->expr(2)); + args.push_back(axis); + } + // Create DaphneIR operation for the built-in function. return builtins.build(loc, func, args); } diff --git a/src/runtime/local/kernels/Map.h b/src/runtime/local/kernels/Map.h index 7b1ab32d2..2402a192b 100644 --- a/src/runtime/local/kernels/Map.h +++ b/src/runtime/local/kernels/Map.h @@ -32,15 +32,15 @@ template struct Map { // We could have a more specialized function pointer here i.e. // (DTRes::VT)(*func)(DTArg::VT). The problem is that this is currently not // supported by kernels.json. - static void apply(DTRes *&res, const DTArg *arg, void *func, DCTX(ctx)) = delete; + static void apply(DTRes *&res, const DTArg *arg, void *func, const int64_t axis, DCTX(ctx)) = delete; }; // **************************************************************************** // Convenience function // **************************************************************************** -template void map(DTRes *&res, const DTArg *arg, void *func, DCTX(ctx)) { - Map::apply(res, arg, func, ctx); +template void map(DTRes *&res, const DTArg *arg, void *func, const int64_t axis, DCTX(ctx)) { + Map::apply(res, arg, func, axis, ctx); } // **************************************************************************** @@ -52,24 +52,68 @@ template void map(DTRes *&res, const DTArg *arg, void // ---------------------------------------------------------------------------- template struct Map, DenseMatrix> { - static void apply(DenseMatrix *&res, const DenseMatrix *arg, void *func, DCTX(ctx)) { + static void apply(DenseMatrix *&res, const DenseMatrix *arg, void *func, const int64_t axis, DCTX(ctx)) { const size_t numRows = arg->getNumRows(); const size_t numCols = arg->getNumCols(); - if (res == nullptr) - res = DataObjectFactory::create>(numRows, numCols, false); + if (res == nullptr) { + if (axis == 0) // row-wise + res = DataObjectFactory::create>(numRows, 1, false); + else if (axis == 1) // column-wise + res = DataObjectFactory::create>(1, numCols, false); + else // element-wise + res = DataObjectFactory::create>(numRows, numCols, false); + } - auto udf = reinterpret_cast(func); + auto udfElem = reinterpret_cast(func); + auto udfMat = reinterpret_cast* (*)(DenseMatrix*)>(func); const VTArg *valuesArg = arg->getValues(); VTRes *valuesRes = res->getValues(); - // TODO(#520) - for (size_t r = 0; r < numRows; r++) { - for (size_t c = 0; c < numCols; c++) - valuesRes[c] = udf(valuesArg[c]); - valuesArg += arg->getRowSkip(); - valuesRes += res->getRowSkip(); + if (axis == 1) { // column-wise + size_t resNumRows = 1; + // Extract each column, apply udf and set result row-wise + for (size_t c = 0; c < numCols; c++) { + auto currentCol = arg->sliceCol(c, c + 1); + // auto currentCol = &DenseMatrix(arg, 0, numRows, c, c + 1); + DenseMatrix *resCol = udfMat(currentCol); + if (c == 0) { + resNumRows = resCol->getNumRows(); + res = DataObjectFactory::create>(resNumRows, numCols, false); + } + VTRes *valuesResCol = resCol->getValues(); + for (size_t r = 0; r < resNumRows; r++) { + valuesRes[c] = valuesResCol[0]; + valuesResCol += resCol->getRowSkip(); + valuesRes += res->getRowSkip(); + } + valuesRes = res->getValues(); + DataObjectFactory::destroy(currentCol); + } + } else { // row-wise or element-wise + size_t resNumCols = 1; + // Extract each row, apply udf and copy values to result + for (size_t r = 0; r < numRows; r++) { + if (axis == 0) { + // auto currentRow = &DenseMatrix(1, numCols, std::make_shared(valuesArg)); + // auto currentRow = &DenseMatrix(arg, r, r + 1, 0, numCols); + auto currentRow = arg->sliceRow(r, r + 1); + DenseMatrix *resRow = udfMat(currentRow); + if (r == 0) { + resNumCols = resRow->getNumCols(); + res = DataObjectFactory::create>(numRows, resNumCols, false); + } + // valuesRes = resRow->getValues(); + memcpy(valuesRes, resRow->getValues(), resNumCols * sizeof(VTRes)); + DataObjectFactory::destroy(currentRow); + } else { + for (size_t c = 0; c < numCols; c++) + valuesRes[c] = udfElem(valuesArg[c]); + valuesArg += arg->getRowSkip(); + } + valuesRes += res->getRowSkip(); + } } } }; @@ -78,8 +122,8 @@ template struct Map, DenseMa // Matrix // ---------------------------------------------------------------------------- -template struct Map, Matrix> { - static void apply(Matrix *&res, const Matrix *arg, void *func, DCTX(ctx)) { +template struct Map, Matrix> { // TODO(#520) + static void apply(Matrix *&res, const Matrix *arg, void *func, const int64_t axis, DCTX(ctx)) { const size_t numRows = arg->getNumRows(); const size_t numCols = arg->getNumCols(); diff --git a/src/runtime/local/kernels/kernels.json b/src/runtime/local/kernels/kernels.json index c663cc499..94d67ab52 100644 --- a/src/runtime/local/kernels/kernels.json +++ b/src/runtime/local/kernels/kernels.json @@ -3183,6 +3183,10 @@ { "type": "void *", "name": "func" + }, + { + "type": "const int64_t", + "name": "axis" } ] }, diff --git a/test/api/cli/codegen/MapOpTest.cpp b/test/api/cli/codegen/MapOpTest.cpp index 32b0075f4..3bd2047ec 100644 --- a/test/api/cli/codegen/MapOpTest.cpp +++ b/test/api/cli/codegen/MapOpTest.cpp @@ -33,3 +33,22 @@ TEST_CASE("mapOp", TAG_CODEGEN) { compareDaphneToStr(result, dirPath + "map.daphne"); compareDaphneToStr(result, dirPath + "map.daphne", "--mlir-codegen", "--no-obj-ref-mgnt"); } + +TEST_CASE("mapOp col", TAG_CODEGEN) { + std::string result = "DenseMatrix(3x4, double)\n" + "1.1 0.0 5.5 -2.2\n" + "1.1 0.0 5.5 -2.2\n" + "1.1 0.0 5.5 -2.2\n"; + + compareDaphneToStr(result, dirPath + "map_col.daphne"); + compareDaphneToStr(result, dirPath + "map_col.daphne", "--mlir-codegen", "--no-obj-ref-mgnt"); +} + +TEST_CASE("mapOp row", TAG_CODEGEN) { + std::string result = "DenseMatrix(2x1, double)\n" + "4.1\n" + "4.3\n"; + + compareDaphneToStr(result, dirPath + "map_row.daphne"); + compareDaphneToStr(result, dirPath + "map_row.daphne", "--mlir-codegen", "--no-obj-ref-mgnt"); +} diff --git a/test/api/cli/codegen/map_col.daphne b/test/api/cli/codegen/map_col.daphne new file mode 100644 index 000000000..a5a38ed4d --- /dev/null +++ b/test/api/cli/codegen/map_col.daphne @@ -0,0 +1,10 @@ +// Performs a MapOp with the UDF `topThree`. Used to compare precompiled +// kernel with codegen. + +def topThree(col:matrix) -> matrix { + return col[:3, 0]; +} + +X = reshape([1.1, 0.0, 5.5, -2.2, 1.1, 0.0, 5.5, -2.2, 1.1, 0.0, 5.5, -2.2, 1.1, 0.0, 5.5, -2.2], 4, 4); + +print(map(X, topThree, 1)); diff --git a/test/api/cli/codegen/map_row.daphne b/test/api/cli/codegen/map_row.daphne new file mode 100644 index 000000000..f30d012cd --- /dev/null +++ b/test/api/cli/codegen/map_row.daphne @@ -0,0 +1,10 @@ +// Performs a MapOp with the UDF `sumOfFirstAndLast`. Used to compare precompiled +// kernel with codegen. + +def sumOfFirstAndLast(row:matrix) -> f64 { + return as.scalar(row[0, 0]) + as.scalar(row[0, 2]); //ncol(row) – 1 +} + +X = reshape([1.1, 0.0, 3.0, 5.5, -2.2, -1.2], 2, 3); + +print(map(X, sumOfFirstAndLast, 0)); diff --git a/test/runtime/local/kernels/MapTest.cpp b/test/runtime/local/kernels/MapTest.cpp index a59613c7d..c4492a90e 100644 --- a/test/runtime/local/kernels/MapTest.cpp +++ b/test/runtime/local/kernels/MapTest.cpp @@ -27,9 +27,9 @@ #define TYPES double, float, int64_t, int32_t, int8_t, uint64_t, uint8_t -template void checkMap(const DTArg *arg, const DTRes *exp, void *func) { +template void checkMap(const DTArg *arg, const DTRes *exp, void *func, int64_t axis) { DTRes *res = nullptr; - map(res, arg, func, nullptr); + map(res, arg, func, axis, nullptr); CHECK(*res == *exp); DataObjectFactory::destroy(res); } @@ -92,15 +92,15 @@ template