From 9e6e9d128b9e4117d9c6dd44ec4760035a5db4ac Mon Sep 17 00:00:00 2001 From: Polykarpos Thomadakis Date: Fri, 4 Apr 2025 18:38:17 -0700 Subject: [PATCH 01/11] [WIP] Removing SetOp --- frontends/comet_dsl/comet.cpp | 2 +- frontends/comet_dsl/mlir/MLIRGen.cpp | 51 +++++------------ .../TensorAlgebraToIndexTree.cpp | 55 ++++++++++++++++-- .../TensorAlgebra/Transforms/TCtoTTGTDyn.cpp | 56 ++++++++++++------- 4 files changed, 100 insertions(+), 64 deletions(-) diff --git a/frontends/comet_dsl/comet.cpp b/frontends/comet_dsl/comet.cpp index 7ad5da23..752b3750 100644 --- a/frontends/comet_dsl/comet.cpp +++ b/frontends/comet_dsl/comet.cpp @@ -427,7 +427,7 @@ int loadAndProcessMLIR(mlir::MLIRContext &context, /// Check if there are missing tensor declaration operations introduced by compound expressions. /// If so, add a new tensor declaration to represent intermediate tensors /// ============================================================================= - optPM.addPass(mlir::comet::createTensorAlgebraCheckImplicitTensorDeclPass()); + // optPM.addPass(mlir::comet::createTensorAlgebraCheckImplicitTensorDeclPass()); /// ============================================================================= /// Check to see if we are dumping to TA dialect. if (emitTA) diff --git a/frontends/comet_dsl/mlir/MLIRGen.cpp b/frontends/comet_dsl/mlir/MLIRGen.cpp index 51991669..f558ccbc 100644 --- a/frontends/comet_dsl/mlir/MLIRGen.cpp +++ b/frontends/comet_dsl/mlir/MLIRGen.cpp @@ -494,29 +494,16 @@ namespace comet_vdump(rhs); comet_vdump(lhs); - /// lookup the output of the binary operation - auto theOutput = symbolTable.lookup(out_format); - if (theOutput == nullptr) - { /// the variable for output of binary operation was not declared by user, - /// we will create a new DenseConstantOp here. - comet_debug() << "creating a new variable declaration, since the user did not declare it\n"; - - double data = 0.0; - auto dataAttribute = mlir::DenseElementsAttr::get(mlir::RankedTensorType::get({1}, builder.getF64Type()), llvm::ArrayRef(data)); - auto denseConst = builder.create(location, returnDataType, dataAttribute); - - theOutput = denseConst; - } comet_vdump(theOutput); auto scalarOp = builder.create(location, returnDataType, rhs, lhs, opAttr); comet_vdump(scalarOp); - builder.create(location, scalarOp, theOutput); + symbolTable.insert(out_format, scalarOp); /// the value returned here will be used in subsequent ops. /// for example, in the code below, 'g' should be returned. /// $ var g = a + b; /// $ print(g); - return theOutput; + return scalarOp; } else if (isa(lhs.getDefiningOp())) @@ -1966,8 +1953,7 @@ namespace LabeledTensorExprAST *lhsLabeledTensorExprAST = llvm::cast(tensor_op->getLHS()); CallExprAST *call = llvm::cast(tensor_op->getRHS()); auto call_res = mlirGen(*call); - auto lhs_tensor = symbolTable.lookup(lhsLabeledTensorExprAST->getTensorName()); - builder.create(loc(tensor_op->loc()), call_res, lhs_tensor); + symbolTable.insert(lhsLabeledTensorExprAST->getTensorName(), call_res); continue; } /// TODO(gkestor): evaluate use of Expr_LabeledTensor for slicing @@ -2109,18 +2095,11 @@ namespace llvm::errs() << __FILE__ << ":" << __LINE__ << " ERROR: Unsupported tensor elementwise addition\n"; /// TODO(gkestor): look at tensor elementwise addition /// auto SemiringAttr = builder.getStringAttr("eltwise_add"); /// this is for standard elementwise addition - /// auto op = builder.create(loc(tensor_op.loc()), mlir::UnrankedTensorType::get(builder.getF64Type()), - /// tensors[1], tensors[0], builder.getStrArrayAttr(formats), ); - /// builder.create(loc(tensor_op.loc()), op.getOperation()->getResult(0), tensors[1]); } else if (binop == TensorOpKind::Tensor_Red_Sub) { llvm::errs() << __FILE__ << ":" << __LINE__ << " ERROR: Unsupported tensor elementwise substraction\n"; - /// TODO(gkestor): look at tensor elementwise subtraction - /// auto op = builder.create(loc(tensor_op.loc()), mlir::UnrankedTensorType::get(builder.getF64Type()), - /// tensors[1], tensors[0], builder.getStrArrayAttr(formats)); - /// builder.create(loc(tensor_op.loc()), op.getOperation()->getResult(0), tensors[1]); } return mlir::success(); @@ -2321,14 +2300,14 @@ namespace mlir::Value lhsLT_op; if ((lhsLT_op = symbolTable.lookup(lhsLT_tensor_name)) != NULL) { - if (isa(lhsLT_op.getDefiningOp())) + // if (isa(lhsLT_op.getDefiningOp())) { tensors.push_back(lhsLT_op); } - else - { - comet_debug() << " not TensorDeclOp\n"; - } + // elsez + // { + // comet_debug() << " not TensorDeclOp\n"; + // } } } @@ -2355,8 +2334,7 @@ namespace MaskingAttr); comet_vdump(op); - /// source is 1st parameter, dest is the second - builder.create(loc(tensor_op.loc()), op.getOperation()->getResult(0), tensors[2]); + symbolTable.insert(exprs[2]->getTensorName(), op); } else if (binop == '-') { @@ -2369,8 +2347,7 @@ namespace SemiringAttr, MaskingAttr); comet_vdump(op); - /// source is 1st parameter, dest is the second - builder.create(loc(tensor_op.loc()), op.getOperation()->getResult(0), tensors[2]); + symbolTable.insert(exprs[2]->getTensorName(), op); } else if (binop == '*' || binop == tok_semiring) { @@ -2384,9 +2361,8 @@ namespace op.getOperation()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); op.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); - /// source is 1st parameter, dest is the second - auto setop = builder.create(loc(tensor_op.loc()), op.getOperation()->getResult(0), tensors[2]); - setop.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); + symbolTable.insert(exprs[2]->getTensorName(), op); + } else if (binop == tok_elews || binop == tok_monoid) { @@ -2394,8 +2370,7 @@ namespace op.getOperation()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); op.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); - auto setop = builder.create(loc(tensor_op.loc()), op.getOperation()->getResult(0), tensors[2]); - setop.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); + symbolTable.insert(exprs[2]->getTensorName(), op); } else { diff --git a/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp b/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp index 1fe4086d..b14818df 100644 --- a/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp +++ b/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp @@ -31,12 +31,20 @@ #include "comet/Dialect/IndexTree/IR/IndexTreeDialect.h" #include "comet/Dialect/IndexTree/Passes.h" #include "comet/Dialect/Utils/Utils.h" +#include "mlir/Dialect/Arith/IR/Arith.h" #include "mlir/Dialect/Func/IR/FuncOps.h" +#include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/AffineExpr.h" +#include "mlir/IR/AffineMap.h" #include "mlir/IR/Block.h" +#include "mlir/IR/BuiltinAttributes.h" +#include "mlir/IR/BuiltinTypeInterfaces.h" #include "mlir/IR/Operation.h" +#include "mlir/IR/Value.h" #include "mlir/IR/ValueRange.h" +#include "mlir/Support/LLVM.h" +#include "llvm/ADT/ArrayRef.h" using namespace mlir; using namespace mlir::indexTree; @@ -235,9 +243,45 @@ mlir::LogicalResult generalIndexOperationRewrite( auto context = rewriter.getContext(); TATensorOp mult_op = llvm::dyn_cast(op); - Value rhs1_tensor = getRealRhs(mult_op.getRhs1()); - Value rhs2_tensor = getRealRhs(mult_op.getRhs2()); - Value lhs_tensor = getRealLhs(op); + Value rhs1_tensor = mult_op.getRhs1(); + Value rhs2_tensor = mult_op.getRhs2(); + Value lhs_tensor; + SmallVector dims; + auto shapeT = cast(mult_op.getResult().getType()); + ArrayAttr indexing_maps = cast(mult_op.getIndexingMaps()); + for(auto [index, v]: enumerate(cast(indexing_maps[2]).getValue().getResults())) + { + if(!cast(rhs1_tensor.getType()).isDynamicDim(index)) + { + continue; + } + + AffineMap map = cast(indexing_maps[0]).getValue(); + if (auto pos = map.getResultPosition(v)) + { + auto dim = rewriter.create(loc, rhs1_tensor, *pos); + dims.push_back(dim); + continue; + } + + map = cast(indexing_maps[1]).getValue(); // try the second map (rhs2) + if (auto pos = map.getResultPosition(v)) + { + auto dim = rewriter.create(loc, rhs2_tensor, *pos); + dims.push_back(dim); + continue; + } + } + + if(auto spTensorT = dyn_cast(op->getResultTypes()[0])) + { + lhs_tensor = rewriter.create(loc, shapeT, ValueRange(dims), false); + } + else if (auto tensorT = dyn_cast(op->getResultTypes()[0])) + { + lhs_tensor = rewriter.create(loc, shapeT, ValueRange(dims)); + // lhs_tensor = rewriter.create(loc, shapeT, ValueRange(dims)); + } comet_vdump(rhs1_tensor); comet_vdump(rhs2_tensor); @@ -253,7 +297,6 @@ mlir::LogicalResult generalIndexOperationRewrite( } } - auto indexing_maps = mult_op.getIndexingMaps(); auto semiring = cast(mult_op.getSemiringAttr()).getValue(); auto tensor_type = op->getResultTypes()[0]; @@ -482,8 +525,8 @@ void LowerTensorAlgebraToIndexTreePass::runOnOperation() comet_pdump(getOperation()->getParentOfType()); mlir::ConversionTarget target(getContext()); - target.addLegalDialect(); - target.addLegalOp(); + target.addLegalDialect(); + target.addLegalOp(); target.addIllegalOp(); diff --git a/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp b/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp index 01adbe4f..efae7814 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp @@ -189,29 +189,46 @@ namespace allPerms.push_back(perm); } - Value lhsDef; - tensorAlgebra::TensorSetOp setnewop; - for (auto u : multop.getOperation()->getResult(0).getUsers()) - { - comet_pdump(u); - if (isa(u)) - { - setnewop = cast(u); - Value dstTensor = u->getOperand(1); - - lhsDef = dstTensor; - comet_vdump(lhsDef); - } - } comet_vdump(setnewop); comet_debug() << "\n"; - Value rhs1Tensor = operands[0], rhs2Tensor = operands[1], lhsTensor = lhsDef; + Value rhs1Tensor = operands[0], rhs2Tensor = operands[1], lhsTensor; auto rhs1TensorType = cast(rhs1Tensor.getType()); auto rhs2TensorType = cast(rhs2Tensor.getType()); - auto lhsTensorType = cast(lhsTensor.getType()); + + SmallVector dims; + auto shapeT = cast(multop.getResult().getType()); + ArrayAttr indexing_maps = cast(multop.getIndexingMaps()); + for(auto [index, v]: enumerate(cast(indexing_maps[2]).getValue().getResults())) + { + if(!cast(rhs1Tensor.getType()).isDynamicDim(index)) + { + continue; + } + + AffineMap map = cast(indexing_maps[0]).getValue(); + if (auto pos = map.getResultPosition(v)) + { + auto dim = rewriter.create(loc, rhs1Tensor, *pos); + dims.push_back(dim); + continue; + } + + map = cast(indexing_maps[1]).getValue(); // try the second map (rhs2) + if (auto pos = map.getResultPosition(v)) + { + auto dim = rewriter.create(loc, rhs2Tensor, *pos); + dims.push_back(dim); + continue; + } + } + + auto zero = rewriter.create(loc, FloatAttr::get(shapeT.getElementType(), 0.0)); + lhsTensor = rewriter.create(loc, shapeT, zero, ValueRange(dims)); + + auto lhsTensorType = cast(lhsTensor.getType()); std::vector allShapes{rhs1Tensor, rhs2Tensor, @@ -710,8 +727,9 @@ namespace loc, printFlopsStr, SmallVector{}, ValueRange{flopsOp}); } - rewriter.replaceAllUsesWith( - op->getResults(), switchOp.getResults()); // Replace the original op with the final result of the matmul or matvec + rewriter.replaceOp(op, switchOp); + // rewriter.replaceAllUsesWith( + // op->getResults(), switchOp.getResults()); // Replace the original op with the final result of the matmul or matvec // rewriter.replaceUsesWithIf(setnewop->getOperand(1), switchOp.getResult(0), [&](OpOperand& use) { // auto user = use.getOwner(); // auto ancestor = switchOp->getBlock()->findAncestorOpInBlock(*user); @@ -719,7 +737,7 @@ namespace // }); // op->replaceAllUsesWith(switchOp); // rewriter.eraseOp(setnewop); - rewriter.eraseOp(op); + // rewriter.eraseOp(op); return success(); } From c70dd56eda14f30b0a7fe0fa18f33fcf4fb31f31 Mon Sep 17 00:00:00 2001 From: Polykarpos Thomadakis Date: Sat, 5 Apr 2025 13:14:36 -0700 Subject: [PATCH 02/11] [WIP] Removing SetOp almost complete. One case is failing --- frontends/comet_dsl/mlir/MLIRGen.cpp | 72 ++++++++----------- .../TensorAlgebraToIndexTree.cpp | 9 ++- .../TensorAlgebraToSCF/TensorAlgebraToSCF.cpp | 36 ++++------ .../IndexTree/Transforms/KernelFusion.cpp | 40 +++++------ 4 files changed, 69 insertions(+), 88 deletions(-) diff --git a/frontends/comet_dsl/mlir/MLIRGen.cpp b/frontends/comet_dsl/mlir/MLIRGen.cpp index f558ccbc..9fd09f76 100644 --- a/frontends/comet_dsl/mlir/MLIRGen.cpp +++ b/frontends/comet_dsl/mlir/MLIRGen.cpp @@ -844,24 +844,9 @@ namespace } else if (isa(e.getDefiningOp())) { - comet_debug() << " is TransposeOp\n"; - - /// get the real transpose op output via the set op. - mlir::Value transposeOut; - mlir::Operation *firstUser = e.getDefiningOp()->getNextNode(); - if (isa(firstUser)) - { - TensorSetOp setOp = cast(firstUser); - transposeOut = setOp.getOperand(1); - } - else - { - llvm::errs() << __FILE__ << ":" << __LINE__ << " ERROR: Transpose has no set_op after it!\n"; - } - /// get the format of transposeOut tensor - formats.push_back(getTensorFormatString(transposeOut.getType())); - tensors.push_back(transposeOut); + formats.push_back(getTensorFormatString(e.getType())); + tensors.push_back(e); } else { @@ -1225,8 +1210,11 @@ namespace auto tens_beta = tensor_op.getBeta(); - auto ret_op = builder.create(loc(tensor_op.loc()), rhs, lhs); - ret_op.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); + auto lhsName = cast(*tensor_op.getLHS()) + .getTensorName(); + symbolTable.insert(lhsName, rhs); + // auto ret_op = builder.create(loc(tensor_op.loc()), rhs, lhs); + // ret_op.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); return rhs; } @@ -1499,7 +1487,7 @@ namespace comet_debug() << " create TransposeOp\n"; mlir::Value t = builder.create(loc(transpose.loc()), lhs_tensor.getType(), rhs_tensor, all_lbls_value, affineMapArrayAttr); - builder.create(loc(transpose.loc()), t, lhs_tensor); + symbolTable.insert(lhsLT.getTensorName(), t) comet_vdump(t); return t; @@ -1622,34 +1610,33 @@ namespace mlir::Type return_type = getType(shape); /// Create Tensor Declarations Ops and populate formats (for lhs) - mlir::Value lhs_tensor; - if (auto tensorT = dyn_cast(rhs_tensor.getType())) - { - auto declOp = builder.create(loc(transpose.loc()), mlir::RankedTensorType::get(shape, tensorT.getElementType()), indices); + // mlir::Value lhs_tensor; + // if (auto tensorT = dyn_cast(rhs_tensor.getType())) + // { + // auto declOp = builder.create(loc(transpose.loc()), mlir::RankedTensorType::get(shape, tensorT.getElementType()), indices); - lhs_tensor = declOp; - } - else if (auto SparseTensorT = dyn_cast(rhs_tensor.getType())) - { - - ArrayRef format = SparseTensorT.getFormat(); - mlir::ShapedType shapedT = mlir::cast(rhs_tensor.getType()); - mlir::Type element_type = shapedT.getElementType(); - return_type = SparseTensorType::get(builder.getContext(), element_type, builder.getIntegerType(defaultSpTensorIndiceBitWidth), shape, format); - auto sp_tensor_type = SparseTensorType::get(builder.getContext(), element_type, builder.getIntegerType(defaultSpTensorIndiceBitWidth), shape, format); + // lhs_tensor = declOp; + // } + // else if (auto SparseTensorT = dyn_cast(rhs_tensor.getType())) + // { + + // ArrayRef format = SparseTensorT.getFormat(); + // mlir::ShapedType shapedT = mlir::cast(rhs_tensor.getType()); + // mlir::Type element_type = shapedT.getElementType(); + // return_type = SparseTensorType::get(builder.getContext(), element_type, builder.getIntegerType(defaultSpTensorIndiceBitWidth), shape, format); + // auto sp_tensor_type = SparseTensorType::get(builder.getContext(), element_type, builder.getIntegerType(defaultSpTensorIndiceBitWidth), shape, format); - /// BoolAttr is true to speficy SparseTensorDeclOp is for temporaries - lhs_tensor = builder.create(loc(transpose.loc()), sp_tensor_type, indices, builder.getBoolAttr(true)); - comet_debug() << "MLIRGen SparseTensorDeclaration creation\n"; - comet_vdump(lhs_tensor); - } + // /// BoolAttr is true to speficy SparseTensorDeclOp is for temporaries + // lhs_tensor = builder.create(loc(transpose.loc()), sp_tensor_type, indices, builder.getBoolAttr(true)); + // comet_debug() << "MLIRGen SparseTensorDeclaration creation\n"; + // comet_vdump(lhs_tensor); + // } comet_debug() << " create TransposeOp\n"; mlir::ShapedType shapedT = mlir::cast(rhs_tensor.getType()); mlir::Value t = builder.create(loc(transpose.loc()), mlir::RankedTensorType::get(shape, shapedT.getElementType()), rhs_tensor, all_labels_val, affineMapArrayAttr); - builder.create(loc(transpose.loc()), t, lhs_tensor); comet_vdump(t); return t; @@ -1928,7 +1915,7 @@ namespace LabeledTensorExprAST *lhsLabeledTensorExprAST = llvm::cast(tensor_op->getLHS()); auto call_res = mlirGen(*call); auto lhs_tensor = symbolTable.lookup(lhsLabeledTensorExprAST->getTensorName()); - builder.create(loc(tensor_op->loc()), call_res, lhs_tensor); + symbolTable.insert(lhsLabeledTensorExprAST->getTensorName(), call_res); } /// TODO: put check here, if the user mis-spells something... @@ -2459,7 +2446,8 @@ namespace /// Build the MLIR op `ta.constant`. This invokes the `DenseConstantOp::build` /// method. auto denseConst = builder.create(loc, lhs_labeledtensor.getType(), dataAttribute); - builder.create(loc, denseConst, lhs_labeledtensor); + symbolTable.insert(tensor_name, denseConst); + // builder.create(loc, denseConst, lhs_labeledtensor); return mlir::success(); } diff --git a/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp b/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp index b14818df..782861da 100644 --- a/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp +++ b/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp @@ -251,7 +251,7 @@ mlir::LogicalResult generalIndexOperationRewrite( ArrayAttr indexing_maps = cast(mult_op.getIndexingMaps()); for(auto [index, v]: enumerate(cast(indexing_maps[2]).getValue().getResults())) { - if(!cast(rhs1_tensor.getType()).isDynamicDim(index)) + if(!shapeT.isDynamicDim(index)) { continue; } @@ -272,6 +272,7 @@ mlir::LogicalResult generalIndexOperationRewrite( continue; } } + assert(dims.size() == shapeT.getNumDynamicDims()); if(auto spTensorT = dyn_cast(op->getResultTypes()[0])) { @@ -280,6 +281,10 @@ mlir::LogicalResult generalIndexOperationRewrite( else if (auto tensorT = dyn_cast(op->getResultTypes()[0])) { lhs_tensor = rewriter.create(loc, shapeT, ValueRange(dims)); + rewriter.create( + loc, + lhs_tensor, + rewriter.getZeroAttr(tensorT.getElementType())); // initialize the tensor to zero // lhs_tensor = rewriter.create(loc, shapeT, ValueRange(dims)); } @@ -526,7 +531,7 @@ void LowerTensorAlgebraToIndexTreePass::runOnOperation() mlir::ConversionTarget target(getContext()); target.addLegalDialect(); - target.addLegalOp(); + target.addLegalOp(); target.addIllegalOp(); diff --git a/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp b/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp index a75eea87..cbeea9e3 100644 --- a/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp +++ b/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp @@ -222,27 +222,29 @@ namespace /// There are tensors for transpose operation: input and output tensors unsigned int tensors_num = 2; - tensorAlgebra::TensorSetOp setOp; Value lhs; - if (isa(inputType)) + if (auto tensorT = dyn_cast(inputType)) { /// for dense comet_debug() << "Dense transpose\n"; - - // auto inputTensorLoadOp = cast(op->getOperand(0).getDefiningOp()); - // auto inputMemref = inputTensorLoadOp.getMemref(); - - for (auto u : op.getOperation()->getResult(0).getUsers()) + SmallVector dims; + for(auto [index, perm]: llvm::enumerate(allPerms[1])) /// for the output tensor, we need to get the dims from the permute order { - if (isa(u)) + if(perm < 0) { - setOp = cast(u); - Value dstTensor = u->getOperand(1); - lhs = dstTensor; + comet_debug() << "Invalid permutation found in transpose\n"; + return failure(); + } + if(tensorT.isDynamicDim(index)) + { + auto dim = rewriter.create(loc, inputTensor, perm); + dims.push_back(dim); } } + lhs = rewriter.create(loc, op.getResult().getType(), dims); + comet_vdump(lhs); // auto outputMemref = lhs.getDefiningOp()->getOperand(0); auto la_transpose = rewriter.create(loc, inputTensor, lhs, llvm::ArrayRef(allPerms[1])); @@ -279,16 +281,6 @@ namespace std::vector> allocs_for_sparse_tensors{tensors_num}; std::vector tensors = {op.getOperation()->getOperand(0)}; - for (auto u : op.getOperation()->getResult(0).getUsers()) - { - if (isa(u)) - { - setOp = cast(u); - mlir::Value lhs = setOp->getOperand(1); /// dest tensor is the 2nd - tensors.push_back(lhs); - } - } - for (unsigned int n = 0; n < tensors_num; n++) { auto tensor_rank = cast(tensors[n].getType()).getRank(); @@ -399,7 +391,6 @@ namespace rewriter.create(loc, func_name, SmallVector{}, ValueRange(allInputs) ); - rewriter.eraseOp(setOp); rewriter.eraseOp(op); return success(); @@ -680,6 +671,7 @@ void LowerTensorAlgebraToSCFPass::runOnOperation() bufferization::BufferizationDialect>(); target.addLegalDialect(); + target.addLegalOp(); target.addIllegalOp itree_args; - for (Value arg : itree.getInputs()) { + for (Value arg : itree->getResults()) { itree_args.push_back(arg); } @@ -502,7 +502,6 @@ IndexTreeOp createNewITree(IndexTreeOp host_itree, mlir::Location &loc) { OpBuilder::InsertionGuard guard(rewriter); - rewriter.setInsertionPoint(host_itree); /// Create the new itree llvm::SmallVector locs; // for (Value arg : itree_arguments) { @@ -977,6 +976,7 @@ void createITree( /// Create the new itree. llvm::SmallVector itree_to_new_compute_op; + rewriter.setInsertionPointAfter(itree_list.back()); IndexTreeOp new_itree = createNewITree(host_itree, tree_types, // itree_arguments, @@ -989,7 +989,6 @@ void createITree( context, rewriter, loc); - /// Fuse other itrees to the new itree fuseITrees(new_itree, num_itrees, @@ -1012,26 +1011,23 @@ void createITree( /// Update the usage of results of itree /// Erase uses of output of kernels other than the last kernel - for (uint32_t tree_i = 0; tree_i < num_itrees - 1; ++tree_i) { - for (auto result : itree_list[tree_i].getResults()) { - for (auto user : result.getUsers()) { - if (auto set_op = dyn_cast(user)) { - rewriter.eraseOp(set_op); - } - } - } - } /// Replace the use of output of the last kernel - uint32_t new_r_i = 0; - uint32_t tree_i = num_itrees - 1; - for (auto old_result : itree_list[tree_i].getResults()) { - rewriter.replaceAllUsesWith(old_result, new_itree.getResult(new_r_i++)); - } - assert(new_r_i == 1 && "Expect to output only one tensor"); - - /// Delete the old itrees - for (auto itree : itree_list) { - rewriter.eraseOp(itree); + // uint32_t new_r_i = 0; + // uint32_t tree_i = num_itrees - 1; + for(auto it : enumerate(ArrayRef(itree_list))) + { + auto index = itree_list.size() - 1 - it.index(); + auto itree_to_remove = it.value(); + rewriter.replaceUsesWithIf( + itree_to_remove.getResults(), + new_itree.getResults()[index], + [&](OpOperand &use) { + auto user = use.getOwner(); + auto ancestor = itree_to_remove->getBlock()->findAncestorOpInBlock(*user); + return (ancestor && new_itree->isBeforeInBlock(ancestor)); + }); + + rewriter.replaceOp(itree_to_remove, new_itree->getOperand(index)); } } From c52f9d2565cef58c221a863b48d674bcd43e9249 Mon Sep 17 00:00:00 2001 From: Polykarpos Thomadakis Date: Sat, 5 Apr 2025 13:55:18 -0700 Subject: [PATCH 03/11] Fixed bug in the creation of temporary traspose output. --- lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp b/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp index cbeea9e3..f497ed8c 100644 --- a/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp +++ b/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp @@ -231,14 +231,9 @@ namespace SmallVector dims; for(auto [index, perm]: llvm::enumerate(allPerms[1])) /// for the output tensor, we need to get the dims from the permute order { - if(perm < 0) + if(tensorT.isDynamicDim(perm)) { - comet_debug() << "Invalid permutation found in transpose\n"; - return failure(); - } - if(tensorT.isDynamicDim(index)) - { - auto dim = rewriter.create(loc, inputTensor, perm); + auto dim = rewriter.create(loc, inputTensor, perm); dims.push_back(dim); } } From 9e8ebdb78ad67d0aaef4ebb9a45b261f5f1ebc11 Mon Sep 17 00:00:00 2001 From: Polykarpos Thomadakis Date: Sat, 5 Apr 2025 15:30:34 -0700 Subject: [PATCH 04/11] [WIP] All tests pass. Need to fix factorization opt --- frontends/comet_dsl/mlir/MLIRGen.cpp | 57 +++++++-- .../TensorAlgebraToIndexTree.cpp | 51 --------- .../TensorAlgebraToSCF/EarlyLowering.cpp | 1 - .../TensorAlgebraToSCF/LateLowering.cpp | 3 +- .../TensorAlgebraToSCF/TensorAlgebraToSCF.cpp | 108 ++++++++---------- .../Transforms/TensorDeclLowering.cpp | 54 +-------- 6 files changed, 104 insertions(+), 170 deletions(-) diff --git a/frontends/comet_dsl/mlir/MLIRGen.cpp b/frontends/comet_dsl/mlir/MLIRGen.cpp index 9fd09f76..1c513562 100644 --- a/frontends/comet_dsl/mlir/MLIRGen.cpp +++ b/frontends/comet_dsl/mlir/MLIRGen.cpp @@ -48,6 +48,7 @@ #include "llvm/ADT/ArrayRef.h" #include "llvm/ADT/STLExtras.h" #include "llvm/ADT/ScopedHashTable.h" +#include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/raw_ostream.h" #include @@ -479,17 +480,55 @@ namespace mlir::StringAttr opAttr = builder.getStringAttr(op); mlir::RankedTensorType returnDataType; - if(mlir::cast(lhs.getType()).getShape() != mlir::cast(rhs.getType()).getShape()) + auto lhsShapedType = llvm::dyn_cast(lhs.getType()); + auto rhsShapedType = llvm::dyn_cast(rhs.getType()); + if(!lhsShapedType && !rhsShapedType) { - returnDataType = mlir::cast(lhs.getType()); - auto bcastRhs = builder.create(location, returnDataType, mlir::cast(rhs.getDefiningOp()).getValueAttr()); + mlir::Type elementType = builder.getF64Type(); + returnDataType = mlir::RankedTensorType::get(1, elementType); + } + else if (lhsShapedType && !rhsShapedType) + { + returnDataType = lhsShapedType; + SmallVector dims; + for(int i = 0; i < lhsShapedType.getRank(); ++i) + { + if(lhsShapedType.isDynamicDim(i)) + { + auto dim = builder.create(location, lhs, i); + dims.push_back(dim); + } + } + auto bcastRhs = builder.create(location, rhs, returnDataType, dims); comet_vdump(bcastRhs); - rhs.replaceAllUsesWith(bcastRhs); + rhs.replaceAllUsesExcept(bcastRhs, bcastRhs); // replace all uses of rhs with bcastRhs except the defining op rhs = bcastRhs; } - else { - mlir::Type elementType = builder.getF64Type(); - returnDataType = mlir::RankedTensorType::get(1, elementType); + else if(!lhsShapedType && rhsShapedType) + { + returnDataType = rhsShapedType; + SmallVector dims; + for(int i = 0; i < rhsShapedType.getRank(); ++i) + { + if(rhsShapedType.isDynamicDim(i)) + { + auto dim = builder.create(location, rhs, i); + dims.push_back(dim); + } + } + auto bcastLhs = builder.create(location, lhs, returnDataType, dims); + comet_vdump(bcastLhs); + lhs.replaceAllUsesExcept(bcastLhs, bcastLhs); // replace all uses of lhs with bcastLhs except the defining op + lhs = bcastLhs; + } + else if (lhsShapedType && rhsShapedType) + { + if (lhsShapedType.getElementType() != rhsShapedType.getElementType()) + { + comet_debug() << "ERROR: mismatched element types for binary operation\n"; + return nullptr; + } + returnDataType = lhsShapedType; } comet_vdump(rhs); comet_vdump(lhs); @@ -497,7 +536,7 @@ namespace comet_vdump(theOutput); auto scalarOp = builder.create(location, returnDataType, rhs, lhs, opAttr); comet_vdump(scalarOp); - symbolTable.insert(out_format, scalarOp); + // symbolTable.insert(out_format, scalarOp); /// the value returned here will be used in subsequent ops. /// for example, in the code below, 'g' should be returned. @@ -1213,7 +1252,6 @@ namespace auto lhsName = cast(*tensor_op.getLHS()) .getTensorName(); symbolTable.insert(lhsName, rhs); - // auto ret_op = builder.create(loc(tensor_op.loc()), rhs, lhs); // ret_op.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); return rhs; @@ -2447,7 +2485,6 @@ namespace /// method. auto denseConst = builder.create(loc, lhs_labeledtensor.getType(), dataAttribute); symbolTable.insert(tensor_name, denseConst); - // builder.create(loc, denseConst, lhs_labeledtensor); return mlir::success(); } diff --git a/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp b/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp index 782861da..80905332 100644 --- a/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp +++ b/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp @@ -133,57 +133,6 @@ bool check_chosen_operations(const std::vector> &allPerms, return false; } -Value getRealLhs(Operation *op) -{ - assert(isa(op) || isa(op) || isa(op) || isa(op)); - Operation *firstUser = nullptr; - for (auto user : op->getResult(0).getUsers()) - { - firstUser = user; - break; - } - - comet_pdump(firstUser); - assert(isa(firstUser)); - TensorSetOp setOp = cast(firstUser); - return setOp.getOperand(1); -} - -Value getRealRhs(Value val) -{ - /// this will return set_op for transpose, but messes up getUsers() or subsequent calls to it. - - /// TODO(gkestor): need to find out why user set_op is not showing up in users of TransposeOp - /// from the for loop below. once resolved, remove getNextNode(). - /// Operation *firstUser; - /// for (auto user : op->getResult(0).getUsers()) - //{ - /// firstUser = user; - /// break; - //} - if(Operation* op = val.getDefiningOp()) - { - Operation *firstUser = op->getNextNode(); - comet_pdump(firstUser); - - if (isa(op)) - { - if (isa(firstUser)) - { - TensorSetOp setOp = cast(firstUser); - return setOp.getOperand(1); - } - else - { - llvm::errs() << "ERROR: Transpose has no set_op after it!\n"; - } - - } - } - - return val; -} - // void buildDefUseInfo(UnitExpression *e) // { // auto lhs = e->getLHS(); diff --git a/lib/Conversion/TensorAlgebraToSCF/EarlyLowering.cpp b/lib/Conversion/TensorAlgebraToSCF/EarlyLowering.cpp index 28eff319..e62aa477 100644 --- a/lib/Conversion/TensorAlgebraToSCF/EarlyLowering.cpp +++ b/lib/Conversion/TensorAlgebraToSCF/EarlyLowering.cpp @@ -251,7 +251,6 @@ namespace tensorAlgebra::TensorFillOp, tensorAlgebra::GetTimeOp, tensorAlgebra::PrintElapsedTimeOp, - tensorAlgebra::TensorSetOp, tensorAlgebra::IndexLabelOp, tensorAlgebra::DenseConstantOp, tensorAlgebra::TensorMultOp, diff --git a/lib/Conversion/TensorAlgebraToSCF/LateLowering.cpp b/lib/Conversion/TensorAlgebraToSCF/LateLowering.cpp index 69e6a0ea..052d8d53 100644 --- a/lib/Conversion/TensorAlgebraToSCF/LateLowering.cpp +++ b/lib/Conversion/TensorAlgebraToSCF/LateLowering.cpp @@ -338,8 +338,7 @@ void LateLoweringPass::runOnOperation() /// PrintOp Lowering insert function call, so mark some operations as a legal Operation target.addLegalOp(); /// Now that the conversion target has been defined, we just need to provide diff --git a/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp b/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp index f497ed8c..706d7717 100644 --- a/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp +++ b/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp @@ -43,6 +43,7 @@ #include "mlir/Support/LLVM.h" #include "mlir/Support/LogicalResult.h" #include "llvm/ADT/SmallVector.h" +#include "llvm/Support/Casting.h" #include using namespace mlir; @@ -87,40 +88,12 @@ namespace DenseElementsAttr constantValue = op.getValue(); Location loc = op.getLoc(); - tensorAlgebra::TensorSetOp setnewop; - bool user_setOp = false; - for (auto u : op.getOperation()->getResult(0).getUsers()) - { - if (isa(u)) - { - setnewop = cast(u); - user_setOp = true; - } - } - /// When lowering the constant operation, we allocate and assign the constant /// values to a corresponding memref allocation. auto tensorType = cast(op.getType()); auto memRefType = convertTensorToMemRef(tensorType); - comet_debug() << "User_setop: " << user_setOp << "/n"; - Value alloc; - if (user_setOp) - { - if (auto toTensor = dyn_cast(setnewop.getOperand(1).getDefiningOp())) - { - auto alloc_op = cast(toTensor->getOperand(0).getDefiningOp()); - comet_vdump(alloc_op); - alloc = alloc_op; - } - else { - alloc = rewriter.create(loc, memRefType); - } - } - else - { - alloc = rewriter.create(loc, memRefType); - } + Value alloc = rewriter.create(loc, memRefType); /// We will be generating constant indices up-to the largest dimension. /// Create these constants up-front to avoid large amounts of redundant @@ -537,32 +510,7 @@ namespace // auto lh = op->getOperand(1).getType(); [[maybe_unused]] auto f64Type = rewriter.getF64Type(); - Value const_index_0 = rewriter.create(loc, 0); comet_vdump(const_index_0); - std::vector alloc_zero_loc = {const_index_0}; - - Value res; - bool res_comes_from_setop = false; - for (auto u : op.getOperation()->getResult(0).getUsers()) - { - comet_debug() << "Users:\n"; - comet_pdump(u); - if (isa(u)) - { - // u->dump(); - // u->getBlock()->dump(); - res = cast(u).getOperation()->getOperand(1); - // (++res.getUsers().begin())->dump(); - if(!res.getUsers().empty() && isa(*(++res.getUsers().begin()))) - { - res = cast(*(++res.getUsers().begin())).getRhs(); - } - res_comes_from_setop = true; - break; - } - } - - assert(res_comes_from_setop && "SetOp is needed to assign the scalar operation result to final variable"); comet_debug() << "Scalar lowering Final step\n"; comet_debug() << "RHS, LHS and result:\n"; @@ -574,22 +522,66 @@ namespace auto arith_op_attr = op.getOpAttr(); std::string op_attr(arith_op_attr.getValue()); comet_debug() << "aritmetic op: " << op_attr << "\n"; + auto areTensors = dyn_cast(rhs.getType()) && dyn_cast(lhs.getType()); + Value res; + if(areTensors) + { + SmallVector dims; + auto lhs_type = dyn_cast(lhs.getType()); + for(int64_t i = 0; i < lhs_type.getRank(); i++) + { + if(lhs_type.isDynamicDim(i)) + { + auto dim = rewriter.create(loc, lhs, i); + dims.push_back(dim); + } + } + res = rewriter.create(loc, rhs.getType(), dims); + } Value res_val; if (op_attr.compare("+") == 0) { - res_val = rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)).getResultTensors()[0]; + if(areTensors) + { + res_val = rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)).getResultTensors()[0]; + } + else + { + res_val = rewriter.create(loc, lhs, rhs).getResult(); + } } else if (op_attr.compare("-") == 0) { - res_val = rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)).getResultTensors()[0]; + if(areTensors) + { + res_val = rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)).getResultTensors()[0]; + } + else + { + res_val = rewriter.create(loc, lhs, rhs).getResult(); + } } else if (op_attr.compare("*") == 0) { - res_val = rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)).getResultTensors()[0]; + if(areTensors) + { + res_val = rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)).getResultTensors()[0]; + } + else + { + res_val = rewriter.create(loc, lhs, rhs).getResult(); + } } else if (op_attr.compare("/") == 0) { - res_val = rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)).getResultTensors()[0]; + if(areTensors) + { + res_val = rewriter.create(loc, ValueRange{lhs, rhs}, ValueRange(res)).getResultTensors()[0]; + } + else + { + res_val = rewriter.create(loc, lhs, rhs).getResult(); + } } else { diff --git a/lib/Dialect/TensorAlgebra/Transforms/TensorDeclLowering.cpp b/lib/Dialect/TensorAlgebra/Transforms/TensorDeclLowering.cpp index 49ad700c..13a34906 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/TensorDeclLowering.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/TensorDeclLowering.cpp @@ -314,36 +314,17 @@ Value insertSparseTensorDeclOp(PatternRewriter & rewriter, comet_debug() << " Users:\n"; comet_pdump(u); - if (isa(u) || - (isa(u) && - isa(cast(u).getOperand(0).getDefiningOp()))) + if (isa(u)) //|| { - if (!isa(u)) - { - comet_debug() << "User of sparse tensor is a set Operation. Src of setOp is transpose\n"; - /// Set the insertion point before its user - rewriter.setInsertionPoint(cast(u).getOperand(0).getDefiningOp()); - } - else - { - comet_debug() << "User of sparse tensor is transpose operation\n"; - /// Set the insertion point before its user - rewriter.setInsertionPoint(u); - } + comet_debug() << "User of sparse tensor is transpose operation\n"; + /// Set the insertion point before its user + rewriter.setInsertionPoint(u); /// Get the freeIndices of the sparse input tensor /// Check the dimension size, if it is integer, format is dense and get dim_size /// If it is ?, get the sparse input and get the definition, and the freeindex, /// tensorAlgebra::TransposeOp transpose_op = cast(u); - tensorAlgebra::TransposeOp transpose_op; - if (isa(u)) - { - transpose_op = cast(u); - } - else - { - transpose_op = cast(cast(u).getOperand(0).getDefiningOp()); - } + tensorAlgebra::TransposeOp transpose_op = cast(u); ArrayAttr indexMaps = transpose_op.getIndexingMaps(); comet_debug() << " we get the indexMaps\n"; @@ -354,16 +335,6 @@ Value insertSparseTensorDeclOp(PatternRewriter & rewriter, comet_debug() << " "; comet_vdump(src_input); mlir::Value dst_input; - for (auto u : op.getOperation()->getResult(0).getUsers()) - { - comet_debug() << " "; - comet_pdump(u); - if (isa(u)) - { - dst_input = u->getOperand(1); /// dest tensor is the 2nd - comet_vdump(dst_input); - } - } /// If in COO format, for every dimension, different dimensions are std::vector dstIndexLocInSrcVec; @@ -643,7 +614,7 @@ Value insertSparseTensorDeclOp(PatternRewriter & rewriter, for (auto u : op->getUsers()) { comet_pdump(u); - if (isa(u) || isa(u)) + if (isa(u)) is_filled = true; } @@ -730,15 +701,6 @@ Value insertSparseTensorDeclOp(PatternRewriter & rewriter, isOutputTensor = true; } } - else if (isa(u1)) - { - comet_debug() << " used in ta.set op\n"; - auto p = u1->getOperand(1); - if(p == op) - { - isOutputTensor = true; - } - } else if (isa(u1)) { comet_debug() << " used in it.Operand op\n"; @@ -1272,7 +1234,6 @@ Value insertSparseTensorDeclOp(PatternRewriter & rewriter, tensorAlgebra::TensorFillOp, tensorAlgebra::GetTimeOp, tensorAlgebra::PrintElapsedTimeOp, - tensorAlgebra::TensorSetOp, tensorAlgebra::SparseOutputTensorDeclOp, tensorAlgebra::TempSparseOutputTensorDeclOp, tensorAlgebra::TensorMultOp, @@ -1326,7 +1287,6 @@ Value insertSparseTensorDeclOp(PatternRewriter & rewriter, tensorAlgebra::SparseTensorConstructOp, tensorAlgebra::SparseOutputTensorDeclOp, tensorAlgebra::TempSparseOutputTensorDeclOp, - tensorAlgebra::TensorSetOp, tensorAlgebra::DenseTensorDeclOp, tensorAlgebra::IndexLabelOp, tensorAlgebra::TensorSortOp, @@ -1374,7 +1334,6 @@ Value insertSparseTensorDeclOp(PatternRewriter & rewriter, tensorAlgebra::TransposeOp, tensorAlgebra::TensorFillOp, tensorAlgebra::SparseTensorConstructOp, - tensorAlgebra::TensorSetOp, tensorAlgebra::DenseTensorDeclOp, tensorAlgebra::SparseOutputTensorDeclOp, tensorAlgebra::IndexLabelOp, @@ -1434,7 +1393,6 @@ Value insertSparseTensorDeclOp(PatternRewriter & rewriter, tensorAlgebra::TransposeOp, tensorAlgebra::TensorFillOp, tensorAlgebra::SparseTensorConstructOp, - tensorAlgebra::TensorSetOp, tensorAlgebra::DenseTensorDeclOp, tensorAlgebra::IndexLabelOp, tensorAlgebra::DenseConstantOp, From 54495be7ac490eb7e47be21987ff50786000bec8 Mon Sep 17 00:00:00 2001 From: Polykarpos Thomadakis Date: Sat, 5 Apr 2025 16:29:16 -0700 Subject: [PATCH 05/11] [#93] Completely removed TensorSetOp operation. Need to make sure it does not affect performance --- frontends/comet_dsl/mlir/MLIRGen.cpp | 4 +- .../comet/Dialect/TensorAlgebra/IR/TAOps.td | 16 - .../TensorAlgebraToSCF/TensorAlgebraToSCF.cpp | 26 +- lib/Dialect/TensorAlgebra/CMakeLists.txt | 2 - .../Transforms/CheckImplicitTensorDecls.cpp | 211 ----- .../TensorAlgebra/Transforms/Passes.cpp | 37 +- .../TensorAlgebra/Transforms/TCtoTTGT.cpp | 785 ------------------ .../TensorAlgebra/Transforms/TCtoTTGTDyn.cpp | 3 +- 8 files changed, 25 insertions(+), 1059 deletions(-) delete mode 100644 lib/Dialect/TensorAlgebra/Transforms/CheckImplicitTensorDecls.cpp delete mode 100644 lib/Dialect/TensorAlgebra/Transforms/TCtoTTGT.cpp diff --git a/frontends/comet_dsl/mlir/MLIRGen.cpp b/frontends/comet_dsl/mlir/MLIRGen.cpp index 1c513562..8c267263 100644 --- a/frontends/comet_dsl/mlir/MLIRGen.cpp +++ b/frontends/comet_dsl/mlir/MLIRGen.cpp @@ -1247,7 +1247,7 @@ namespace return nullptr; comet_debug() << " get rhs\n"; - auto tens_beta = tensor_op.getBeta(); + // auto tens_beta = tensor_op.getBeta(); auto lhsName = cast(*tensor_op.getLHS()) .getTensorName(); @@ -1645,7 +1645,6 @@ namespace } } - mlir::Type return_type = getType(shape); /// Create Tensor Declarations Ops and populate formats (for lhs) // mlir::Value lhs_tensor; @@ -1952,7 +1951,6 @@ namespace { LabeledTensorExprAST *lhsLabeledTensorExprAST = llvm::cast(tensor_op->getLHS()); auto call_res = mlirGen(*call); - auto lhs_tensor = symbolTable.lookup(lhsLabeledTensorExprAST->getTensorName()); symbolTable.insert(lhsLabeledTensorExprAST->getTensorName(), call_res); } /// TODO: put check here, if the user mis-spells something... diff --git a/include/comet/Dialect/TensorAlgebra/IR/TAOps.td b/include/comet/Dialect/TensorAlgebra/IR/TAOps.td index 148dc97c..b9d1c50e 100644 --- a/include/comet/Dialect/TensorAlgebra/IR/TAOps.td +++ b/include/comet/Dialect/TensorAlgebra/IR/TAOps.td @@ -539,22 +539,6 @@ def TensorDimOp : TA_Op<"dim", [Pure]>{ //let hasVerifier = 1; } - -def TensorSetOp : TA_Op<"set_op">{ - - let summary = ""; - let description = [{ - }]; - - let arguments = (ins AnyTypeOf<[TA_AnyTensor,F64MemRef]>:$lhs, - AnyTypeOf<[TA_AnyTensor,F64MemRef]>:$rhs); - - //let hasCanonicalizer = 1; - - //TODO(gkestor): add verifier - //let hasVerifier = 1; -} - def TransposeOp : TA_Op<"transpose",[Pure]> { let summary = "transpose operation"; diff --git a/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp b/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp index 706d7717..697660c2 100644 --- a/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp +++ b/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp @@ -594,26 +594,6 @@ namespace } }; /// ScalarOpsLowering -class ConvertSetOp : public OpConversionPattern { - using OpConversionPattern::OpConversionPattern; - - LogicalResult - matchAndRewrite(TensorSetOp op, OpAdaptor adaptor, - ConversionPatternRewriter &rewriter) const final { - - auto opAdaptor = llvm::cast(adaptor); - Value lhs = opAdaptor.getLhs(); - Value rhs = opAdaptor.getRhs(); - rewriter.replaceUsesWithIf(rhs, lhs, [&](OpOperand& use) { - auto user = use.getOwner(); - auto ancestor = op->getBlock()->findAncestorOpInBlock(*user); - return (ancestor && op->isBeforeInBlock(ancestor)); - }); - rewriter.eraseOp(op); - return success(); - } -}; - } /// end anonymous namespace. /// This is a partial lowering to linear algebra of the tensor algebra operations that are @@ -662,8 +642,7 @@ void LowerTensorAlgebraToSCFPass::runOnOperation() target.addIllegalOp(); + tensorAlgebra::DenseConstantOp>(); /// Now that the conversion target has been defined, we just need to provide /// the set of patterns that will lower the TA operations. @@ -671,8 +650,7 @@ void LowerTensorAlgebraToSCFPass::runOnOperation() patterns.insert(&getContext()); + ConstantOpLowering>(&getContext()); /// With the target and rewrite patterns defined, we can now attempt the /// conversion. The conversion will signal failure if any of our `illegal` /// operations were not converted successfully. diff --git a/lib/Dialect/TensorAlgebra/CMakeLists.txt b/lib/Dialect/TensorAlgebra/CMakeLists.txt index 209b7767..c12b85ff 100644 --- a/lib/Dialect/TensorAlgebra/CMakeLists.txt +++ b/lib/Dialect/TensorAlgebra/CMakeLists.txt @@ -3,11 +3,9 @@ add_llvm_library(COMETTensorAlgebraDialect Transforms/Transforms.cpp Transforms/LinalgTransforms.cpp - Transforms/TCtoTTGT.cpp Transforms/TCtoTTGTDyn.cpp Transforms/Passes.cpp - Transforms/CheckImplicitTensorDecls.cpp Transforms/TensorDeclLowering.cpp Transforms/WorkspaceOptimizations.cpp diff --git a/lib/Dialect/TensorAlgebra/Transforms/CheckImplicitTensorDecls.cpp b/lib/Dialect/TensorAlgebra/Transforms/CheckImplicitTensorDecls.cpp deleted file mode 100644 index 9f19069e..00000000 --- a/lib/Dialect/TensorAlgebra/Transforms/CheckImplicitTensorDecls.cpp +++ /dev/null @@ -1,211 +0,0 @@ -//===- CheckImplicitTensorDecl.cpp - check if it is needed to add tensor declarations introduced by compound expressions------------------===// -// -// Copyright 2022 Battelle Memorial Institute -// -// Redistribution and use in source and binary forms, with or without modification, -// are permitted provided that the following conditions are met: -// -// 1. Redistributions of source code must retain the above copyright notice, this list of conditions -// and the following disclaimer. -// -// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions -// and the following disclaimer in the documentation and/or other materials provided with the distribution. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED -// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR -// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, -// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE -// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, -// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -// -//===----------------------------------------------------------------------===// -// -/// This file implements a pass that adds temporary tensor declaration introduced by compound expressions in the COMET DSL -//===----------------------------------------------------------------------===// - -#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" -#include "comet/Dialect/TensorAlgebra/Passes.h" -#include "comet/Dialect/Utils/Utils.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" - -#include -#include -#include -#include - -#include "llvm/Support/Debug.h" - -using namespace mlir; -using namespace mlir::tensorAlgebra; - -#define DEBUG_TYPE "check-implicit-tensor-decl" - -// *********** For debug purpose *********// -// #define COMET_DEBUG_MODE -#include "comet/Utils/debug.h" -#undef COMET_DEBUG_MODE -// *********** For debug purpose *********// - -using namespace mlir; - -namespace -{ - /// Add tensor decl ops for tmp result - struct TensorAlgebraCheckImplicitTensorDeclPass - : public PassWrapper> - { - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TensorAlgebraCheckImplicitTensorDeclPass) - void runOnOperation() override; - }; -} /// namespace - -template -bool isNeedTensorDecl(t op) -{ - bool isUsedInSetSource = true; - mlir::Value result = op->getResult(0); - comet_debug() << " "; - comet_vdump(result); - for (auto u1 : result.getUsers()) - { - comet_debug() << " "; - comet_pdump(u1); - /// If not used as source tensor of set_op, it is tmp result - if (isa(u1)) - { - auto p = cast(u1).getOperation(); - if(result == p->getOperand(0)) - { - isUsedInSetSource = false; - } - } - } - return isUsedInSetSource; -} - -template -void addTensorDecl(t op) -{ - OpBuilder builder(op); - auto location = op.getLoc(); - op->getOperands(); - std::vector lbls_value; - mlir::Value ret_value; - ArrayAttr imaps = op.getIndexingMaps(); - - /// Retrieve the size of the tensor based on the affinity maps - /// 1. Check the affinity map results of the LHS e.g. - /// 2. Find which dimensions match on the RHS(1,2) e.g. (d0,d1,d2) -> (d0, d1), (d0,d1,d2) -> (d1, d2) - /// 3. The result will be the dimension of the matching indices - /// e.g For matmulop RHS1: (d0,d1,d2) -> (d0, d1), RHS2: (d0,d1,d2) -> (d1, d2), LHS : (d0,d1,d2) -> (d0,d2) - /// So we get (dim(RHS1, 0), dim(RHS2, 1)) because of d0, d2 respectively - auto res_map = cast(imaps[imaps.size() - 1]).getValue(); - for (auto v : res_map.getResults()) - { - for (size_t i = 0; i < imaps.size() - 1; i++) - { - auto map = cast(imaps[i]).getValue(); - if (auto pos = map.getResultPosition(v)) - { - lbls_value.push_back(builder.create(location, op->getOperand(i), *pos)); - } - } - } - ret_value = op.getOperation()->getResult(0); - - // mlir::ArrayAttr opFormatsArrayAttr = op.getFormats(); - // unsigned int i = opFormatsArrayAttr.size() - 1; - // std::string ret_format_local(cast(opFormatsArrayAttr[i]).getValue()); - // ret_format = ret_format_local; - - mlir::Value itensor; - if (isa(ret_value.getType())) - { - itensor = builder.create(location, ret_value.getType(), lbls_value); - builder.create(location, itensor, builder.getF64FloatAttr(0)); - } - else if (isa(ret_value.getType())) - { - /// It is a temporal tensor declaration generated by compound expressions, BoolAttr is true - /// to identify SparseTensorDeclOp is for temporaries - itensor = builder.create(location, ret_value.getType(), lbls_value, true); - } - else - { - comet_debug() << "Error: cannot handle the tensor type for TensorDeclLowering\n"; - op.emitOpError("cannot handle the tensor type for TensorDeclLowering"); - return; - } - comet_debug() << "PreLowering SparseTensorDeclaration creation\n"; - comet_vdump(itensor); - op.replaceAllUsesWith(itensor); - - builder.setInsertionPointAfter(op); -#ifdef DEBUG_MODE_ADDTEMPTENSORDECLARATION - auto setop = builder.create(location, ret_value, itensor); - comet_debug() << " "; - comet_vdump(setop); - - comet_debug() << " "; - comet_vdump(ret_value); -#else - builder.create(location, ret_value, itensor); -#endif -} - -void TensorAlgebraCheckImplicitTensorDeclPass::runOnOperation() -{ - func::FuncOp func = getOperation(); - comet_debug() << "Before TensorAlgebraCheckImplicitTensorDeclPass\n"; - - /// The walker proceeds in post-order, but we need to process outer loops first - /// to control the number of outer parallel loops, so push candidate loops to - /// the front of a deque. - func.walk([&](mlir::Operation *cur_op) - { - comet_debug() << " find a transpose op\n"; - /// if the output is not used as a source tensor of a set op - /// Need to store use a sparse/dense tensor decl op to store the result - if (auto op = dyn_cast(cur_op)) - { - if (isNeedTensorDecl(op)) - { - addTensorDecl(op); - } - } - else if (auto op = dyn_cast(cur_op)) - { - if (isNeedTensorDecl(op)) - { - addTensorDecl(op); - } - } - else if (auto op = dyn_cast(cur_op)) - { - if (isNeedTensorDecl(op)) - { - addTensorDecl(op); - } - } - else if (auto op = dyn_cast(cur_op)) - { - if (isNeedTensorDecl(op)) - { - addTensorDecl(op); - } - } - else if (auto op = dyn_cast(cur_op)) - { - if (isNeedTensorDecl(op)) - { - addTensorDecl(op); - } - } }); - comet_debug() << "After TensorAlgebraCheckImplicitTensorDeclPass\n"; -} - -std::unique_ptr mlir::comet::createTensorAlgebraCheckImplicitTensorDeclPass() -{ - return std::make_unique(); -} diff --git a/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp b/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp index f7ecdc3c..fa509d1f 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp @@ -24,6 +24,7 @@ #include "comet/Dialect/Utils/Utils.h" #include "mlir/IR/BuiltinTypes.h" +#include "mlir/IR/ValueRange.h" #include "mlir/Pass/Pass.h" #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Arith/IR/Arith.h" @@ -31,8 +32,11 @@ #include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/Bufferization/IR/Bufferization.h" +#include "llvm/ADT/ArrayRef.h" +#include "llvm/ADT/STLExtras.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/Casting.h" +#include "llvm/Support/LogicalResult.h" #include #include @@ -62,7 +66,7 @@ namespace MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FindOptimalTCFactorizationPass) void runOnOperation() override; - void FindOptimalTCFactorization(tensorAlgebra::TensorSetOp op); + LogicalResult FindOptimalTCFactorization(tensorAlgebra::TensorMultOp op); }; /// class FindOptimalTCFactorizationPass } /// End anonymous namespace @@ -220,14 +224,12 @@ optimalOrder(ArrayRef inLTOps, Operation *outLTOp, return std::make_tuple(minResult, minSumLabels, minLHSTensorShapes); } -void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::TensorSetOp op) +mlir::LogicalResult FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::TensorMultOp op) { OpBuilder builder(op); comet_pdump(op); - auto operands = op->getOperands(); auto loc = op->getLoc(); - auto lhsOp = operands[0].getDefiningOp(); /// TensorMultOp - + auto lhsOp = op; std::vector MultOpsToRemove; std::vector LTOpsToRemove; @@ -240,11 +242,11 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T std::map> lblMaps; /// collect all operands from series of ta.tc ops - if (isa(lhsOp)) + if (isa(op)) { std::stack stack; - Value currValue = operands[0]; + Value currValue = op; comet_vdump(currValue); Operation *curr = currValue.getDefiningOp(); while ((curr && isa(curr)) || !stack.empty()) @@ -580,33 +582,34 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::T Value tcop = builder.create(loc, newType, newRhs1, newRhs2, all_labels, affineMapArrayAttr, SemiringAttr, MaskingAttr, nullptr); - tcop.getDefiningOp()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); - tcop.getDefiningOp()->setAttr("__beta__", builder.getF64FloatAttr(0.0)); comet_debug() << "New operation " << tcop << "\n"; newRhs1 = tcop; } - - mlir::tensorAlgebra::TensorSetOp newSetOp = builder.create(loc, newRhs1, operands[1]); - newSetOp->setAttr("__beta__", builder.getF64FloatAttr(0.0)); + comet_debug() << "are they previous multop\n"; - for (auto oldTcOp : MultOpsToRemove) + MultOpsToRemove.front()->replaceAllUsesWith(ValueRange(newRhs1)); /// replace the last multop with the new one + for (auto oldTcOp : ArrayRef(MultOpsToRemove).drop_front()) { - comet_debug() << "Calling removeAllUsers\n"; removeAllUsers(oldTcOp); } + return success(); } comet_debug() << "MulOpFactorization end\n"; - return; + return failure(); } void FindOptimalTCFactorizationPass::runOnOperation() { comet_debug() << " start FindOptimalTCFactorizationPass pass \n"; func::FuncOp func = getOperation(); + SmallVector tensorMultOps; - func.walk([&](tensorAlgebra::TensorSetOp op) - { FindOptimalTCFactorization(op); }); + do { + tensorMultOps.clear(); // clear the vector to find the next round of tensor mult ops + func.walk([&](tensorAlgebra::TensorMultOp op) + { tensorMultOps.push_back(op); }); + }while(!tensorMultOps.empty() && succeeded(FindOptimalTCFactorization(tensorMultOps.back()))); } void STCRemoveDeadOpsPass::runOnOperation() diff --git a/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGT.cpp b/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGT.cpp deleted file mode 100644 index fc06d88b..00000000 --- a/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGT.cpp +++ /dev/null @@ -1,785 +0,0 @@ -//===- TCtoTTGT.cpp ------===// -// -// Copyright 2022 Battelle Memorial Institute -// -// Redistribution and use in source and binary forms, with or without modification, -// are permitted provided that the following conditions are met: -// -// 1. Redistributions of source code must retain the above copyright notice, this list of conditions -// and the following disclaimer. -// -// 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions -// and the following disclaimer in the documentation and/or other materials provided with the distribution. -// -// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED -// WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR -// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, -// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE -// GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, -// WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE -// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -// -//===----------------------------------------------------------------------===// -// -// This file implements reformulation of tensor contraction operations as Transpose-Transpose-GEMM-Transpose -//===----------------------------------------------------------------------===// - -#include "comet/Dialect/TensorAlgebra/IR/TADialect.h" -#include "comet/Dialect/TensorAlgebra/Passes.h" -#include "comet/Dialect/Utils/Utils.h" - -#include "mlir/Dialect/Linalg/IR/Linalg.h" -#include "mlir/Dialect/Linalg/Transforms/Transforms.h" -#include "mlir/Dialect/Arith/IR/Arith.h" -#include "mlir/Dialect/Func/IR/FuncOps.h" -#include "mlir/Dialect/Bufferization/IR/Bufferization.h" - -#include -#include -#include -#include -#include - -#include "mlir/Dialect/MemRef/IR/MemRef.h" -#include "mlir/IR/BuiltinTypes.h" - -using namespace mlir; -using namespace mlir::linalg; -using namespace mlir::arith; -using namespace mlir::bufferization; - -using namespace mlir::tensorAlgebra; - -// *********** For debug purpose *********// -// #define COMET_DEBUG_MODE -#include "comet/Utils/debug.h" -#undef COMET_DEBUG_MODE -// *********** For debug purpose *********// - -const StringLiteral kLinalgTransformMarker = "__with_tiling__"; - -template -static bool arePermutations(const std::vector &vec1, - const std::vector &vec2) -{ - if (vec1.size() != vec2.size()) - { - return false; - } - std::vector taken(vec1.size(), false); - for (size_t i = 0; i < vec1.size(); i++) - { - auto it = std::find(vec2.begin(), vec2.end(), vec1[i]); - if (it == vec2.end()) - { - return false; - } - if (taken[std::distance(vec2.begin(), it)] == true) - { - return false; - } - taken[std::distance(vec2.begin(), it)] = true; - } - return true; -} - -/// Detect whether memref dims [dim, dim + extent) can be reshaped without -/// copies. -[[maybe_unused]] static bool isReshapableDimBand(unsigned dim, unsigned extent, - ArrayRef sizes, - ArrayRef strides) -{ - assert(sizes.size() == strides.size() && "mismatched ranks"); - /// off by 1 indexing to avoid out of bounds - for (auto idx = dim, e = dim + extent; idx + 1 < e; ++idx) - { - /// Only bands of static shapes are reshapable. This is due to the fact that - /// there is no relation between dynamic sizes and dynamic strides: we do not - /// have enough information to know whether a "-1" size corresponds to the - /// proper symbol in the AffineExpr of a stride. - if (ShapedType::isDynamic(sizes[dim + 1])) - return false; - /// simplify on the fly and catch more reshapable cases. - if (strides[idx] != strides[idx + 1] * sizes[idx + 1]) - return false; - } - return true; -} - -static IndexVector getIndexRange(unsigned lo, unsigned hi, unsigned step = 1) -{ - IndexVector result; - for (unsigned i = lo; i < hi; i += step) - { - result.push_back(i); - } - return result; -} - -//===----------------------------------------------------------------------===// -/// TAEarlyLoweringTTGTPass -//===----------------------------------------------------------------------===// - -/// This is a partial lowering to linear algebra of the tensor algebra -/// operations that are computationally intensive (like matmul for example...) -/// while keeping the rest of the code in the TA dialect. -namespace -{ - - struct TensorContractionOpLoweringTTGT : public ConversionPattern - { - TensorContractionOpLoweringTTGT(MLIRContext *ctx, bool isSelectBestPerm, int whatPerm, bool printFlops) - : ConversionPattern(tensorAlgebra::TensorMultOp::getOperationName(), 1, ctx), - isSelectBestPerm(isSelectBestPerm), whatPerm(whatPerm), printFlops{printFlops} {} - - /** - * @brief Latest implementation with following optimizations: - * - if no transpose is required there won't be any copy operations - * - if any operand is 2 dimensional no reshape - * - does not copy C - */ - LogicalResult - matchAndRewrite(Operation *op, ArrayRef operands, - ConversionPatternRewriter &rewriter) const final - { - comet_pdump(op); - assert(isa(op)); - - auto ctx = rewriter.getContext(); - auto loc = op->getLoc(); - auto multop = cast(op); - auto alphaAttr = multop.getOperation()->getAttr("__alpha__"); - auto betaAttr = multop.getOperation()->getAttr("__beta__"); - - Operation *startTime = nullptr; - std::string getTimeStr = "getTime"; - auto f64Type = rewriter.getF64Type(); - if (printFlops) - { - startTime = rewriter.create( - op->getLoc(), getTimeStr, SmallVector{f64Type}); - } - - ArrayAttr indexMaps = multop.getIndexingMaps(); - std::vector> allPerms; - - /// Find summation indices - for (const auto &map : indexMaps) - { - auto affineMap = cast(map).getValue(); - std::vector perm; - for (size_t i = 0; i < affineMap.getNumResults(); i++) - { - auto expr = affineMap.getResult(i); - perm.push_back(llvm::cast(expr).getPosition()); - } - - allPerms.push_back(perm); - } - - comet_pdump(op); - comet_debug() << "\n"; - // auto rhs1Tensor = cast(operands[0].getDefiningOp()); - // auto rhs2Tensor = cast(operands[1].getDefiningOp()); - comet_debug() << "\n"; - Value lhsDef; - tensorAlgebra::TensorSetOp setnewop; - for (auto u : multop.getOperation()->getResult(0).getUsers()) - { - comet_pdump(u); - if (isa(u)) - { - setnewop = cast(u); - Value dstTensor = u->getOperand(1); - - lhsDef = dstTensor; - comet_vdump(lhsDef); - } - } - // auto lhsTensor = cast(lhsDef.getDefiningOp()); - - comet_vdump(setnewop); - comet_debug() << "\n"; - - Value rhs1Memref, rhs2Memref, lhsMemref; - if(auto tensor_rhs1 = dyn_cast(operands[0].getType())) - { - rhs1Memref = rewriter.createOrFold(loc, MemRefType::get(tensor_rhs1.getShape(), tensor_rhs1.getElementType()), operands[0]); - } - else if(mlir::isa(operands[0].getType())) - { - rhs1Memref = operands[0]; - } - else - { - assert(false && "Unexpected type"); - } - - if(auto tensor_rhs2 = dyn_cast(operands[1].getType())) - { - rhs2Memref = rewriter.createOrFold(loc, MemRefType::get(tensor_rhs2.getShape(), tensor_rhs2.getElementType()), operands[1]); - } - else if(mlir::isa(operands[1].getType())) - { - rhs2Memref = operands[1]; - } - else - { - assert(false && "Unexpected type"); - } - - if(auto tensor_lhs = dyn_cast(lhsDef.getType())) - { - lhsMemref = rewriter.createOrFold(loc, MemRefType::get(tensor_lhs.getShape(), tensor_lhs.getElementType()), lhsDef); - } - else if(mlir::isa(lhsDef.getType())) - { - lhsMemref = lhsDef; - } - else - { - assert(false && "Unexpected type"); - } - - auto rhs1MemrefType = cast(rhs1Memref.getType()); - auto rhs2MemrefType = cast(rhs2Memref.getType()); - auto lhsMemrefType = cast(lhsMemref.getType()); - - std::vector allShapes{rhs1MemrefType.getShape(), - rhs2MemrefType.getShape(), - lhsMemrefType.getShape()}; - - ContractionPlan plan{allPerms[0], allShapes[0], allPerms[1], - allShapes[1], allPerms[2], allShapes[2]}; - - /// computeBestPermutations identifies the optimal index permutation for TTGT - /// it should enable and disable to heuristic - IndexVector rhs1OutPerm, rhs2OutPerm, lhsOutPerm; - std::tie(rhs1OutPerm, rhs2OutPerm, lhsOutPerm) = plan.computePermutations(isSelectBestPerm, whatPerm); - - comet_debug() << "Best permutation : " << plan.bestPermStr_ << "\n"; - - std::set rhsIndices(allPerms[0].begin(), allPerms[0].end()); - rhsIndices.insert(allPerms[1].begin(), allPerms[1].end()); - std::set lhsIndices(allPerms[2].begin(), allPerms[2].end()); - - std::vector sumIndices; - - std::set_difference(rhsIndices.begin(), rhsIndices.end(), - lhsIndices.begin(), lhsIndices.end(), - std::inserter(sumIndices, sumIndices.begin())); - - std::vector rhs1InPerm = getIdentityPermutation(allPerms[0].size()); - std::vector rhs2InPerm = getIdentityPermutation(allPerms[1].size()); - std::vector lhsInPerm = getIdentityPermutation(allPerms[2].size()); - - AffineMapAttr rhs1OutMapAttr = AffineMapAttr::get(AffineMap::getPermutationMap(rhs1OutPerm, ctx)); - [[maybe_unused]] AffineMap rhs1InMap = AffineMap::getPermutationMap(rhs1InPerm, ctx); /// [PT] Remove? - [[maybe_unused]] AffineMap rhs1OutMap = AffineMap::getPermutationMap(rhs1OutPerm, ctx); /// [PT] Remove? - - AffineMapAttr rhs2OutMapAttr = - AffineMapAttr::get(AffineMap::getPermutationMap(rhs2OutPerm, ctx)); - [[maybe_unused]] AffineMap rhs2InMap = AffineMap::getPermutationMap(rhs2InPerm, ctx); /// [PT] Remove? - [[maybe_unused]] AffineMap rhs2OutMap = AffineMap::getPermutationMap(rhs2OutPerm, ctx); /// [PT] Remove? - - AffineMapAttr lhsOutMapAttr = - AffineMapAttr::get(AffineMap::getPermutationMap(lhsOutPerm, ctx)); - [[maybe_unused]] AffineMap lhsInMap = AffineMap::getPermutationMap(lhsInPerm, ctx); /// [PT] Remove? - [[maybe_unused]] AffineMap lhsOutMap = AffineMap::getPermutationMap(lhsOutPerm, ctx); /// [PT] Remove? - - Value rhs1Alloc = rhs1Memref; - Value rhs2Alloc = rhs2Memref; - Value lhsAlloc = lhsMemref; - - std::vector rhs1OutPerm_int64(rhs1OutPerm.begin(), rhs1OutPerm.end()); - std::vector rhs2OutPerm_int64(rhs2OutPerm.begin(), rhs2OutPerm.end()); - std::vector lhsOutPerm_int64(lhsOutPerm.begin(), lhsOutPerm.end()); - - /// Do transpose if needed - if (!rhs1OutMapAttr.getValue().isIdentity()) - { - auto shape = rhs1MemrefType.getShape(); - std::vector operands; - std::vector rhs1Dims; - for (auto idx : rhs1OutPerm) - { - rhs1Dims.push_back(shape[idx]); - if (rhs1MemrefType.isDynamicDim(idx)) - { - operands.push_back(rewriter.create(loc, rhs1Memref, idx)); - } - } - - rhs1Alloc = insertAllocAndDeallocDynamic( - MemRefType::get(rhs1Dims, rhs1MemrefType.getElementType()), operands, loc, - rewriter); - -#ifdef DEBUG_MODE_TTGT - auto rhs1LinalgCopy = rewriter.create(loc, rhs1Memref, rhs1Alloc, llvm::ArrayRef(rhs1OutPerm_int64)); - comet_debug() << "\n"; - comet_vdump(rhs1LinalgCopy); -#else - rewriter.create(loc, rhs1Memref, rhs1Alloc, llvm::ArrayRef(rhs1OutPerm_int64)); -#endif - } - - if (!rhs2OutMapAttr.getValue().isIdentity()) - { - std::vector operands; - std::vector rhs2Dims; - auto shape = rhs2MemrefType.getShape(); - for (auto idx : rhs2OutPerm) - { - rhs2Dims.push_back(shape[idx]); - if (rhs2MemrefType.isDynamicDim(idx)) - { - operands.push_back(rewriter.create(loc, rhs2Memref, idx)); - } - } - - rhs2Alloc = insertAllocAndDeallocDynamic( - MemRefType::get(rhs2Dims, rhs2MemrefType.getElementType()), operands, loc, - rewriter); -#ifdef DEBUG_MODE_TTGT - auto rhs2LinalgCopy = rewriter.create(loc, rhs2Memref, rhs2Alloc, llvm::ArrayRef(rhs2OutPerm_int64)); - comet_debug() << " rhs2LinalgCopy op: " << __LINE__ << "\n"; - comet_vdump(rhs2LinalgCopy); -#else - rewriter.create(loc, rhs2Memref, rhs2Alloc, llvm::ArrayRef(rhs2OutPerm_int64)); - -#endif - } - - bool useLHSTranspose = false; - if (!lhsOutMapAttr.getValue().isIdentity()) - { - std::vector operands; - std::vector lhsDims; - auto shape = lhsMemrefType.getShape(); - for (auto idx : lhsOutPerm) - { - lhsDims.push_back(shape[idx]); - if (lhsMemrefType.isDynamicDim(idx)) - { - operands.push_back(rewriter.create(loc, lhsMemref, idx)); - } - } - - lhsAlloc = insertAllocAndDeallocDynamic( - MemRefType::get(lhsDims, lhsMemrefType.getElementType()), operands, loc, - rewriter); - useLHSTranspose = true; - double beta_val = cast(betaAttr).getValueAsDouble(); - - if (beta_val == 0) - { - Value constantOp = rewriter.create(loc, rewriter.getF64FloatAttr(0.0)); - rewriter.create(loc, constantOp, lhsAlloc); - } - else - { - rewriter.create(loc, lhsMemref, lhsAlloc, llvm::ArrayRef(lhsOutPerm_int64)); - } - } - - Value rhs1Reshape = rhs1Alloc; - Value rhs2Reshape = rhs2Alloc; - Value lhsReshape = lhsAlloc; - - unsigned mIdxSize = plan.m_indices_.size(); - unsigned nIdxSize = plan.n_indices_.size(); - unsigned kIdxSize = plan.k_indices_.size(); - - bool isRHS1SumPermutation = arePermutations(allPerms[0], sumIndices); - bool isRHS2SumPermutation = arePermutations(allPerms[1], sumIndices); - - comet_debug() << __LINE__ << "mIdxSize, nIdxSize, kIdxSize: " - << mIdxSize << ", " - << nIdxSize << ", " - << kIdxSize << " isRHS1SumPermutation, isRHS2SumPermutation: " - << isRHS1SumPermutation << ", " - << isRHS2SumPermutation << "\n"; - - /// Do reshape if needed - if (isRHS1SumPermutation) - { - auto resultShape = rhs1MemrefType.getShape(); - - auto rhs1AffineMap = AffineMap::getPermutationMap( - getIdentityPermutation(resultShape.size()), ctx); - - SmallVector rhs1IndexingMap{rhs1AffineMap}; - - SmallVector reassociationIndices = - getReassociationIndices(rhs1IndexingMap); - - comet_debug() << "\n"; - rhs1Reshape = rewriter.create( - loc, rhs1Alloc, reassociationIndices); - comet_vdump(rhs1Reshape); - } - else if (rhs1MemrefType.getShape().size() != 2) - { - auto resultShape = rhs1MemrefType.getShape(); - /// Construct combined shape of 2D memrefc - std::vector rhs1_0, rhs1_1; - - if (plan.swapAB_) - { - rhs1_0 = getIndexRange(0, kIdxSize); - rhs1_1 = getIndexRange(kIdxSize, kIdxSize + mIdxSize); - } - else - { - rhs1_0 = getIndexRange(0, mIdxSize); - rhs1_1 = getIndexRange(mIdxSize, mIdxSize + kIdxSize); - } - - auto rhs1AffineMap = AffineMap::getPermutationMap( - getIdentityPermutation(resultShape.size()), ctx); - auto rhs1Subset0 = rhs1AffineMap.getSubMap(rhs1_0); - auto rhs1Subset1 = rhs1AffineMap.getSubMap(rhs1_1); - - SmallVector rhs1IndexingMap; - - rhs1IndexingMap.push_back(rhs1Subset0); - rhs1IndexingMap.push_back(rhs1Subset1); - - SmallVector reassociationIndices = - getReassociationIndices(rhs1IndexingMap); - - comet_debug() << "\n"; - comet_debug() << " rhs1Alloc: \n"; - comet_vdump(rhs1Alloc); - comet_vdump(rhs1MemrefType); - - rhs1Reshape = rewriter.create( - loc, rhs1Alloc, reassociationIndices); - comet_debug() << " Before rhs1Reshape: \n"; - comet_vdump(rhs1Reshape); - comet_debug() << " After rhs1Reshape: \n"; - } - - if (isRHS2SumPermutation && rhs2MemrefType.getShape().size() != 1) - { - auto resultShape = rhs2MemrefType.getShape(); - - auto rhs2AffineMap = AffineMap::getPermutationMap( - getIdentityPermutation(resultShape.size()), ctx); - - SmallVector rhs2IndexingMap{rhs2AffineMap}; - - SmallVector reassociationIndices = - getReassociationIndices(rhs2IndexingMap); - - rhs2Reshape = rewriter.create( - loc, rhs2Alloc, reassociationIndices); - - comet_debug() << "\n"; - comet_vdump(rhs2Reshape); - } - else if (rhs2MemrefType.getShape().size() != 2 && rhs2MemrefType.getShape().size() != 1) - { - auto resultShape = rhs2MemrefType.getShape(); - - /// Construct combined shape of 2D memref - std::vector rhs2_0, rhs2_1; - - if (plan.swapAB_) - { - rhs2_0 = getIndexRange(0, nIdxSize); - rhs2_1 = getIndexRange(nIdxSize, nIdxSize + kIdxSize); - } - else - { - rhs2_0 = getIndexRange(0, kIdxSize); - rhs2_1 = getIndexRange(kIdxSize, kIdxSize + nIdxSize); - } - - auto rhs2AffineMap = AffineMap::getPermutationMap( - getIdentityPermutation(resultShape.size()), ctx); - auto rhs2Subset0 = rhs2AffineMap.getSubMap(rhs2_0); - auto rhs2Subset1 = rhs2AffineMap.getSubMap(rhs2_1); - - SmallVector rhs2IndexingMap; - - rhs2IndexingMap.push_back(rhs2Subset0); - rhs2IndexingMap.push_back(rhs2Subset1); - - SmallVector reassociationIndices = - getReassociationIndices(rhs2IndexingMap); - - rhs2Reshape = rewriter.create( - loc, rhs2Alloc, reassociationIndices); - - comet_debug() << "\n"; - comet_vdump(rhs2Reshape); - } - - bool expandLHS = false; - /// Keep the reassociation indices that will be used for collapsing the LHS tensor - /// The exact same indices can be used to re-expand it back to its original rank (after the potential transpose operation) - SmallVector lhsReassociationIndices; - - comet_debug() << "\n"; - if (isRHS1SumPermutation || (isRHS2SumPermutation && rhs2MemrefType.getShape().size() != 1)) - { - comet_debug() << "\n"; - auto resultShape = lhsMemrefType.getShape(); - - auto lhsAffineMap = AffineMap::getPermutationMap( - getIdentityPermutation(resultShape.size()), ctx); - - SmallVector lhsIndexingMap{lhsAffineMap}; - - SmallVector reassociationIndices = - getReassociationIndices(lhsIndexingMap); - - /// TODO(gkestor): should it be expandop? - lhsReshape = rewriter.create( - loc, lhsAlloc, reassociationIndices); - - comet_debug() << "\n"; - comet_vdump(lhsReshape); - expandLHS = true; - lhsReassociationIndices = reassociationIndices; - } - else if (lhsMemrefType.getShape().size() != 2 && lhsMemrefType.getShape().size() != 1) - { - comet_debug() << "\n"; - auto resultShape = lhsMemrefType.getShape(); - /// Construct combined shape of 2D memref - std::vector lhs_0, lhs_1; - if (plan.swapAB_) - { - lhs_0 = getIndexRange(0, nIdxSize); - lhs_1 = getIndexRange(nIdxSize, nIdxSize + mIdxSize); - } - else - { - lhs_0 = getIndexRange(0, mIdxSize); - lhs_1 = getIndexRange(mIdxSize, mIdxSize + nIdxSize); - } - - auto lhsAffineMap = AffineMap::getPermutationMap( - getIdentityPermutation(resultShape.size()), ctx); - auto lhsSubset0 = lhsAffineMap.getSubMap(lhs_0); - auto lhsSubset1 = lhsAffineMap.getSubMap(lhs_1); - - SmallVector lhsIndexingMap; - - lhsIndexingMap.push_back(lhsSubset0); - lhsIndexingMap.push_back(lhsSubset1); - - SmallVector reassociationIndices = - getReassociationIndices(lhsIndexingMap); - - lhsReshape = rewriter.create( - loc, lhsAlloc, reassociationIndices); - comet_debug() << "\n"; - comet_vdump(lhsReshape); - - expandLHS = true; - lhsReassociationIndices = reassociationIndices; - } - - comet_debug() << "\n"; - /// Create linalg matmul op - linalg::MatmulOp matmulop; - linalg::MatvecOp matvecop; - - if (isRHS1SumPermutation) - { - comet_debug() << "\n"; - matvecop = rewriter.create( - loc, ValueRange{rhs2Reshape, rhs1Reshape}, - ValueRange{lhsReshape}); - comet_debug() << "\n"; - comet_vdump(matvecop); - - matvecop.getOperation()->setAttr("__alpha__", alphaAttr); - matvecop.getOperation()->setAttr("__beta__", betaAttr); - - /// TODO(gkestor): Add attribute to the linalg.matvec operations - /// matvecop.setAttr(kLinalgTransformMarker, rewriter.getStringAttr(kLinalgTransformMarker)); - } - else if (isRHS2SumPermutation) - { - comet_debug() << "\n"; - matvecop = rewriter.create( - loc, ValueRange{rhs1Reshape, rhs2Reshape}, - ValueRange{lhsReshape}); - comet_debug() << "\n"; - comet_vdump(rhs1Reshape); - comet_vdump(rhs2Reshape); - comet_vdump(lhsReshape); - comet_vdump(matvecop); - - matvecop.getOperation()->setAttr("__alpha__", alphaAttr); - matvecop.getOperation()->setAttr("__beta__", betaAttr); - - /// TODO(gkestor): Add attribute to the linalg.matvec operations - /// matvecop.setAttr(kLinalgTransformMarker, rewriter.getStringAttr(kLinalgTransformMarker)); - } - else - { - comet_debug() << "\n"; - if (plan.swapAB_) - { - /// TODO(gkestor) - there is error with the building process - matmulop = rewriter.create( - loc, ValueRange{rhs2Reshape, rhs1Reshape}, - ValueRange{lhsReshape}); - comet_debug() << "\n"; - comet_vdump(matmulop); - } - else - { - matmulop = rewriter.create( - loc, ValueRange{rhs1Reshape, rhs2Reshape}, - ValueRange{lhsReshape}); - comet_debug() << "\n"; - comet_vdump(rhs1Reshape); - comet_vdump(rhs2Reshape); - comet_vdump(lhsReshape); - comet_vdump(matmulop); - } - comet_debug() << "\n"; - /// Add attribute to the linalg.matmul operations - matmulop.getOperation()->setAttr(kLinalgTransformMarker, - rewriter.getStringAttr(kLinalgTransformMarker)); - matmulop.getOperation()->setAttr("__alpha__", alphaAttr); - matmulop.getOperation()->setAttr("__beta__", betaAttr); - } - - Value lhsExpand = lhsReshape; - if (expandLHS) /// LHS tensor was collapsed and now needs to be re-expanded using the same reassociation indices - { - auto expandedTensorType = MemRefType::get(cast(lhsAlloc.getType()).getShape(), cast(lhsAlloc.getType()).getElementType()); - - comet_debug() << "\nExpanded:\n"; - lhsExpand = rewriter.create( - loc, expandedTensorType, lhsReshape, lhsReassociationIndices); -#ifdef DEBUG_MODE_TTGT - comet_debug() << "\n"; - comet_vdump(lhsExpand); -#endif - } - - /// Copy back the result if needed - if (lhsAlloc != lhsMemref && useLHSTranspose) - { - std::vector revLhsOutPerm(lhsOutPerm_int64.size()); - for (size_t i = 0; i < revLhsOutPerm.size(); i++) - revLhsOutPerm[lhsOutPerm_int64[i]] = i; - -#ifdef DEBUG_MODE_TTGT - auto lhsFinalCopy = - rewriter.create(loc, lhsExpand, lhsMemref, llvm::ArrayRef(revLhsOutPerm)); - comet_debug() << "\n"; - comet_vdump(lhsFinalCopy); -#else - rewriter.create(loc, lhsExpand, lhsMemref, llvm::ArrayRef(revLhsOutPerm)); - -#endif - } - - if (printFlops) - { - auto endTime = rewriter.create( - loc, getTimeStr, SmallVector{f64Type}); - - auto start = startTime->getResult(0); - auto end = endTime.getResult(0); - - Value totalTimeValue = - rewriter.create(loc, f64Type, end, start); - - double opNums = 2.0 * plan.m_size_ * plan.n_size_ * plan.k_size_; - - Value numFlopsOp = - rewriter.create(loc, FloatAttr::get(f64Type, opNums)); - - Value flopsOp = - rewriter.create(loc, f64Type, numFlopsOp, totalTimeValue); - - /// call @print_flops(%flops) : (f64) -> () - std::string printFlopsStr = "print_flops"; - /// auto printFlopsCall = - rewriter.create( - loc, printFlopsStr, SmallVector{}, ValueRange{flopsOp}); - } - - rewriter.eraseOp(setnewop); - rewriter.eraseOp(op); - return success(); - } - - private: - bool isSelectBestPerm; - int whatPerm; - bool printFlops; - }; /// namespace - - struct TALoweringTTGTPass - : public PassWrapper> - { - MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TALoweringTTGTPass) - TALoweringTTGTPass(bool isSelectBestPerm, int whatPerm, bool printFlops) : isSelectBestPerm(isSelectBestPerm), whatPerm(whatPerm), printFlops{printFlops} {}; - void runOnOperation() override; - - private: - bool isSelectBestPerm; - int whatPerm; - bool printFlops; - }; - -} /// end anonymous namespace. - -void TALoweringTTGTPass::runOnOperation() -{ - func::FuncOp function = getOperation(); - auto module = function.getOperation()->getParentOfType(); - auto *ctx = &getContext(); - - auto getTimeFunc = FunctionType::get(ctx, {}, {FloatType::getF64(ctx)}); - auto printFlopFunc = FunctionType::get(ctx, {FloatType::getF64(ctx)}, {}); - - /// func @getTime() -> f64 - if (this->printFlops && !hasFuncDeclaration(module, "getTime")) - { - mlir::func::FuncOp func1 = mlir::func::FuncOp::create(function.getLoc(), "getTime", getTimeFunc, - ArrayRef{}); - func1.setPrivate(); - module.push_back(func1); - } - - /// func @print_flops(%flops) : (f64) -> () - if (this->printFlops && !hasFuncDeclaration(module, "print_flops")) - { - mlir::func::FuncOp func1 = mlir::func::FuncOp::create(function.getLoc(), "print_flops", - printFlopFunc, ArrayRef{}); - func1.setPrivate(); - module.push_back(func1); - } - - RewritePatternSet patterns(&getContext()); - patterns.insert(&getContext(), isSelectBestPerm, whatPerm, printFlops); - - ConversionTarget target(getContext()); - target.addLegalDialect(); - - if (failed(applyPartialConversion(function, target, std::move(patterns)))) - { - llvm::errs() << "Failed to applyPartialConversion in TALoweringTTGTPass\n"; - signalPassFailure(); - } -} - -/// Create a pass for lowering operations in the `LinAlg` and `Std` dialects, -/// for a subset of the TA IR (e.g. matmul). -/// ordering of permutation starts with one -std::unique_ptr mlir::comet::createLoweringTTGTPass(bool isSelectBestPerm, int whatPerm, bool printFlops) -{ - return std::make_unique(isSelectBestPerm, whatPerm, printFlops); -} diff --git a/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp b/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp index efae7814..1da5e370 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp @@ -225,7 +225,8 @@ namespace } } - auto zero = rewriter.create(loc, FloatAttr::get(shapeT.getElementType(), 0.0)); + auto zero = rewriter.create(loc, rewriter.getZeroAttr( + shapeT.getElementType())); lhsTensor = rewriter.create(loc, shapeT, zero, ValueRange(dims)); auto lhsTensorType = cast(lhsTensor.getType()); From 30f29510c9c469f510c636a5593fa73f1f3b19d1 Mon Sep 17 00:00:00 2001 From: Polykarpos Thomadakis Date: Mon, 21 Apr 2025 14:29:24 -0700 Subject: [PATCH 06/11] Added missing attributes in --opt-multiop-factorize pass --- frontends/comet_dsl/mlir/MLIRGen.cpp | 16 ++++------------ lib/Dialect/TensorAlgebra/Transforms/Passes.cpp | 2 ++ 2 files changed, 6 insertions(+), 12 deletions(-) diff --git a/frontends/comet_dsl/mlir/MLIRGen.cpp b/frontends/comet_dsl/mlir/MLIRGen.cpp index 8c267263..0119bf57 100644 --- a/frontends/comet_dsl/mlir/MLIRGen.cpp +++ b/frontends/comet_dsl/mlir/MLIRGen.cpp @@ -1247,12 +1247,12 @@ namespace return nullptr; comet_debug() << " get rhs\n"; - // auto tens_beta = tensor_op.getBeta(); + auto tens_beta = tensor_op.getBeta(); auto lhsName = cast(*tensor_op.getLHS()) .getTensorName(); symbolTable.insert(lhsName, rhs); - // ret_op.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); + lhs.getDefiningOp()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); return rhs; } @@ -1266,16 +1266,8 @@ namespace mlir::Value tensor_op; std::vector labels; - - if ((tensor_op = symbolTable.lookup(tensor_name)) != NULL) - { - if (!isa(tensor_op.getDefiningOp()) && !isa(tensor_op.getDefiningOp())) - { - emitError(loc(lbl_tensor.loc()), "Tensor declaration required '") - << tensor_name << "'"; - } - } - else + + if ((tensor_op = symbolTable.lookup(tensor_name)) == NULL) { emitError(loc(lbl_tensor.loc()), "Unknown variable LabeledTensorExprAST'"); diff --git a/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp b/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp index fa509d1f..3e437472 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp @@ -582,6 +582,8 @@ mlir::LogicalResult FindOptimalTCFactorizationPass::FindOptimalTCFactorization(t Value tcop = builder.create(loc, newType, newRhs1, newRhs2, all_labels, affineMapArrayAttr, SemiringAttr, MaskingAttr, nullptr); + tcop.getDefiningOp()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); + tcop.getDefiningOp()->setAttr("__beta__", builder.getF64FloatAttr(0.0)); comet_debug() << "New operation " << tcop << "\n"; newRhs1 = tcop; } From 652e847abb6562b88663dd0218a5e610b790792a Mon Sep 17 00:00:00 2001 From: Polykarpos Thomadakis Date: Mon, 7 Apr 2025 09:28:26 -0700 Subject: [PATCH 07/11] Introduce proper verifiers for operations with index labels as operands --- frontends/comet_dsl/mlir/MLIRGen.cpp | 50 ++- .../comet/Dialect/TensorAlgebra/IR/TAOps.td | 341 +----------------- lib/Dialect/TensorAlgebra/IR/TADialect.cpp | 102 ++++++ .../TensorAlgebra/Transforms/Passes.cpp | 12 +- lib/Dialect/Utils/Utils.cpp | 2 +- 5 files changed, 158 insertions(+), 349 deletions(-) diff --git a/frontends/comet_dsl/mlir/MLIRGen.cpp b/frontends/comet_dsl/mlir/MLIRGen.cpp index 0119bf57..4205d252 100644 --- a/frontends/comet_dsl/mlir/MLIRGen.cpp +++ b/frontends/comet_dsl/mlir/MLIRGen.cpp @@ -38,6 +38,7 @@ #include "mlir/Dialect/Func/IR/FuncOps.h" #include "mlir/IR/BuiltinAttributes.h" #include "mlir/IR/BuiltinTypeInterfaces.h" +#include "mlir/IR/ValueRange.h" #include "mlir/IR/Verifier.h" #include "mlir/IR/Attributes.h" #include "mlir/IR/Builders.h" @@ -160,8 +161,9 @@ namespace std::vector getOperationResultIndexLabels(mlir::Operation *op) { auto cast_op = cast(op); - return cast_op.getResultIndexLabels(); + return std::vector(cast_op.getResultIndexLabels().begin(), cast_op.getResultIndexLabels().end()); } + // Implementation of a simple MLIR emission from the Tensor Algebra AST. /// // This will emit operations that are specific to the Tensor Algebra language, @@ -953,18 +955,22 @@ namespace assert(tensors.size() == 2 && " less than 2 input tensors for ta.mul or ta.elews_mul\n"); - std::vector labels; + std::vector labels, newLhsLabels, newRhsLabels, newRetLabels; for (auto i : lhs_lbls) { + labels.push_back(all_lbls_value[i]); + newLhsLabels.push_back(all_lbls_value[i]); } for (auto i : rhs_lbls) { labels.push_back(all_lbls_value[i]); + newRhsLabels.push_back(all_lbls_value[i]); } for (auto i : ret_lbls) { labels.push_back(all_lbls_value[i]); + newRetLabels.push_back(all_lbls_value[i]); } mlir::StringAttr SemiringAttr; @@ -978,14 +984,14 @@ namespace SemiringAttr = builder.getStringAttr("noop_plusxy"); // this is for standard elementwise addition MaskingAttr = builder.getStringAttr("none"); // default for standard elementwise addition return builder.create(location, ret_tensor_type, tensors[0], tensors[1], - labels, affineMapArrayAttr, SemiringAttr, + newLhsLabels, newRhsLabels, newRetLabels, affineMapArrayAttr, SemiringAttr, MaskingAttr); case '-': comet_debug() << "creating TensorSubtractOp\n"; SemiringAttr = builder.getStringAttr("noop_minus"); // this is for standard elementwise subtraction MaskingAttr = builder.getStringAttr("none"); // default for standard elementwise subtraction return builder.create(location, ret_tensor_type, tensors[0], tensors[1], - labels, affineMapArrayAttr, SemiringAttr, + newLhsLabels, newRhsLabels, newRetLabels, affineMapArrayAttr, SemiringAttr, MaskingAttr); case '*': { @@ -997,7 +1003,7 @@ namespace SemiringAttr = builder.getStringAttr("plusxy_times"); // this is for standard matrix multiplication MaskingAttr = builder.getStringAttr("none"); // default for standard matrix multiplication mlir::Value tcop = builder.create(location, ret_tensor_type, tensors[0], tensors[1], - labels, affineMapArrayAttr, SemiringAttr, + newLhsLabels, newRhsLabels, newRetLabels, affineMapArrayAttr, SemiringAttr, MaskingAttr, nullptr); // TODO: masking is an optional operand tcop.getDefiningOp()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); tcop.getDefiningOp()->setAttr("__beta__", builder.getF64FloatAttr(0.0)); @@ -1014,7 +1020,7 @@ namespace comet_debug() << "\n"; auto SemiringAttr = builder.getStringAttr("noop_times"); /// this is for standard element-wise multiplication MaskingAttr = builder.getStringAttr("none"); /// default for standard element-wise multiplication - mlir::Value tcop = builder.create(location, ret_tensor_type, tensors[0], tensors[1], labels, + mlir::Value tcop = builder.create(location, ret_tensor_type, tensors[0], tensors[1], newLhsLabels, newRhsLabels, newRetLabels, affineMapArrayAttr, SemiringAttr, MaskingAttr); @@ -1476,7 +1482,18 @@ namespace std::vector all_lbls = rhs_lbls; all_lbls.insert(all_lbls.end(), lhs_lbls.begin(), lhs_lbls.end()); - std::vector all_lbls_value; + std::vector all_lbls_value, lhs_lbls_value, rhs_lbls_value; + + for (auto s : lhs_lbls) + { + lhs_lbls_value.push_back(symbolTable.lookup(s)); + } + + for (auto s : rhs_lbls) + { + rhs_lbls_value.push_back(symbolTable.lookup(s)); + } + for (auto s : all_lbls) { all_lbls_value.push_back(symbolTable.lookup(s)); @@ -1516,7 +1533,7 @@ namespace comet_debug() << " create TransposeOp\n"; mlir::Value t = builder.create(loc(transpose.loc()), lhs_tensor.getType(), - rhs_tensor, all_lbls_value, affineMapArrayAttr); + rhs_tensor, rhs_lbls_value, lhs_lbls_value, affineMapArrayAttr); symbolTable.insert(lhsLT.getTensorName(), t) comet_vdump(t); @@ -1570,7 +1587,7 @@ namespace /// Secondly, Look at the lhs /// Collect labels values - std::vector lhs_labels_val; + std::vector lhs_labels_val, rhs_labels_val; std::vector all_labels_val; for (const auto &lbl_str : rhs_lbls) { @@ -1578,6 +1595,7 @@ namespace { if (isa(var.getDefiningOp())) { + rhs_labels_val.push_back(var); all_labels_val.push_back(var); } else @@ -1665,7 +1683,7 @@ namespace comet_debug() << " create TransposeOp\n"; mlir::ShapedType shapedT = mlir::cast(rhs_tensor.getType()); mlir::Value t = builder.create(loc(transpose.loc()), mlir::RankedTensorType::get(shape, shapedT.getElementType()), - rhs_tensor, all_labels_val, affineMapArrayAttr); + rhs_tensor, rhs_labels_val, lhs_labels_val, affineMapArrayAttr); comet_vdump(t); return t; @@ -2208,7 +2226,7 @@ namespace auto rhs2_lbls = rhs2LT->getLabelNames(); auto all_lbls = rhsBinOp->getLabels(); - std::vector lhs_lbls_value; + std::vector lhs_lbls_value, rhs1_lbls_value, rhs2_lbls_value; std::vector all_lbls_value; for (auto n : lhs_lbls) { @@ -2216,10 +2234,12 @@ namespace } for (auto n : rhs1_lbls) { + rhs1_lbls_value.push_back(symbolTable.lookup(n)); all_lbls_value.push_back(symbolTable.lookup(n)); } for (auto n : rhs2_lbls) { + rhs2_lbls_value.push_back(symbolTable.lookup(n)); all_lbls_value.push_back(symbolTable.lookup(n)); } for (auto n : lhs_lbls) @@ -2343,7 +2363,7 @@ namespace auto op = builder.create(loc(tensor_op.loc()), tensors[2].getType(), tensors[0], tensors[1], - all_lbls_value, + rhs1_lbls_value, rhs2_lbls_value, lhs_lbls_value, affineMapArrayAttr, SemiringAttr, MaskingAttr); @@ -2357,7 +2377,7 @@ namespace auto op = builder.create(loc(tensor_op.loc()), tensors[2].getType(), tensors[0], tensors[1], - all_lbls_value, + rhs1_lbls_value, rhs2_lbls_value, lhs_lbls_value, affineMapArrayAttr, SemiringAttr, MaskingAttr); @@ -2369,7 +2389,7 @@ namespace auto op = builder.create(loc(tensor_op.loc()), tensors[2].getType(), tensors[0], tensors[1], - all_lbls_value, + rhs1_lbls_value, rhs2_lbls_value, lhs_lbls_value, affineMapArrayAttr, SemiringAttr, MaskingAttr, maskVal); @@ -2381,7 +2401,7 @@ namespace } else if (binop == tok_elews || binop == tok_monoid) { - auto op = builder.create(loc(tensor_op.loc()), tensors[2].getType(), tensors[0], tensors[1], all_lbls_value, affineMapArrayAttr, SemiringAttr, MaskingAttr); + auto op = builder.create(loc(tensor_op.loc()), tensors[2].getType(), tensors[0], tensors[1], rhs1_lbls_value, rhs2_lbls_value, lhs_lbls_value, affineMapArrayAttr, SemiringAttr, MaskingAttr); op.getOperation()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); op.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); diff --git a/include/comet/Dialect/TensorAlgebra/IR/TAOps.td b/include/comet/Dialect/TensorAlgebra/IR/TAOps.td index b9d1c50e..7ca72a07 100644 --- a/include/comet/Dialect/TensorAlgebra/IR/TAOps.td +++ b/include/comet/Dialect/TensorAlgebra/IR/TAOps.td @@ -351,7 +351,7 @@ def TensorMultOp : TA_Op<"mul", [Pure, AttrSizedOperandSegments]>{ let description = [{ }]; - let arguments = (ins TA_AnyTensor:$rhs1, TA_AnyTensor:$rhs2, Variadic:$index_labels, + let arguments = (ins TA_AnyTensor:$rhs1, TA_AnyTensor:$rhs2, Variadic:$rhs1_index_labels, Variadic:$rhs2_index_labels, Variadic:$result_index_labels, AffineMapArrayAttr:$indexing_maps, StrAttr:$semiring, OptionalAttr:$MaskType, /// TODO: try to use: DefaultValuedAttr Optional:$mask); @@ -359,80 +359,10 @@ def TensorMultOp : TA_Op<"mul", [Pure, AttrSizedOperandSegments]>{ /// Return value let results = (outs TA_AnyTensor); - let extraClassDeclaration = [{ - template - size_t getTensorRank(T tensor) - { - if(auto denseT = mlir::dyn_cast(tensor.getType())) - { - return denseT.getRank(); - } - else if(auto sparseT = mlir::dyn_cast(tensor.getType())) - { - return sparseT.getRank(); - } - else - { - assert(false&& "Unexpected Type"); - } - } - - std::vector getRhs1IndexLabels() - { - std::vector labels; - size_t start = 0; - - size_t end = getTensorRank(getRhs1()); - - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - - } - - std::vector getRhs2IndexLabels() - { - std::vector labels; - - - size_t start =getTensorRank(getRhs1()); - - size_t end = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); - - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - - } - - std::vector getResultIndexLabels() - { - std::vector labels; - size_t end = getIndexLabels().size(); - - size_t start = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); - - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - } - - }]; - //TODO(gkestor): add verifier - //let hasVerifier = 1; - + let hasVerifier = 1; } -def TensorElewsMultOp : TA_Op<"elews_mul", [Pure]>{ +def TensorElewsMultOp : TA_Op<"elews_mul", [Pure, AttrSizedOperandSegments]>{ let summary = ""; let description = [{ @@ -440,83 +370,14 @@ def TensorElewsMultOp : TA_Op<"elews_mul", [Pure]>{ let arguments = (ins TA_AnyTensor:$rhs1, TA_AnyTensor:$rhs2, - Variadic:$index_labels, + Variadic:$rhs1_index_labels, Variadic:$rhs2_index_labels, Variadic:$result_index_labels, AffineMapArrayAttr:$indexing_maps, StrAttr:$semiring, OptionalAttr:$MaskType); let results = (outs TA_AnyTensor); - let extraClassDeclaration = [{ - template - size_t getTensorRank(T tensor) - { - if(auto denseT = mlir::dyn_cast(tensor.getType())) - { - return denseT.getRank(); - } - else if(auto sparseT = mlir::dyn_cast(tensor.getType())) - { - return sparseT.getRank(); - } - else - { - assert(false&& "Unexpected Type"); - } - } - - std::vector getRhs1IndexLabels() - { - std::vector labels; - size_t start = 0; - size_t end = getTensorRank(getRhs1()); - - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - - } - - std::vector getRhs2IndexLabels() - { - std::vector labels; - size_t start = getTensorRank(getRhs1()); - - size_t end = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); - - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - - } - - std::vector getResultIndexLabels() - { - std::vector labels; - size_t end = getIndexLabels().size(); - - size_t start = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); - - - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - } - - }]; - - //TODO(gkestor): add verifier - //let hasVerifier = 1; - + let hasVerifier = 1; } def TensorDimOp : TA_Op<"dim", [Pure]>{ @@ -539,64 +400,14 @@ def TensorDimOp : TA_Op<"dim", [Pure]>{ //let hasVerifier = 1; } -def TransposeOp : TA_Op<"transpose",[Pure]> { +def TransposeOp : TA_Op<"transpose",[Pure, AttrSizedOperandSegments]> { let summary = "transpose operation"; - let arguments = (ins TA_AnyTensor:$rhs, Variadic:$index_labels, AffineMapArrayAttr:$indexing_maps); + let arguments = (ins TA_AnyTensor:$rhs, Variadic:$rhs_index_labels, Variadic:$result_index_labels, AffineMapArrayAttr:$indexing_maps); let results = (outs TA_AnyTensor); - - - let extraClassDeclaration = [{ - template - size_t getTensorRank(T tensor) - { - if(auto denseT = mlir::dyn_cast(tensor.getType())) - { - return denseT.getRank(); - } - else if(auto sparseT = mlir::dyn_cast(tensor.getType())) - { - return sparseT.getRank(); - } - else - { - assert(false&& "Unexpected Type"); - } - } - - std::vector getRhsIndexLabels() - { - std::vector labels; - size_t start = 0; - size_t end = static_cast(getTensorRank(getRhs())); - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - - } - - std::vector getResultIndexLabels() - { - std::vector labels; - size_t start = static_cast(getTensorRank(getRhs())); - size_t end = getIndexLabels().size(); - - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - } - - }]; - //TODO(gkestor): add verifier - //let hasVerifier = 1; + let hasVerifier = 1; } def ReduceOp : TA_Op<"reduce",[Pure]> { @@ -614,7 +425,7 @@ def ReduceOp : TA_Op<"reduce",[Pure]> { } def TensorAddOp : TA_Op<"add", - [Pure]> { + [Pure, AttrSizedOperandSegments]> { let summary = "element-wise addition operation"; let description = [{ The "add" operation performs element-wise addition between two tensors. @@ -623,79 +434,17 @@ def TensorAddOp : TA_Op<"add", let arguments = (ins TA_AnyTensor:$rhs1, TA_AnyTensor:$rhs2, - Variadic:$index_labels, + Variadic:$rhs1_index_labels, Variadic:$rhs2_index_labels, Variadic:$result_index_labels, AffineMapArrayAttr:$indexing_maps, StrAttr:$semiring, OptionalAttr:$MaskType); let results = (outs TA_AnyTensor); - - let extraClassDeclaration = [{ - template - size_t getTensorRank(T tensor) - { - if(auto denseT = mlir::dyn_cast(tensor.getType())) - { - return denseT.getRank(); - } - else if(auto sparseT = mlir::dyn_cast(tensor.getType())) - { - return sparseT.getRank(); - } - else - { - assert(false&& "Unexpected Type"); - } - } - - std::vector getRhs1IndexLabels() - { - std::vector labels; - size_t start = 0; - size_t end = getTensorRank(getRhs1()); - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - - } - - std::vector getRhs2IndexLabels() - { - std::vector labels; - size_t start = getTensorRank(getRhs1()); - size_t end = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); - - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - - } - - std::vector getResultIndexLabels() - { - std::vector labels; - size_t start = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); - size_t end = getIndexLabels().size(); - - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - } - - }]; + let hasVerifier = 1; } def TensorSubtractOp : TA_Op<"subtract", - [Pure]> { + [Pure, AttrSizedOperandSegments]> { let summary = "element-wise subtract operation"; let description = [{ The "subtract" operation performs element-wise subtract between two tensors. @@ -704,75 +453,13 @@ def TensorSubtractOp : TA_Op<"subtract", let arguments = (ins TA_AnyTensor:$rhs1, TA_AnyTensor:$rhs2, - Variadic:$index_labels, + Variadic:$rhs1_index_labels, Variadic:$rhs2_index_labels, Variadic:$result_index_labels, AffineMapArrayAttr:$indexing_maps, StrAttr:$semiring, OptionalAttr:$MaskType); let results = (outs TA_AnyTensor); - - let extraClassDeclaration = [{ - template - size_t getTensorRank(T tensor) - { - if(auto denseT = mlir::dyn_cast(tensor.getType())) - { - return denseT.getRank(); - } - else if(auto sparseT = mlir::dyn_cast(tensor.getType())) - { - return sparseT.getRank(); - } - else - { - assert(false&& "Unexpected Type"); - } - } - - std::vector getRhs1IndexLabels() - { - std::vector labels; - size_t start = 0; - size_t end = getTensorRank(getRhs1()); - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - - } - - std::vector getRhs2IndexLabels() - { - std::vector labels; - size_t start = getTensorRank(getRhs1()); - size_t end = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); - - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - - } - - std::vector getResultIndexLabels() - { - std::vector labels; - size_t start = getTensorRank(getRhs1()) + getTensorRank(getRhs2()); - size_t end = getIndexLabels().size(); - - for(size_t i = start; i < end; i++ ) - { - labels.push_back(getIndexLabels()[i]); - } - - return labels; - } - - }]; + let hasVerifier = 1; } //TODO(gkestor): support other datatypes diff --git a/lib/Dialect/TensorAlgebra/IR/TADialect.cpp b/lib/Dialect/TensorAlgebra/IR/TADialect.cpp index c87c57db..65bfd75d 100644 --- a/lib/Dialect/TensorAlgebra/IR/TADialect.cpp +++ b/lib/Dialect/TensorAlgebra/IR/TADialect.cpp @@ -41,6 +41,7 @@ #include "mlir/Dialect/Bufferization/IR/DstBufferizableOpInterfaceImpl.h" #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h" #include "llvm/ADT/TypeSwitch.h" +#include "llvm/Support/LogicalResult.h" using namespace mlir; using namespace mlir::tensorAlgebra; @@ -273,6 +274,107 @@ void TensorDimOp::build(mlir::OpBuilder &builder, mlir::OperationState &result, } +LogicalResult TensorMultOp::verify() +{ + if(getRhs1IndexLabels().size() != static_cast(mlir::cast(getRhs1().getType()).getRank())) + { + emitError("number of index labels for RHS1 does not match its rank"); + return failure(); + } + if(getRhs2IndexLabels().size() != static_cast(mlir::cast(getRhs2().getType()).getRank())) + { + emitError("number of index labels for RHS2 does not match its rank"); + return failure(); + } + if(getResultIndexLabels().size() != static_cast(mlir::cast(getResult().getType()).getRank())) + { + emitError("number of index labels for result does not match its rank"); + return failure(); + } + + return success(); +} + +LogicalResult TensorAddOp::verify() +{ + if(getRhs1IndexLabels().size() != static_cast(mlir::cast(getRhs1().getType()).getRank())) + { + emitError("number of index labels for RHS1 does not match its rank"); + return failure(); + } + if(getRhs2IndexLabels().size() != static_cast(mlir::cast(getRhs2().getType()).getRank())) + { + emitError("number of index labels for RHS2 does not match its rank"); + return failure(); + } + if(getResultIndexLabels().size() != static_cast(mlir::cast(getResult().getType()).getRank())) + { + emitError("number of index labels for result does not match its rank"); + return failure(); + } + + return success(); +} + +LogicalResult TensorSubtractOp::verify() +{ + if(getRhs1IndexLabels().size() != static_cast(mlir::cast(getRhs1().getType()).getRank())) + { + emitError("number of index labels for RHS1 does not match its rank"); + return failure(); + } + if(getRhs2IndexLabels().size() != static_cast(mlir::cast(getRhs2().getType()).getRank())) + { + emitError("number of index labels for RHS2 does not match its rank"); + return failure(); + } + if(getResultIndexLabels().size() != static_cast(mlir::cast(getResult().getType()).getRank())) + { + emitError("number of index labels for result does not match its rank"); + return failure(); + } + + return success(); +} + +LogicalResult TensorElewsMultOp::verify() +{ + if(getRhs1IndexLabels().size() != static_cast(mlir::cast(getRhs1().getType()).getRank())) + { + emitError("number of index labels for RHS1 does not match its rank"); + return failure(); + } + if(getRhs2IndexLabels().size() != static_cast(mlir::cast(getRhs2().getType()).getRank())) + { + emitError("number of index labels for RHS2 does not match its rank"); + return failure(); + } + if(getResultIndexLabels().size() != static_cast(mlir::cast(getResult().getType()).getRank())) + { + emitError("number of index labels for result does not match its rank"); + return failure(); + } + + return success(); +} + +LogicalResult TransposeOp::verify() +{ + if(getRhsIndexLabels().size() != static_cast(mlir::cast(getRhs().getType()).getRank())) + { + emitError("number of index labels for RHS1 does not match its rank"); + return failure(); + } + + if(getResultIndexLabels().size() != static_cast(mlir::cast(getResult().getType()).getRank())) + { + emitError("number of index labels for result does not match its rank"); + return failure(); + } + + return success(); +} + /// Sort Op bool hasFuncDeclaration(ModuleOp &module, std::string funcName) { diff --git a/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp b/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp index 3e437472..62ecd22c 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp @@ -256,7 +256,7 @@ mlir::LogicalResult FindOptimalTCFactorizationPass::FindOptimalTCFactorization(t auto multop = cast(curr); stack.push(curr); MultOpsToRemove.push_back(multop.getOperation()); - std::vector labels = multop.getRhs2IndexLabels(); + auto labels = multop.getRhs2IndexLabels(); std::vector labelVec; for (size_t i = 0; i < labels.size(); i++) @@ -507,11 +507,11 @@ mlir::LogicalResult FindOptimalTCFactorizationPass::FindOptimalTCFactorization(t std::vector lhs_lbls; std::vector rhs_lbls; - std::vector all_labels; + std::vector newRhs1Labels, newRhs2Labels, newResultLabels; for (auto e : new_rhs_lbls_value) { - all_labels.push_back(e); + newRhs1Labels.push_back(e); auto index = std::find(new_all_lbls_value.begin(), new_all_lbls_value.end(), e) - new_all_lbls_value.begin(); comet_debug() << index << " " << new_all_lbls_value[index] << "\n"; @@ -520,7 +520,7 @@ mlir::LogicalResult FindOptimalTCFactorizationPass::FindOptimalTCFactorization(t for (auto e : new_lhs_lbls_value) { - all_labels.push_back(e); + newRhs2Labels.push_back(e); auto index = std::find(new_all_lbls_value.begin(), new_all_lbls_value.end(), e) - new_all_lbls_value.begin(); comet_debug() << index << " " << new_all_lbls_value[index] << "\n"; lhs_lbls.push_back(index); @@ -530,7 +530,7 @@ mlir::LogicalResult FindOptimalTCFactorizationPass::FindOptimalTCFactorization(t for (auto e : newSumLabels) { - all_labels.push_back(e); + newResultLabels.push_back(e); auto index = std::find(new_all_lbls_value.begin(), new_all_lbls_value.end(), e) - new_all_lbls_value.begin(); comet_debug() << index << " " << new_all_lbls_value[index] << "\n"; @@ -580,7 +580,7 @@ mlir::LogicalResult FindOptimalTCFactorizationPass::FindOptimalTCFactorization(t auto SemiringAttr = builder.getStringAttr("plusxy_times"); auto MaskingAttr = builder.getStringAttr("none"); Value tcop = builder.create(loc, newType, newRhs1, newRhs2, - all_labels, affineMapArrayAttr, SemiringAttr, + newRhs1Labels, newRhs2Labels, newResultLabels, affineMapArrayAttr, SemiringAttr, MaskingAttr, nullptr); tcop.getDefiningOp()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); tcop.getDefiningOp()->setAttr("__beta__", builder.getF64FloatAttr(0.0)); diff --git a/lib/Dialect/Utils/Utils.cpp b/lib/Dialect/Utils/Utils.cpp index f2a27110..990390d0 100644 --- a/lib/Dialect/Utils/Utils.cpp +++ b/lib/Dialect/Utils/Utils.cpp @@ -2237,7 +2237,7 @@ namespace mlir auto MaskingAttr = rewriter.getStringAttr("none"); auto tc = rewriter.create(loc, lhsTensor.getType(), rhs1Tensor, rhs2Tensor, - lhsLabels, affineMapArrayAttr, + rhs1Labels, rhs2Labels, lhsLabels, affineMapArrayAttr, SemiringAttr, MaskingAttr, nullptr); /// TODO: masking is an optional operand tc.getOperation()->setAttr("__alpha__", rewriter.getF64FloatAttr(alpha)); From 81668b77e0e9e8b2af6259ce88969b265d6a7adf Mon Sep 17 00:00:00 2001 From: Polykarpos Thomadakis Date: Wed, 23 Apr 2025 11:12:48 -0700 Subject: [PATCH 08/11] 1. All tensor operations (Add, Sub, Mult, EltwiseMult, Tranpose) now accept an optional operand `lhs` which can be used to specify the output storage of the operation that will occur during bufferization, similar to linalg's paradigm. For example, an expression C = Ax + Bx would previously work as: A' = Ax, B' = Bx, C = A' + B'. Thus, each sub-expression would cause new allocations (A', B' tensors). In contrast, by being able to optionally specify the output of each operation, we can store the result of the LHS of an expression into the output storage and accumulate the rest, i.e. C = Ax, B' = Bx, C = C + B'. Thus, we can remove some of the intermediate allocations. Note that, we could eventually also avoid the allocation of B' if we change operations to accumulate into C (C[i,j] += op) instead of storing (C[i,j] = op) as we do in matmul 2. Fixed fusion + parallel, but only without applying dimension reduction 3. Made dimension reduction optional --- frontends/comet_dsl/comet.cpp | 2 +- frontends/comet_dsl/mlir/MLIRGen.cpp | 49 +++-- .../comet/Dialect/TensorAlgebra/IR/TAOps.td | 6 +- include/comet/Dialect/Utils/Utils.h | 14 +- .../IndexTreeToSCF/IndexTreeToSCF.cpp | 2 +- .../TensorAlgebraToIndexTree.cpp | 23 ++- .../TensorAlgebraToSCF/TensorAlgebraToSCF.cpp | 8 +- .../IndexTree/Transforms/KernelFusion.cpp | 8 +- .../TensorAlgebra/Transforms/Passes.cpp | 3 + .../TensorAlgebra/Transforms/TCtoTTGTDyn.cpp | 7 +- lib/Dialect/Utils/Utils.cpp | 172 +++++++++--------- 11 files changed, 164 insertions(+), 130 deletions(-) diff --git a/frontends/comet_dsl/comet.cpp b/frontends/comet_dsl/comet.cpp index 752b3750..e45694b9 100644 --- a/frontends/comet_dsl/comet.cpp +++ b/frontends/comet_dsl/comet.cpp @@ -467,7 +467,7 @@ int loadAndProcessMLIR(mlir::MLIRContext &context, /// Concretize the domains of all the index variables optPM.addPass(mlir::comet::createIndexTreeDomainConcretizationPass()); - if (OptKernelFusion || OptDimensionReduction) + if (OptDimensionReduction) { /// Reduce intermediate tensors' dimension after kernel fusion optPM.addPass(mlir::comet::createIndexTreeDimensionReductionPass()); diff --git a/frontends/comet_dsl/mlir/MLIRGen.cpp b/frontends/comet_dsl/mlir/MLIRGen.cpp index 4205d252..e9391f2a 100644 --- a/frontends/comet_dsl/mlir/MLIRGen.cpp +++ b/frontends/comet_dsl/mlir/MLIRGen.cpp @@ -52,11 +52,13 @@ #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" #include "llvm/Support/raw_ostream.h" +#include #include #include #include #include #include /// for random num generation +#include #include /// for seed of random num generation #include @@ -82,7 +84,7 @@ int32_t defaultSpTensorIndiceBitWidth = 64; // TODO: We should be able to pass t using StringSet = std::set; // *********** For debug purpose *********// -//#define COMET_DEBUG_MODE +// #define COMET_DEBUG_MODE #include "comet/Utils/debug.h" #undef COMET_DEBUG_MODE // *********** For debug purpose *********// @@ -349,7 +351,7 @@ namespace // Emit a binary operation mlir::Value mlirGen(BinaryExprAST &binop, - const std::set &out_lbls = {}, std::string out_format = "") + const std::set &out_lbls = {}, std::string out_format = "", mlir::Value out_val = nullptr) { comet_debug() << " mlirGen for BinaryExprAST \n"; /// First emit the operations for each side of the operation before emitting @@ -442,10 +444,10 @@ namespace << __LINE__ << " rhsAST is not Expr_BinOp/Expr_LabeledTensor/CallExprAST/TransposeExprAST \n"; } - mlir::Value lhs = mlirGen(*binop.getLHS(), rhs_lbls); + mlir::Value lhs = mlirGen(*binop.getLHS(), rhs_lbls, "", binop.getOp() != '*' ? out_val : nullptr); if (!lhs) return nullptr; - mlir::Value rhs = mlirGen(*binop.getRHS(), lhs_lbls); + mlir::Value rhs = mlirGen(*binop.getRHS(), lhs_lbls, "", nullptr); if (!rhs) return nullptr; auto location = loc(binop.loc()); @@ -535,7 +537,6 @@ namespace comet_vdump(rhs); comet_vdump(lhs); - comet_vdump(theOutput); auto scalarOp = builder.create(location, returnDataType, rhs, lhs, opAttr); comet_vdump(scalarOp); // symbolTable.insert(out_format, scalarOp); @@ -977,6 +978,10 @@ namespace mlir::StringAttr MaskingAttr; /// Derive the operation name from the binary operator. At the moment we /// only support '+', '-','*'. + if(out_val && out_val.getType() != ret_tensor_type) + { + out_val = nullptr; + } switch (binop.getOp()) { case '+': @@ -985,6 +990,7 @@ namespace MaskingAttr = builder.getStringAttr("none"); // default for standard elementwise addition return builder.create(location, ret_tensor_type, tensors[0], tensors[1], newLhsLabels, newRhsLabels, newRetLabels, affineMapArrayAttr, SemiringAttr, + out_val, MaskingAttr); case '-': comet_debug() << "creating TensorSubtractOp\n"; @@ -992,6 +998,7 @@ namespace MaskingAttr = builder.getStringAttr("none"); // default for standard elementwise subtraction return builder.create(location, ret_tensor_type, tensors[0], tensors[1], newLhsLabels, newRhsLabels, newRetLabels, affineMapArrayAttr, SemiringAttr, + out_val, MaskingAttr); case '*': { @@ -1004,6 +1011,7 @@ namespace MaskingAttr = builder.getStringAttr("none"); // default for standard matrix multiplication mlir::Value tcop = builder.create(location, ret_tensor_type, tensors[0], tensors[1], newLhsLabels, newRhsLabels, newRetLabels, affineMapArrayAttr, SemiringAttr, + out_val, MaskingAttr, nullptr); // TODO: masking is an optional operand tcop.getDefiningOp()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); tcop.getDefiningOp()->setAttr("__beta__", builder.getF64FloatAttr(0.0)); @@ -1022,6 +1030,7 @@ namespace MaskingAttr = builder.getStringAttr("none"); /// default for standard element-wise multiplication mlir::Value tcop = builder.create(location, ret_tensor_type, tensors[0], tensors[1], newLhsLabels, newRhsLabels, newRetLabels, affineMapArrayAttr, SemiringAttr, + out_val, MaskingAttr); comet_vdump(tcop); @@ -1248,17 +1257,17 @@ namespace } comet_debug() << "\n"; - auto rhs = mlirGen(*tensor_op.getRHS(), out_labels, out_format); + auto rhs = mlirGen(*tensor_op.getRHS(), out_labels, out_format, lhs); if (!rhs) return nullptr; comet_debug() << " get rhs\n"; - auto tens_beta = tensor_op.getBeta(); + // auto tens_beta = tensor_op.getBeta(); auto lhsName = cast(*tensor_op.getLHS()) .getTensorName(); symbolTable.insert(lhsName, rhs); - lhs.getDefiningOp()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); + // lhs.getDefiningOp()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); return rhs; } @@ -1284,13 +1293,13 @@ namespace /// Dispatch codegen for the right expression subclass using RTTI. mlir::Value mlirGen(ExprAST &expr, - const std::set out_lbls = {}, std::string out_format = "") + const std::set out_lbls = {}, std::string out_format = "", mlir::Value val = nullptr) { comet_debug() << " mlirGen ExprAST " << expr.getKind() << " \n"; switch (expr.getKind()) { case tensorAlgebra::ExprAST::Expr_BinOp: - return mlirGen(cast(expr), out_lbls, out_format); + return mlirGen(cast(expr), out_lbls, out_format, val); case tensorAlgebra::ExprAST::Expr_Var: return mlirGen(cast(expr)); case tensorAlgebra::ExprAST::Expr_Literal: @@ -1298,7 +1307,7 @@ namespace case tensorAlgebra::ExprAST::Expr_Call: return mlirGen(cast(expr)); case tensorAlgebra::ExprAST::Expr_Transpose: - return mlirGen(cast(expr)); + return mlirGen(cast(expr), out_lbls, out_format, val); case tensorAlgebra::ExprAST::Expr_Num: return mlirGen(cast(expr)); case tensorAlgebra::ExprAST::Expr_LabeledTensor: @@ -1533,15 +1542,15 @@ namespace comet_debug() << " create TransposeOp\n"; mlir::Value t = builder.create(loc(transpose.loc()), lhs_tensor.getType(), - rhs_tensor, rhs_lbls_value, lhs_lbls_value, affineMapArrayAttr); - symbolTable.insert(lhsLT.getTensorName(), t) + rhs_tensor, rhs_lbls_value, lhs_lbls_value, affineMapArrayAttr, lhs_tensor); + symbolTable.insert(lhsLT.getTensorName(), t); comet_vdump(t); return t; } /// Handle tranpose(A[i,j], {j, i}) in DSL, when no lhs_LabeledTensor has been created. - mlir::Value mlirGen(TransposeExprAST &transpose) + mlir::Value mlirGen(TransposeExprAST &transpose, const std::set out_lbls = {}, std::string out_format = "", mlir::Value val= nullptr) { comet_debug() << "TransposeExprAST with no lhs labeled tensor \n"; @@ -1682,8 +1691,13 @@ namespace comet_debug() << " create TransposeOp\n"; mlir::ShapedType shapedT = mlir::cast(rhs_tensor.getType()); + auto out_type = mlir::RankedTensorType::get(shape, shapedT.getElementType()); + if(val && val.getType() != out_type) + { + val = nullptr; + } mlir::Value t = builder.create(loc(transpose.loc()), mlir::RankedTensorType::get(shape, shapedT.getElementType()), - rhs_tensor, rhs_labels_val, lhs_labels_val, affineMapArrayAttr); + rhs_tensor, rhs_labels_val, lhs_labels_val, affineMapArrayAttr, val); comet_vdump(t); return t; @@ -2366,6 +2380,7 @@ namespace rhs1_lbls_value, rhs2_lbls_value, lhs_lbls_value, affineMapArrayAttr, SemiringAttr, + tensors[2], MaskingAttr); comet_vdump(op); @@ -2380,6 +2395,7 @@ namespace rhs1_lbls_value, rhs2_lbls_value, lhs_lbls_value, affineMapArrayAttr, SemiringAttr, + tensors[2], MaskingAttr); comet_vdump(op); symbolTable.insert(exprs[2]->getTensorName(), op); @@ -2392,6 +2408,7 @@ namespace rhs1_lbls_value, rhs2_lbls_value, lhs_lbls_value, affineMapArrayAttr, SemiringAttr, + tensors[2], MaskingAttr, maskVal); op.getOperation()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); op.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); @@ -2401,7 +2418,7 @@ namespace } else if (binop == tok_elews || binop == tok_monoid) { - auto op = builder.create(loc(tensor_op.loc()), tensors[2].getType(), tensors[0], tensors[1], rhs1_lbls_value, rhs2_lbls_value, lhs_lbls_value, affineMapArrayAttr, SemiringAttr, MaskingAttr); + auto op = builder.create(loc(tensor_op.loc()), tensors[2].getType(), tensors[0], tensors[1], rhs1_lbls_value, rhs2_lbls_value, lhs_lbls_value, affineMapArrayAttr, SemiringAttr, tensors[2], MaskingAttr); op.getOperation()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); op.getOperation()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); diff --git a/include/comet/Dialect/TensorAlgebra/IR/TAOps.td b/include/comet/Dialect/TensorAlgebra/IR/TAOps.td index 7ca72a07..dfe74bba 100644 --- a/include/comet/Dialect/TensorAlgebra/IR/TAOps.td +++ b/include/comet/Dialect/TensorAlgebra/IR/TAOps.td @@ -353,6 +353,7 @@ def TensorMultOp : TA_Op<"mul", [Pure, AttrSizedOperandSegments]>{ let arguments = (ins TA_AnyTensor:$rhs1, TA_AnyTensor:$rhs2, Variadic:$rhs1_index_labels, Variadic:$rhs2_index_labels, Variadic:$result_index_labels, AffineMapArrayAttr:$indexing_maps, StrAttr:$semiring, + Optional:$lhs, OptionalAttr:$MaskType, /// TODO: try to use: DefaultValuedAttr Optional:$mask); @@ -373,6 +374,7 @@ def TensorElewsMultOp : TA_Op<"elews_mul", [Pure, AttrSizedOperandSegments]>{ Variadic:$rhs1_index_labels, Variadic:$rhs2_index_labels, Variadic:$result_index_labels, AffineMapArrayAttr:$indexing_maps, StrAttr:$semiring, + Optional:$lhs, OptionalAttr:$MaskType); let results = (outs TA_AnyTensor); @@ -403,7 +405,7 @@ def TensorDimOp : TA_Op<"dim", [Pure]>{ def TransposeOp : TA_Op<"transpose",[Pure, AttrSizedOperandSegments]> { let summary = "transpose operation"; - let arguments = (ins TA_AnyTensor:$rhs, Variadic:$rhs_index_labels, Variadic:$result_index_labels, AffineMapArrayAttr:$indexing_maps); + let arguments = (ins TA_AnyTensor:$rhs, Variadic:$rhs_index_labels, Variadic:$result_index_labels, AffineMapArrayAttr:$indexing_maps, Optional:$lhs); let results = (outs TA_AnyTensor); @@ -437,6 +439,7 @@ def TensorAddOp : TA_Op<"add", Variadic:$rhs1_index_labels, Variadic:$rhs2_index_labels, Variadic:$result_index_labels, AffineMapArrayAttr:$indexing_maps, StrAttr:$semiring, + Optional:$lhs, OptionalAttr:$MaskType); let results = (outs TA_AnyTensor); @@ -456,6 +459,7 @@ def TensorSubtractOp : TA_Op<"subtract", Variadic:$rhs1_index_labels, Variadic:$rhs2_index_labels, Variadic:$result_index_labels, AffineMapArrayAttr:$indexing_maps, StrAttr:$semiring, + Optional:$lhs, OptionalAttr:$MaskType); let results = (outs TA_AnyTensor); diff --git a/include/comet/Dialect/Utils/Utils.h b/include/comet/Dialect/Utils/Utils.h index 4f841241..30be14c5 100644 --- a/include/comet/Dialect/Utils/Utils.h +++ b/include/comet/Dialect/Utils/Utils.h @@ -166,13 +166,13 @@ namespace mlir const std::vector &rhs2Labels, const std::vector &lhsLabels); - void createTensorContraction(Location loc, Value rhs1Tensor, - ArrayRef rhs1Labels, - Value rhs2Tensor, - ArrayRef rhs2Labels, Value lhsTensor, - ArrayRef lhsLabels, - ConversionPatternRewriter &rewriter, - double beta = 0.0); + // void createTensorContraction(Location loc, Value rhs1Tensor, + // ArrayRef rhs1Labels, + // Value rhs2Tensor, + // ArrayRef rhs2Labels, Value lhsTensor, + // ArrayRef lhsLabels, + // ConversionPatternRewriter &rewriter, + // double beta = 0.0); std::vector constructPermutationMapAttr(const std::vector &rhs_labels, const std::vector &lhs_labels); diff --git a/lib/Conversion/IndexTreeToSCF/IndexTreeToSCF.cpp b/lib/Conversion/IndexTreeToSCF/IndexTreeToSCF.cpp index 03bdd91a..eacf49ed 100644 --- a/lib/Conversion/IndexTreeToSCF/IndexTreeToSCF.cpp +++ b/lib/Conversion/IndexTreeToSCF/IndexTreeToSCF.cpp @@ -1224,7 +1224,7 @@ namespace /// RHS operand is a constant value. return tensor; } else if((tensor_type = llvm::dyn_cast(tensor.getType()))){ - return rewriter.create(loc, tensor_type.getElementType(), tensor, crds); + return rewriter.create(loc, tensor_type.getElementType(), tensor, positions); } else { // LHS may not be constant (i.e. if we are inserting into a tensor that we need to resize), // so cannot directly lower like we can the RHS diff --git a/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp b/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp index 80905332..ed1ad8ba 100644 --- a/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp +++ b/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp @@ -194,7 +194,8 @@ mlir::LogicalResult generalIndexOperationRewrite( Value rhs1_tensor = mult_op.getRhs1(); Value rhs2_tensor = mult_op.getRhs2(); - Value lhs_tensor; + Value lhs_tensor = mult_op.getLhs(); + SmallVector dims; auto shapeT = cast(mult_op.getResult().getType()); ArrayAttr indexing_maps = cast(mult_op.getIndexingMaps()); @@ -223,18 +224,22 @@ mlir::LogicalResult generalIndexOperationRewrite( } assert(dims.size() == shapeT.getNumDynamicDims()); - if(auto spTensorT = dyn_cast(op->getResultTypes()[0])) - { - lhs_tensor = rewriter.create(loc, shapeT, ValueRange(dims), false); - } - else if (auto tensorT = dyn_cast(op->getResultTypes()[0])) + if(!lhs_tensor) { - lhs_tensor = rewriter.create(loc, shapeT, ValueRange(dims)); - rewriter.create( + + if(auto spTensorT = dyn_cast(op->getResultTypes()[0])) + { + lhs_tensor = rewriter.create(loc, shapeT, ValueRange(dims), false); + } + else if (auto tensorT = dyn_cast(op->getResultTypes()[0])) + { + lhs_tensor = rewriter.create(loc, shapeT, ValueRange(dims)); + rewriter.create( loc, lhs_tensor, rewriter.getZeroAttr(tensorT.getElementType())); // initialize the tensor to zero - // lhs_tensor = rewriter.create(loc, shapeT, ValueRange(dims)); + // lhs_tensor = rewriter.create(loc, shapeT, ValueRange(dims)); + } } comet_vdump(rhs1_tensor); diff --git a/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp b/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp index 697660c2..1ac8c2b2 100644 --- a/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp +++ b/lib/Conversion/TensorAlgebraToSCF/TensorAlgebraToSCF.cpp @@ -195,7 +195,7 @@ namespace /// There are tensors for transpose operation: input and output tensors unsigned int tensors_num = 2; - Value lhs; + Value lhs = op.getLhs(); if (auto tensorT = dyn_cast(inputType)) { /// for dense @@ -210,8 +210,10 @@ namespace dims.push_back(dim); } } - - lhs = rewriter.create(loc, op.getResult().getType(), dims); + if(!lhs) + { + lhs = rewriter.create(loc, op.getResult().getType(), dims); + } comet_vdump(lhs); // auto outputMemref = lhs.getDefiningOp()->getOperand(0); diff --git a/lib/Dialect/IndexTree/Transforms/KernelFusion.cpp b/lib/Dialect/IndexTree/Transforms/KernelFusion.cpp index fe4b15a4..d2f8c293 100644 --- a/lib/Dialect/IndexTree/Transforms/KernelFusion.cpp +++ b/lib/Dialect/IndexTree/Transforms/KernelFusion.cpp @@ -832,10 +832,10 @@ void fuseITrees(IndexTreeOp new_itree, OpBuilder::InsertionGuard guard(rewriter); rewriter.setInsertionPoint(yield_op); /// TODO: Parallel execution does not seem to work properly when fusion is in place... - for(auto indexOp : host_index_ops) - { - indexOp.setIsParallel(false); - } + // for(auto indexOp : host_index_ops) + // { + // indexOp.setIsParallel(false); + // } /// Fuse each other itree to the new itree. for (uint32_t tree_i = 1; tree_i < num_itrees; ++tree_i) { /// Create Index Ops: 1) Record the common index ops, then 2) add the new index ops. diff --git a/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp b/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp index 62ecd22c..9dfe3d86 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp @@ -230,6 +230,7 @@ mlir::LogicalResult FindOptimalTCFactorizationPass::FindOptimalTCFactorization(t comet_pdump(op); auto loc = op->getLoc(); auto lhsOp = op; + auto out_val = op.getLhs(); std::vector MultOpsToRemove; std::vector LTOpsToRemove; @@ -581,6 +582,7 @@ mlir::LogicalResult FindOptimalTCFactorizationPass::FindOptimalTCFactorization(t auto MaskingAttr = builder.getStringAttr("none"); Value tcop = builder.create(loc, newType, newRhs1, newRhs2, newRhs1Labels, newRhs2Labels, newResultLabels, affineMapArrayAttr, SemiringAttr, + nullptr, MaskingAttr, nullptr); tcop.getDefiningOp()->setAttr("__alpha__", builder.getF64FloatAttr(1.0)); tcop.getDefiningOp()->setAttr("__beta__", builder.getF64FloatAttr(0.0)); @@ -591,6 +593,7 @@ mlir::LogicalResult FindOptimalTCFactorizationPass::FindOptimalTCFactorization(t comet_debug() << "are they previous multop\n"; MultOpsToRemove.front()->replaceAllUsesWith(ValueRange(newRhs1)); /// replace the last multop with the new one + cast(newRhs1.getDefiningOp()).getLhsMutable().assign(out_val); for (auto oldTcOp : ArrayRef(MultOpsToRemove).drop_front()) { removeAllUsers(oldTcOp); diff --git a/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp b/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp index 1da5e370..9b1e8930 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp @@ -193,7 +193,7 @@ namespace comet_vdump(setnewop); comet_debug() << "\n"; - Value rhs1Tensor = operands[0], rhs2Tensor = operands[1], lhsTensor; + Value rhs1Tensor = operands[0], rhs2Tensor = operands[1], lhsTensor = multop.getLhs(); auto rhs1TensorType = cast(rhs1Tensor.getType()); auto rhs2TensorType = cast(rhs2Tensor.getType()); @@ -227,7 +227,10 @@ namespace auto zero = rewriter.create(loc, rewriter.getZeroAttr( shapeT.getElementType())); - lhsTensor = rewriter.create(loc, shapeT, zero, ValueRange(dims)); + if(!lhsTensor) + { + lhsTensor = rewriter.create(loc, shapeT, zero, ValueRange(dims)); + } auto lhsTensorType = cast(lhsTensor.getType()); diff --git a/lib/Dialect/Utils/Utils.cpp b/lib/Dialect/Utils/Utils.cpp index 990390d0..62c1f47a 100644 --- a/lib/Dialect/Utils/Utils.cpp +++ b/lib/Dialect/Utils/Utils.cpp @@ -2159,92 +2159,92 @@ namespace mlir return result; } - void createTensorContraction(Location loc, Value rhs1Tensor, - ArrayRef rhs1Labels, - Value rhs2Tensor, - ArrayRef rhs2Labels, Value lhsTensor, - ArrayRef lhsLabels, - ConversionPatternRewriter &rewriter, - double beta) - { - - std::map allLabels; - double alpha = 1.0; - auto rhs1AlphaAttr = rhs1Tensor.getDefiningOp()->getAttr("__alpha__"); - auto rhs2AlphaAttr = rhs2Tensor.getDefiningOp()->getAttr("__alpha__"); - - alpha *= cast(rhs1AlphaAttr).getValueAsDouble(); - alpha *= cast(rhs2AlphaAttr).getValueAsDouble(); - - unsigned idx = 0; - for (auto lbl : rhs1Labels) - { - if (allLabels.find(lbl.getDefiningOp()) == allLabels.end()) - { - allLabels[lbl.getDefiningOp()] = - getAffineDimExpr(idx++, rewriter.getContext()); - } - } - - for (auto lbl : rhs2Labels) - { - if (allLabels.find(lbl.getDefiningOp()) == allLabels.end()) - { - allLabels[lbl.getDefiningOp()] = - getAffineDimExpr(idx++, rewriter.getContext()); - } - } - - for (auto lbl : lhsLabels) - { - if (allLabels.find(lbl.getDefiningOp()) == allLabels.end()) - { - allLabels[lbl.getDefiningOp()] = - getAffineDimExpr(idx++, rewriter.getContext()); - } - } - - std::vector rhs1Exprs; - std::vector rhs2Exprs; - std::vector lhsExprs; - - for (const auto &lbl : rhs1Labels) - { - rhs1Exprs.push_back(allLabels[lbl.getDefiningOp()]); - } - - for (const auto &lbl : rhs2Labels) - { - rhs2Exprs.push_back(allLabels[lbl.getDefiningOp()]); - } - - for (const auto &lbl : lhsLabels) - { - lhsExprs.push_back(allLabels[lbl.getDefiningOp()]); - } - - auto context = rewriter.getContext(); - SmallVector affineMaps{ - mlir::AffineMap::get(idx, 0, rhs1Exprs, context), - mlir::AffineMap::get(idx, 0, rhs2Exprs, context), - mlir::AffineMap::get(idx, 0, lhsExprs, context)}; - - auto affineMapArrayAttr = rewriter.getAffineMapArrayAttr(affineMaps); - comet_debug() << "\n"; - - - auto SemiringAttr = rewriter.getStringAttr("none"); - auto MaskingAttr = rewriter.getStringAttr("none"); - auto tc = rewriter.create(loc, lhsTensor.getType(), - rhs1Tensor, rhs2Tensor, - rhs1Labels, rhs2Labels, lhsLabels, affineMapArrayAttr, - SemiringAttr, MaskingAttr, - nullptr); /// TODO: masking is an optional operand - tc.getOperation()->setAttr("__alpha__", rewriter.getF64FloatAttr(alpha)); - tc.getOperation()->setAttr("__beta__", rewriter.getF64FloatAttr(beta)); - comet_debug() << " "; - comet_vdump(tc); - } + // void createTensorContraction(Location loc, Value rhs1Tensor, + // ArrayRef rhs1Labels, + // Value rhs2Tensor, + // ArrayRef rhs2Labels, Value lhsTensor, + // ArrayRef lhsLabels, + // ConversionPatternRewriter &rewriter, + // double beta) + // { + + // std::map allLabels; + // double alpha = 1.0; + // auto rhs1AlphaAttr = rhs1Tensor.getDefiningOp()->getAttr("__alpha__"); + // auto rhs2AlphaAttr = rhs2Tensor.getDefiningOp()->getAttr("__alpha__"); + + // alpha *= cast(rhs1AlphaAttr).getValueAsDouble(); + // alpha *= cast(rhs2AlphaAttr).getValueAsDouble(); + + // unsigned idx = 0; + // for (auto lbl : rhs1Labels) + // { + // if (allLabels.find(lbl.getDefiningOp()) == allLabels.end()) + // { + // allLabels[lbl.getDefiningOp()] = + // getAffineDimExpr(idx++, rewriter.getContext()); + // } + // } + + // for (auto lbl : rhs2Labels) + // { + // if (allLabels.find(lbl.getDefiningOp()) == allLabels.end()) + // { + // allLabels[lbl.getDefiningOp()] = + // getAffineDimExpr(idx++, rewriter.getContext()); + // } + // } + + // for (auto lbl : lhsLabels) + // { + // if (allLabels.find(lbl.getDefiningOp()) == allLabels.end()) + // { + // allLabels[lbl.getDefiningOp()] = + // getAffineDimExpr(idx++, rewriter.getContext()); + // } + // } + + // std::vector rhs1Exprs; + // std::vector rhs2Exprs; + // std::vector lhsExprs; + + // for (const auto &lbl : rhs1Labels) + // { + // rhs1Exprs.push_back(allLabels[lbl.getDefiningOp()]); + // } + + // for (const auto &lbl : rhs2Labels) + // { + // rhs2Exprs.push_back(allLabels[lbl.getDefiningOp()]); + // } + + // for (const auto &lbl : lhsLabels) + // { + // lhsExprs.push_back(allLabels[lbl.getDefiningOp()]); + // } + + // auto context = rewriter.getContext(); + // SmallVector affineMaps{ + // mlir::AffineMap::get(idx, 0, rhs1Exprs, context), + // mlir::AffineMap::get(idx, 0, rhs2Exprs, context), + // mlir::AffineMap::get(idx, 0, lhsExprs, context)}; + + // auto affineMapArrayAttr = rewriter.getAffineMapArrayAttr(affineMaps); + // comet_debug() << "\n"; + + + // auto SemiringAttr = rewriter.getStringAttr("none"); + // auto MaskingAttr = rewriter.getStringAttr("none"); + // auto tc = rewriter.create(loc, lhsTensor.getType(), + // rhs1Tensor, rhs2Tensor, + // rhs1Labels, rhs2Labels, lhsLabels, affineMapArrayAttr, + // SemiringAttr, MaskingAttr, + // nullptr); /// TODO: masking is an optional operand + // tc.getOperation()->setAttr("__alpha__", rewriter.getF64FloatAttr(alpha)); + // tc.getOperation()->setAttr("__beta__", rewriter.getF64FloatAttr(beta)); + // comet_debug() << " "; + // comet_vdump(tc); + // } std::vector constructPermutationMapAttr(const std::vector &rhs_labels, const std::vector &lhs_labels) From aa5f3d87fabbd8e41e8792f0e3df548b2e39378e Mon Sep 17 00:00:00 2001 From: Polykarpos Thomadakis Date: Thu, 24 Apr 2025 11:28:49 -0700 Subject: [PATCH 09/11] During lowering from TA to IT, properly handle cases where one or more of the operands in a ComputeOp are also found in the LHS --- .../TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp b/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp index ed1ad8ba..4ee45565 100644 --- a/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp +++ b/lib/Conversion/TensorAlgebraToIndexTree/TensorAlgebraToIndexTree.cpp @@ -263,6 +263,14 @@ mlir::LogicalResult generalIndexOperationRewrite( Region* body = &itree_op.getRegion(); loc = body->getLoc(); Block* block = rewriter.createBlock(body, {}, TypeRange(lhs_tensor.getType()), {lhs_tensor.getLoc()}); + if(rhs1_tensor == lhs_tensor) + { + rhs1_tensor = block->getArgument(0); + } + if(rhs2_tensor == lhs_tensor) + { + rhs2_tensor = block->getArgument(0); + } lhs_tensor = block->getArgument(0); comet_vdump(itree_op); comet_pdump(block); From 1d786e2b72c877f25a0a4474fc13acd4e4b2e816 Mon Sep 17 00:00:00 2001 From: Polykarpos Thomadakis Date: Thu, 24 Apr 2025 13:01:25 -0700 Subject: [PATCH 10/11] 1. When possible, store the intermediate results of an expression in the output (lhs) instead of creating new allocations. 2. Fixed issue with opt-dense-transpose not lowering properly when its result is returend from a function. 3. Added significant more functionality for cometpy and make sure we avoid unnecessary copies. --- frontends/comet_dsl/mlir/MLIRGen.cpp | 26 ++ frontends/numpy-scipy/cometpy/MLIRGen/ops.py | 173 +++++++----- frontends/numpy-scipy/cometpy/comet.py | 254 +++++++++++++----- .../kernels/test_triangleCount_SandiaLL.py | 1 - .../integration_tests/opts/test_fusion.py | 4 +- .../Transforms/LinalgTransforms.cpp | 3 +- 6 files changed, 319 insertions(+), 142 deletions(-) diff --git a/frontends/comet_dsl/mlir/MLIRGen.cpp b/frontends/comet_dsl/mlir/MLIRGen.cpp index e9391f2a..96740875 100644 --- a/frontends/comet_dsl/mlir/MLIRGen.cpp +++ b/frontends/comet_dsl/mlir/MLIRGen.cpp @@ -51,6 +51,7 @@ #include "llvm/ADT/ScopedHashTable.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/StringRef.h" +#include "llvm/Support/Casting.h" #include "llvm/Support/raw_ostream.h" #include #include @@ -445,6 +446,31 @@ namespace } mlir::Value lhs = mlirGen(*binop.getLHS(), rhs_lbls, "", binop.getOp() != '*' ? out_val : nullptr); + mlir::Value lhs_res_val = nullptr; + if(auto tensor_op = mlir::dyn_cast_if_present(lhs.getDefiningOp())) + { + lhs_res_val = tensor_op.getLhs(); + } + else if(auto tensor_op = mlir::dyn_cast_if_present(lhs.getDefiningOp())) + { + lhs_res_val = tensor_op.getLhs(); + } + else if(auto tensor_op = mlir::dyn_cast_if_present(lhs.getDefiningOp())) + { + lhs_res_val = tensor_op.getLhs(); + } + else if(auto tensor_op = mlir::dyn_cast_if_present(lhs.getDefiningOp())) + { + lhs_res_val = tensor_op.getLhs(); + } + else if(auto tensor_op = mlir::dyn_cast_if_present(lhs.getDefiningOp())) + { + lhs_res_val = tensor_op.getLhs(); + } + if(mlir::isa(binop.getLHS()) && lhs_res_val) + { + out_val = lhs; + } if (!lhs) return nullptr; mlir::Value rhs = mlirGen(*binop.getRHS(), lhs_lbls, "", nullptr); diff --git a/frontends/numpy-scipy/cometpy/MLIRGen/ops.py b/frontends/numpy-scipy/cometpy/MLIRGen/ops.py index 46897ff3..f946b792 100644 --- a/frontends/numpy-scipy/cometpy/MLIRGen/ops.py +++ b/frontends/numpy-scipy/cometpy/MLIRGen/ops.py @@ -15,17 +15,19 @@ def __repr__(self): class Operation: - def __init__(self, operands, result_types, startsBody = False, endsBody = False): + def __init__(self, block, index_in_block, operands, result_types, startsBody = False, endsBody = False): self.startsBody = startsBody self.endsBody = endsBody self.operands = operands - self.results = [Symbol(t) for t in result_types] + self.index_in_block = index_in_block + self.block = block + self.results = [Symbol(t, self, block) for t in result_types] class BinaryElementwiseOp(Operation): - def __init__(self, lhs, rhs, _type = None): - super().__init__([lhs, rhs], [_type if _type else lhs.type]) + def __init__(self, block, iiblock, lhs, rhs, _type = None): + super().__init__(block, iiblock, [lhs, rhs], [_type if _type else lhs.type]) self.lhs = lhs self.rhs = rhs @@ -46,8 +48,8 @@ def dump(self, name): class SubOp(BinaryElementwiseOp): - def __init__(self, lhs, rhs): - super().__init__(lhs, rhs) + def __init__(self, block, iiblock, lhs, rhs): + super().__init__(block, iiblock, lhs, rhs) def dump(self): name = self.text.render( @@ -62,8 +64,8 @@ def dump(self): class MulOp(BinaryElementwiseOp): - def __init__(self, lhs, rhs): - super().__init__(lhs, rhs) + def __init__(self, block, iiblock, lhs, rhs): + super().__init__(block, iiblock, lhs, rhs) def dump(self): name = self.text.render( @@ -77,7 +79,7 @@ def dump(self): ) class TensorIndexBasedOp(Operation): - def __init__(self, name, inputs, inputs_indices, res_indices, res_format, alpha = None, beta = None , mask = None, mask_type = None, semiring = None): + def __init__(self, block, iiblock, name, inputs, inputs_indices, res_indices, res_format, res_val = None, alpha = None, beta = None , mask = None, mask_type = None, semiring = None): self.name = name self.inputs = inputs self.inputs_indices = inputs_indices @@ -86,6 +88,7 @@ def __init__(self, name, inputs, inputs_indices, res_indices, res_format, alpha self.alpha = alpha self.beta = beta self.mask = mask + self.res_val = res_val self.res_indices = res_indices all_indices = inputs_indices[0][:] @@ -123,7 +126,6 @@ def __init__(self, name, inputs, inputs_indices, res_indices, res_format, alpha affine_maps.append(affine_map_res) self.indexing_maps = ",".join(affine_maps) - if res_format != DENSE: for input in inputs: if isinstance(input.type, TASparseTensorType): @@ -131,46 +133,54 @@ def __init__(self, name, inputs, inputs_indices, res_indices, res_format, alpha res_type = TASparseTensorType(res_shape, inputs[0].type.element_type, indices_type, res_format) else: res_type = TensorType(res_shape, inputs[0].type.element_type) - - super().__init__([*inputs, *inputs_indices, res_indices], [res_type]) + if self.res_val: + super().__init__(block, iiblock, [*inputs, *inputs_indices, res_indices, self.res_val], [res_type]) + else: + super().__init__(block, iiblock, [*inputs, *inputs_indices, res_indices], [res_type]) def dump(self): - num_indices = len(self.res_indices) - for indices in self.inputs_indices : - num_indices += len(indices) + input_index_label_nums = ", ".join([str(len(indices)) for indices in self.inputs_indices]) + all_index_label_nums = ", ".join([input_index_label_nums, (str(len(self.res_indices)))]) input_types = ", ".join([ ",".join([f'{input.type}' for input in self.inputs]) , ", ".join([f'{index.type}' for indices in self.inputs_indices for index in indices]) , ", ".join([f'{oper.type}' for oper in self.res_indices])]) if self.mask != None: input_types += f", {self.mask.type}" + if self.res_val != None: + input_types += f", {self.res_val.type}" + if self.mask_type: + segment_sizes = ", ".join([", ".join('1' * len(self.inputs)), all_index_label_nums, '1' if self.res_val else '0', '1' if self.mask_type != 'none' else '0']) + else: + segment_sizes = ", ".join([", ".join('1' * len(self.inputs)), all_index_label_nums, '1' if self.res_val else '0']) return self.tensor_binary_op_text.render( ssa = self.results[0].ssa, name = self.name, - inputs = ",".join([f'%{input.ssa}' for input in self.inputs]), - inputs_indices = ",".join([f'%{index.ssa}' for indices in self.inputs_indices for index in indices]), - res_indices = ",".join([f'%{input.ssa}' for input in self.res_indices]), + inputs = ", ".join([f'%{input.ssa}' for input in self.inputs]), + inputs_indices = ", ".join([f'%{index.ssa}' for indices in self.inputs_indices for index in indices]), + res_indices = ", ".join([f'%{input.ssa}' for input in self.res_indices]), mask_type = self.mask_type, mask = self.mask, semiring = self.semiring, input_types = input_types, res_type = self.results[0].type, affine_maps = self.indexing_maps, - operand_segment_sizes = ", ".join([", ".join('1' * len(self.inputs)), str(num_indices), '1' if self.mask_type and self.mask_type != 'none' else '0']), + operand_segment_sizes = segment_sizes, + res_val = self.res_val, alpha = self.alpha, beta = self.beta, ) tensor_binary_op_text = jinja2.Template( - '%{{ssa}} = "{{name}}" ({{inputs}}, {{inputs_indices}}, {{res_indices}} {%if mask!=None%}, %{{mask.ssa}}{%endif%} ) <{ {% if mask_type%} MaskType = "{{mask_type}}", {% endif%} indexing_maps = [{{affine_maps}}] {% if semiring%}, operandSegmentSizes = array, semiring="{{semiring}}" {%endif%}}> {%if alpha != None or beta != None %} { {%endif%} {%if alpha!= None%} __alpha__ = {{alpha}} : f64 {%endif%} {%if alpha!= None and beta!= None %}, {%endif%} {%if beta!= None %} __beta__ = {{beta}}: f64 {%endif%} {%if alpha!= None or beta!= None %}} {%endif%} : ({{input_types}}) -> {{res_type}}', + '%{{ssa}} = "{{name}}" ({{inputs}}, {{inputs_indices}}, {{res_indices}}{%if res_val!=None%}, %{{res_val.ssa}}{%endif%}{%if mask!=None%}, %{{mask.ssa}}{%endif%}) <{ {% if mask_type%} MaskType = "{{mask_type}}", {% endif%} indexing_maps = [{{affine_maps}}], operandSegmentSizes = array{% if semiring%} , semiring="{{semiring}}" {%endif%}}> {%if alpha != None or beta != None %} { {%endif%} {%if alpha!= None%} __alpha__ = {{alpha}} : f64 {%endif%} {%if alpha!= None and beta!= None %}, {%endif%} {%if beta!= None %} __beta__ = {{beta}}: f64 {%endif%} {%if alpha!= None or beta!= None %}} {%endif%} : ({{input_types}}) -> {{res_type}}', undefined=jinja2.StrictUndefined, ) class TensorBinaryOp(TensorIndexBasedOp): - def __init__(self, name, lhs, rhs, lhs_indices, rhs_indices, res_indices, alpha, beta, mask, masktype, semiring, res_format): - super().__init__(name, [lhs, rhs], [lhs_indices, rhs_indices], res_indices, res_format, alpha, beta, mask, masktype, semiring) + def __init__(self, block, iiblock, name, lhs, rhs, lhs_indices, rhs_indices, res_indices, res_val, alpha, beta, mask, masktype, semiring, res_format): + super().__init__(block, iiblock, name, [lhs, rhs], [lhs_indices, rhs_indices], res_indices, res_format, res_val, alpha, beta, mask, masktype, semiring) self.lhs = lhs self.lhs_indices = lhs_indices self.rhs = rhs @@ -185,8 +195,8 @@ def dump(self): class TensorElewiseBinaryOp(TensorBinaryOp): - def __init__(self, name, lhs, rhs, lhs_indices, alpha, beta, semiring, format): - super().__init__(name, lhs, rhs, lhs_indices, lhs_indices, lhs_indices, alpha, beta, None, "none", semiring, format) + def __init__(self, block, iiblock, name, lhs, rhs, lhs_indices, res_val, alpha, beta, semiring, format): + super().__init__(block, iiblock, name, lhs, rhs, lhs_indices, lhs_indices, lhs_indices, res_val, alpha, beta, None, None, semiring, format) def dump(self): return super().dump() @@ -194,8 +204,8 @@ def dump(self): class TensorIndexLabelOp(Operation): - def __init__(self): - super().__init__([], [TensorIndexType()]) + def __init__(self, block, iiblock): + super().__init__(block, iiblock, [], [TensorIndexType()]) def dump(self): return self.text.render( @@ -209,32 +219,32 @@ def dump(self): class TensorAddOp(TensorElewiseBinaryOp): - def __init__(self, lhs, rhs, res_indices, format, alpha = 1.0, beta= 0.0, semiring= "noop_plusxy"): - super().__init__('ta.add', lhs, rhs, res_indices, alpha, beta, semiring, format) + def __init__(self, block, iiblock, lhs, rhs, res_indices, format, res_val, alpha = 1.0, beta= 0.0, semiring= "noop_plusxy"): + super().__init__(block, iiblock, 'ta.add', lhs, rhs, res_indices, res_val, alpha, beta, semiring, format) def dump(self): return super().dump() class TensorSubOp(TensorElewiseBinaryOp): - def __init__(self, lhs, rhs, res_indices, format, alpha = 1.0, beta= 0.0, semiring= "noop_minus"): - super().__init__('ta.subtract', lhs, rhs, res_indices, alpha, beta, semiring, format) + def __init__(self, block, iiblock, lhs, rhs, res_indices, format, res_val, alpha = 1.0, beta= 0.0, semiring= "noop_minus"): + super().__init__(block, iiblock, 'ta.subtract', lhs, rhs, res_indices, res_val, alpha, beta, semiring, format) def dump(self): return super().dump() class TensorMulOp(TensorElewiseBinaryOp): - def __init__(self, lhs, rhs, res_indices, format, alpha = 1.0, beta= 0.0, semiring= "noop_times"): - super().__init__('ta.elews_mul', lhs, rhs, res_indices, alpha, beta, semiring, format) + def __init__(self, block, iiblock, lhs, rhs, res_indices, format, res_val, alpha = 1.0, beta= 0.0, semiring= "noop_times"): + super().__init__(block, iiblock, 'ta.elews_mul', lhs, rhs, res_indices, res_val, alpha, beta, semiring, format) def dump(self): return super().dump() class TensorTransposeOp(TensorIndexBasedOp): - def __init__(self, input, input_indices, res_indices, res_format): - super().__init__('ta.transpose', [input], [input_indices], res_indices, res_format) + def __init__(self, block, iiblock, input, input_indices, res_indices, res_format, res_val): + super().__init__(block, iiblock, 'ta.transpose', [input], [input_indices], res_indices, res_format, res_val) self.input = input self.input_indices = input_indices self.res_indices = res_indices @@ -245,8 +255,8 @@ def dump(self): class TensorSumOp(Operation): - def __init__(self, value): - super().__init__([value], [value.type.element_type]) + def __init__(self, block, iiblock, value): + super().__init__(block, iiblock, [value], [value.type.element_type]) self.value = value def dump(self): @@ -263,8 +273,8 @@ def dump(self): class TensorSetOp(Operation): - def __init__(self, src, dst, beta): - super().__init__([src, dst], []) + def __init__(self, block, iiblock, src, dst, beta): + super().__init__(block, iiblock, [src, dst], []) self.src = src self.dst = dst self.beta = beta @@ -286,8 +296,8 @@ def dump(self): class TensorPrintOp(Operation): - def __init__(self, input): - super().__init__([input], []) + def __init__(self, block, iiblock, input): + super().__init__(block, iiblock, [input], []) self.input = input def dump(self): @@ -304,14 +314,14 @@ def dump(self): class TensorMatMultOp(TensorBinaryOp): - def __init__(self, lhs, rhs, lhs_indices, rhs_indices, res_indices, res_format, alpha = 1.0, beta = 0.0, mask = None, masktype='none', semiring='plusxy_times'): - super().__init__('ta.mul', lhs, rhs, lhs_indices, rhs_indices, res_indices, alpha, beta, mask, masktype, semiring, res_format) + def __init__(self, block, iiblock, lhs, rhs, lhs_indices, rhs_indices, res_indices, res_format, res_val, alpha = 1.0, beta = 0.0, mask = None, masktype='none', semiring='plusxy_times'): + super().__init__(block, iiblock, 'ta.mul', lhs, rhs, lhs_indices, rhs_indices, res_indices, res_val, alpha, beta, mask, masktype, semiring, res_format) class AddOp(BinaryElementwiseOp): - def __init__(self, lhs, rhs): - super().__init__(lhs, rhs) + def __init__(self, bloc, iiblock, lhs, rhs): + super().__init__(bloc, iiblock, lhs, rhs) def dump(self): name = self.text.render( @@ -327,7 +337,7 @@ def dump(self): class ConstantOp(Operation): - def __init__(self, val): + def __init__(self, block, iiblock, val): self.value = val if isinstance(val, int): @@ -335,7 +345,7 @@ def __init__(self, val): elif isinstance(val, float): _type = 'f64' - super().__init__([], [_type]) + super().__init__(block, iiblock, [], [_type]) def dump(self): return self.text.render( @@ -351,8 +361,8 @@ def dump(self): class ToMemrefOp(Operation): - def __init__(self, src): - super().__init__([src], [MemrefType(src.type.shape, src.type.element_type)]) + def __init__(self, block, iiblock, src): + super().__init__(block, iiblock, [src], [MemrefType(src.type.shape, src.type.element_type)]) self.src = src @@ -368,11 +378,30 @@ def dump(self): '%{{ssa}} = bufferization.to_memref %{{src}}: {{res_type}} ' ) +class ToTensorOp(Operation): + + def __init__(self, block, iiblock, src): + super().__init__(block, iiblock, [src], [TensorType(src.type.shape, src.type.element_type)]) + self.src = src + + + def dump(self): + return self.text.render( + ssa = self.results[0].ssa, + src = self.src.ssa, + src_type = self.src.type, + ) + + + text = jinja2.Template ( + '%{{ssa}} = bufferization.to_tensor %{{src}} restrict writable : {{src_type}} ' + ) + class LoadOp(Operation): - def __init__(self, src, indices): - super().__init__([src] + indices, [src.type.element_type]) + def __init__(self, block, iiblock, src, indices): + super().__init__(block, iiblock, [src] + indices, [src.type.element_type]) self.src = src self.indices = indices @@ -391,8 +420,8 @@ def dump(self): class StoreOp(Operation): - def __init__(self, value, dst, indices): - super().__init__([value, dst] + indices, []) + def __init__(self, block, iiblock, value, dst, indices): + super().__init__(block, iiblock, [value, dst] + indices, []) self.value = value self.dst = dst self.indices = indices @@ -414,9 +443,11 @@ def dump(self): class Symbol: curr_ssa = 0 - def __init__(self, _type): + def __init__(self, _type, defining_op, block): self.ssa = Symbol.curr_ssa self.type = _type + self.defining_op = defining_op + self.block = block Symbol.curr_ssa += 1 class SymbolTable: @@ -436,8 +467,8 @@ def find(self, id): class YieldOp(Operation): - def __init__(self, values): - super().__init__(values, [v.type for v in values], endsBody=True) + def __init__(self, block, iiblock, values): + super().__init__(block, iiblock, values, [v.type for v in values], endsBody=True) self. values = values def dump(self): @@ -453,8 +484,8 @@ def dump(self): class ReduceOp(Operation): - def __init__(self, values): - super().__init__(values, [v.type for v in values], endsBody=True) + def __init__(self, block, iiblock, values): + super().__init__(block, iiblock, values, [v.type for v in values], endsBody=True) self. values = values def dump(self): @@ -471,8 +502,8 @@ def dump(self): class ReturnOp(Operation): - def __init__(self, values): - super().__init__(values, [v.type for v in values], endsBody=True) + def __init__(self, block, iiblock, values): + super().__init__(block, iiblock, values, [v.type for v in values], endsBody=True) self. values = values def dump(self): @@ -489,8 +520,8 @@ def dump(self): class OperationWithBody(Operation): - def __init__(self, operands, return_types): - super().__init__(operands, return_types, True) + def __init__(self, block, iiblock, operands, return_types): + super().__init__(block, iiblock, operands, return_types, True) self.body = [] self.identation = 0 @@ -514,8 +545,8 @@ def dump(self) -> str: class ModuleOp(OperationWithBody): - def __init__(self): - super().__init__([], []) + def __init__(self, block, iiblock): + super().__init__(block, iiblock, [], []) def dump(self) -> str: op = self.text.render() @@ -536,12 +567,14 @@ class FuncOp(OperationWithBody): def __init__( self, + block, + iiblock, func_name: str, inputs, return_types, private: bool, ) -> None: - super().__init__([], return_types) + super().__init__(block, iiblock, [], return_types) self.func_name = func_name self.inputs = inputs self.return_types = return_types @@ -559,12 +592,12 @@ def dump(self) -> str: class ForOp(OperationWithBody): - def __init__(self, lb, ub, step): - super().__init__([lb, ub, step], []) + def __init__(self, block, iiblock, lb, ub, step): + super().__init__(block, iiblock, [lb, ub, step], []) self.lb = lb self.ub = ub self.step = step - self.iv = Symbol('index') + self.iv = Symbol('index', None, self) def dump(self): op = self.text.render( @@ -584,12 +617,12 @@ def dump(self): class ForAllOp(OperationWithBody): - def __init__(self, lb, ub, step): - super().__init__([lb, ub, step], []) + def __init__(self, block, iiblock, lb, ub, step): + super().__init__(block, iiblock, [lb, ub, step], []) self.lb = lb self.ub = ub self.step = step - self.iv = Symbol('index') + self.iv = Symbol('index', None, self) def dump(self): my_render = self.text.render( diff --git a/frontends/numpy-scipy/cometpy/comet.py b/frontends/numpy-scipy/cometpy/comet.py index 71ba31c1..92478109 100644 --- a/frontends/numpy-scipy/cometpy/comet.py +++ b/frontends/numpy-scipy/cometpy/comet.py @@ -45,7 +45,7 @@ from cometpy.MLIRGen.types import * #import time -def lower_einsum_expresion(args, kwargs, visitor): +def lower_einsum_expresion(args, kwargs, visitor, res_val = None): char_index_to_symbol_index = {} all_input_indices, res_indices = args[0].value.split('->') inputs_indices = all_input_indices.split(',') @@ -54,11 +54,12 @@ def lower_einsum_expresion(args, kwargs, visitor): input_symbol_indices = [] for c in input_indices: if c not in char_index_to_symbol_index: - char_index_to_symbol_index[c] = visitor.build(ops.TensorIndexLabelOp()).results[0] + char_index_to_symbol_index[c] = visitor.build(ops.TensorIndexLabelOp).results[0] input_symbol_indices.append(char_index_to_symbol_index[c]) inputs_symbol_indices.append(input_symbol_indices) operands = [visitor.visit(arg) for arg in args[1:len(inputs_indices)+1]] + operand_expr = [arg for arg in args[1:len(inputs_indices)+1]] semiring = None masktype = "none" mask = None @@ -85,7 +86,11 @@ def lower_einsum_expresion(args, kwargs, visitor): input = operands[0] input_indices = inputs_symbol_indices[0] res_format = visitor.sp_mattr_conversions[input.type.format] - res = visitor.build(ops.TensorTransposeOp(input, input_indices, res_symbol_indices, res_format)).results[0] + if isinstance(input, MemrefType): + input = visitor.build(ops.ToTensorOp, input).results[0] + if isinstance(operand_expr[0], ast.Name): + visitor.symbol_table.insert(operand_expr[0].id, input) + res = visitor.build(ops.TensorTransposeOp, input, input_indices, res_symbol_indices, res_format, res_val).results[0] lhs = res elif all_same and inputs_indices[0] == res_indices: # Elementwise multiplication @@ -100,12 +105,21 @@ def lower_einsum_expresion(args, kwargs, visitor): semiring = "noop_times" lhs, lhs_indices = operands[0], inputs_symbol_indices[0] - for rhs in operands[1:]: + if isinstance(lhs, MemrefType): + lhs = visitor.build(ops.ToTensorOp, lhs).results[0] + if isinstance(operand_expr[0], ast.Name): + visitor.symbol_table.insert(operand_expr[0].id, lhs) + + for i, rhs in enumerate(operands[1:]): res_format = visitor.sp_elw_mult_conversions[lhs.type.format][rhs.type.format] + if isinstance(rhs, MemrefType): + rhs = visitor.build(ops.ToTensorOp, rhs).results[0] + if isinstance(operand_expr[1 + i], ast.Name): + visitor.symbol_table.insert(operand_expr[0].id, rhs) if semiring: - res = visitor.build(ops.TensorMulOp(lhs, rhs, lhs_indices, res_format, semiring=semiring)).results[0] + res = visitor.build(ops.TensorMulOp, lhs, rhs, lhs_indices, res_format, res_val, 1.0, 0.0, semiring).results[0] else: - res = visitor.build(ops.TensorMulOp(lhs, rhs, lhs_indices, res_format)).results[0] + res = visitor.build(ops.TensorMulOp, lhs, rhs, lhs_indices, res_format, res_val).results[0] lhs = res else: # Tensor contraction @@ -127,10 +141,19 @@ def lower_einsum_expresion(args, kwargs, visitor): semiring += "plusxy" elif s2 == "second": semiring += "second" - + actual_res_val = None lhs, lhs_indices = operands[0], inputs_symbol_indices[0] + if isinstance(lhs, MemrefType): + lhs = visitor.build(ops.ToTensorOp, lhs).results[0] + if isinstance(operand_expr[0], ast.Name): + visitor.symbol_table.insert(operand_expr[0].id, lhs) for (i, (rhs, rhs_indices)) in enumerate(zip(operands[1:], inputs_symbol_indices[1:])): + if isinstance(rhs, MemrefType): + rhs = visitor.build(ops.ToTensorOp, rhs).results[0] + if isinstance(operand_expr[1 + i], ast.Name): + visitor.symbol_table.insert(operand_expr[0].id, rhs) if i == len(operands) - 2: + actual_res_val = res_val res_symbol_indices = [] for symbol in res_indices: res_symbol_indices.append(char_index_to_symbol_index[symbol]) @@ -144,10 +167,11 @@ def lower_einsum_expresion(args, kwargs, visitor): res_symbol_indices.append(indice) res_format = visitor.sp_matmult_conversions[lhs.type.format][rhs.type.format] + print('EINSUM', actual_res_val) if semiring: - res = visitor.build(ops.TensorMatMultOp(lhs, rhs, lhs_indices, rhs_indices, res_symbol_indices, res_format, 1.0, 0.0, mask, masktype, semiring=semiring)).results[0] + res = visitor.build(ops.TensorMatMultOp, lhs, rhs, lhs_indices, rhs_indices, res_symbol_indices, res_format, actual_res_val, 1.0, 0.0, mask, masktype, semiring).results[0] else: - res = visitor.build(ops.TensorMatMultOp(lhs, rhs, lhs_indices, rhs_indices, res_symbol_indices, res_format, 1.0, 0.0, mask, masktype)).results[0] + res = visitor.build(ops.TensorMatMultOp, lhs, rhs, lhs_indices, rhs_indices, res_symbol_indices, res_format, actual_res_val, 1.0, 0.0, mask, masktype).results[0] lhs, lhs_indices = res, res_symbol_indices if lhs.type.format != DENSE: @@ -159,19 +183,23 @@ def lower_einsum_expresion(args, kwargs, visitor): class NewAstParser(ast.NodeVisitor): - + ADDOP = 0 + SUBOP = 1 + MULOP = 2 + MATMULOP = 3 identation_default = 2 + def __init__(self,inputs): self.symbol_table = ops.SymbolTable() self.inputs = inputs - self.insertion_point = [self] + self.blocks = [self] + self.insertion_point = [self.blocks[-1], 0] self.identation = 0 self.return_type = [] self.pragma = '' self.body = [] self.need_opt_comp_workspace = False - self.build(ops.ModuleOp()) - + self.build(ops.ModuleOp) # Output formats when multiply matrices of different formats self.sp_matmult_conversions = { @@ -271,34 +299,50 @@ def __init__(self,inputs): self.sp_elw_add_sub_conversions[CSR][DENSE] = CSR self.sp_elw_add_sub_conversions[COO][DENSE] = COO + def set_insertion_point(self, op): + self.insertion_point = [op.block, op.index_in_block] + + def reset_insertion_point(self, insertion_point): + self.insertion_point = insertion_point + + def get_insertion_point(self): + return self.insertion_point def dump(self): return "\n".join([stmt.dump() for stmt in self.body]) - def build(self, op): + def build(self, op_ctor, *args): # self.mlir_code.append(" " * self.identation + op.dump()) - self.insertion_point[-1].body.append(op) + op = op_ctor(self.insertion_point[0], self.insertion_point[1], *args) + self.insertion_point[0].body.insert(self.insertion_point[1], op) + self.insertion_point[1] += 1 + for operation in self.insertion_point[0].body[self.insertion_point[1]:]: + operation.index_in_block += 1 if op.startsBody: self.identation += NewAstParser.identation_default op.identation = self.identation - self.insertion_point.append(op) + self.blocks.append(op) + self.insertion_point = [self.blocks[-1], 0] elif op.endsBody: self.identation -= NewAstParser.identation_default # self.mlir_code.append(" " * self.identation + "}") - self.insertion_point.pop() + self.blocks.pop() + self.insertion_point = [self.blocks[-1], len(self.blocks[-1].body)] return op def visit_FunctionDef(self, node): args = [] for arg, val in zip(node.args.args, self.inputs): if isinstance(val, np.ndarray) or sp.sparse.issparse(val): - new_arg = ops.Symbol(mlir_type_from_ndarray(val)) + new_arg = ops.Symbol(mlir_type_from_ndarray(val), None, None) else: - new_arg = ops.Symbol(mlir_type_from_python_type(val)) + new_arg = ops.Symbol(mlir_type_from_python_type(val), None, None) self.symbol_table.insert(arg.arg, new_arg) args.append(new_arg) - funcOp = self.build(ops.FuncOp(node.name, args, [], False)) + funcOp = self.build(ops.FuncOp, node.name, args, [], False) + for input in funcOp.inputs: + input.block = funcOp # for arg in args: # self.build(ops.TensorPrintOp(arg)) for stmt in node.body: @@ -308,7 +352,7 @@ def visit_FunctionDef(self, node): if self.return_type: funcOp.return_types = self.return_type else: - self.build(ops.ReturnOp([])) + self.build(ops.ReturnOp, []) # self.mlir_code.append("" *self.identation + "}") def visit_Comment(self, node): @@ -317,60 +361,124 @@ def visit_Comment(self, node): def visit_Assign(self, node): self.pragma = '' - rhs = self.visit(node.value) if isinstance(node.targets[0], ast.Name): + rhs = self.visit(node.value) self.symbol_table.insert(node.targets[0].id, rhs) elif isinstance(node.targets[0], ast.Subscript): - print(node.targets[0].slice) mem = self.visit(node.targets[0].value) indices = self.visit(node.targets[0].slice) if indices: # + rhs = self.visit(node.value) if not isinstance(indices, list): indices = [indices] if isinstance(mem.type, TensorType): - mem = self.build(ops.ToMemrefOp(mem)) - self.build(ops.StoreOp(rhs, mem.results[0], indices)) + if mem.block != self.insertion_point[0]: + ip = self.get_insertion_point() + ip_to_insert = self.insertion_point[0] + while mem.block != ip_to_insert.block: + print(mem.block) + print(ip_to_insert) + ip_to_insert = ip_to_insert.block + self.set_insertion_point(ip_to_insert) + memref = self.build(ops.ToMemrefOp, mem) + self.reset_insertion_point(ip) + mem = memref + else: + mem = self.build(ops.ToMemrefOp, mem) + + self.build(ops.StoreOp, rhs, mem.results[0], indices) + self.symbol_table.insert(node.targets[0].value.id, mem.results[0]) + else: + self.build(ops.StoreOp, rhs, mem, indices) else: - self.build(ops.TensorSetOp(rhs, mem, 0.0)) - self.build(ops.TensorPrintOp(mem)) + if isinstance(node.value, ast.BinOp): + if isinstance(mem.type, MemrefType): + if isinstance(node.targets[0].value, ast.Name): + mem = self.build(ops.ToTensorOp, mem).results[0] + self.symbol_table.insert(node.targets[0].value.id, mem) + else: + assert(False) + rhs = self.visit_BinOp(node.value, mem) + elif isinstance(node.value, ast.Call): + rhs = self.visit_Call(node.value, mem) + else: + rhs = self.visit(node.value) + # self.build(ops.TensorSetOp(rhs, mem, 0.0)) + self.symbol_table.insert(node.targets[0].value.id, rhs) + self.build(ops.TensorPrintOp, rhs) # self.mlir_code.append(store.dump()) # lhs = self.visit(node.targets[0].) # if indices: # store = StoreOp(load.results[0], lhs) - def visit_BinOp(self, node): + def visit_BinOp(self, node, res_val = None): self.pragma = '' + op = None + elewise = True + if isinstance(node.op, ast.Add): + op = self.ADDOP + elif isinstance(node.op, ast.Sub): + op = self.SUBOP + elif isinstance(node.op, ast.Mult): + op = self.MULOP + elif isinstance(node.op, ast.MatMult): + elewise = False + op = self.MATMULOP + else: + raise Exception("Unexpected binary operation") + + if isinstance(node.left, ast.BinOp): + left = self.visit_BinOp(node.left, res_val if elewise else None) + if hasattr(left.defining_op, 'res_val') and left.defining_op.res_val != None: + res_val = left + elif isinstance(node.left, ast.Call): + left = self.visit_Call(node.left, res_val if elewise else None) + if hasattr(left.defining_op, 'res_val') and left.defining_op.res_val != None: + res_val = left + else: + left = self.visit(node.left) - left = self.visit(node.left) + if isinstance(left.type, MemrefType) : + if isinstance(node.left, ast.Name): + left = self.build(ops.ToTensorOp, left).results[0] + self.symbol_table.insert(node.left.id, left) + else: + assert(False) right = self.visit(node.right) + if isinstance(right.type, MemrefType): + if isinstance(node.right, ast.Name): + right = self.build(ops.ToTensorOp, right).results[0] + self.symbol_table.insert(node.right.id, right) + else: + assert(False) val = None res_format = DENSE - if isinstance(node.op, ast.Add): + if op == self.ADDOP: if isinstance(left.type, ShapedType): - indices = [self.build(ops.TensorIndexLabelOp()).results[0] for d in left.type.shape] + indices = [self.build(ops.TensorIndexLabelOp).results[0] for d in left.type.shape] res_format = self.sp_elw_add_sub_conversions[left.type.format][right.type.format] - val = self.build(ops.TensorAddOp(left, right, indices, res_format)) + val = self.build(ops.TensorAddOp, left, right, indices, res_format, res_val) else: - val = self.build(ops.AddOp(left, right)) - elif isinstance(node.op, ast.Sub): + val = self.build(ops.AddOp, left, right) + elif op == self.SUBOP: if isinstance(left.type, ShapedType): - indices = [self.build(ops.TensorIndexLabelOp()).results[0] for d in left.type.shape] + indices = [self.build(ops.TensorIndexLabelOp).results[0] for d in left.type.shape] res_format = self.sp_elw_add_sub_conversions[left.type.format][right.type.format] - val = self.build(ops.TensorSubOp(left, right, indices, res_format)) + val = self.build(ops.TensorSubOp, left, right, indices, res_format, res_val) else: - val = self.build(ops.SubOp(left, right)) - elif isinstance(node.op, ast.Mult): + val = self.build(ops.SubOp, left, right) + elif op == self.MULOP: if isinstance(left.type, ShapedType): - indices = [self.build(ops.TensorIndexLabelOp()).results[0] for d in left.type.shape] + indices = [self.build(ops.TensorIndexLabelOp).results[0] for d in left.type.shape] res_format = self.sp_elw_mult_conversions[left.type.format][right.type.format] - val = self.build(ops.TensorMulOp(left, right, indices, res_format)) + val = self.build(ops.TensorMulOp, left, right, indices, res_format, res_val) else: - val = self.build(ops.MulOp(left, right)) - elif isinstance(node.op, ast.MatMult): + val = self.build(ops.MulOp, left, right) + elif op == self.MATMULOP: assert(isinstance(left.type, ShapedType) and isinstance(right.type, ShapedType)) - lhs_indices = [self.build(ops.TensorIndexLabelOp()).results[0] for d in left.type.shape] - rhs_indices = [lhs_indices[-1]] + [self.build(ops.TensorIndexLabelOp()).results[0] for d in right.type.shape[1:]] + lhs_indices = [self.build(ops.TensorIndexLabelOp).results[0] for d in left.type.shape] + rhs_indices = [lhs_indices[-1]] + [self.build(ops.TensorIndexLabelOp).results[0] for d in right.type.shape[1:]] if len(lhs_indices) ==2 and len(rhs_indices) == 2: res_indices = [lhs_indices[0], rhs_indices[1]] elif len(lhs_indices) == 2 and len(rhs_indices) == 1: @@ -378,7 +486,7 @@ def visit_BinOp(self, node): elif len(lhs_indices) == 1 and len(rhs_indices) == 2: res_indices = [rhs_indices[1]] res_format = self.sp_matmult_conversions[left.type.format][right.type.format] - val = self.build(ops.TensorMatMultOp(left, right, lhs_indices, rhs_indices, res_indices, res_format)) + val = self.build(ops.TensorMatMultOp, left, right, lhs_indices, rhs_indices, res_indices, res_format, res_val) if res_format != DENSE: self.need_opt_comp_workspace = True if val: @@ -390,32 +498,32 @@ def visit_For(self, node): pragma = self.pragma assert(isinstance(node.iter, ast.Call) and node.iter.func.id == 'range') if len(node.iter.args) == 1: - lb = self.build(ops.ConstantOp(0)) + lb = self.build(ops.ConstantOp, 0) lb = lb.results[0] ub = self.visit(node.iter.args[0]) - step = self.build(ops.ConstantOp(1)) + step = self.build(ops.ConstantOp, 1) step = step.results[0] elif len(node.iter.args) == 2: lb = self.visit(node.iter.args[0]) ub = self.visit(node.iter.args[1]) - step = self.build(ops.ConstantOp(1)) + step = self.build(ops.ConstantOp, 1) step = step.results[0] else: lb = self.visit(node.iter.args[0]) ub = self.visit(node.iter.args[1]) step = self.visit(node.iter.args[2]) if pragma == "#pragma parallel": - forOp = self.build(ops.ForAllOp(lb, ub, step)) + forOp = self.build(ops.ForAllOp, lb, ub, step) else: - forOp = self.build(ops.ForOp(lb, ub, step)) + forOp = self.build(ops.ForOp, lb, ub, step) self.symbol_table.insert(node.target.id, forOp.iv) for stmt in node.body: self.visit(stmt) if pragma == "#pragma parallel": - self.build(ops.ReduceOp([])) + self.build(ops.ReduceOp, []) else: - self.build(ops.YieldOp([])) + self.build(ops.YieldOp, []) @@ -423,14 +531,13 @@ def visit_Subscript(self, node): self.pragma = '' mem = self.visit(node.value) - assert(mem != None) indices = self.visit(node.slice) if not isinstance(indices, list): indices = [indices] if isinstance(node.value.ctx, ast.Load): if isinstance(mem.type, TensorType): - mem = self.build(ops.ToMemrefOp(mem)).results[0] - load = self.build(ops.LoadOp(mem, indices)) + mem = self.build(ops.ToMemrefOp, mem).results[0] + load = self.build(ops.LoadOp, mem, indices) return load.results[0] # elif isinstance(node.value.ctx, ast.Store): # store = StoreOp(mem, indices) @@ -442,7 +549,7 @@ def visit_Tuple(self, node): - def visit_Call(self, node): + def visit_Call(self, node, res_val = None): if isinstance(node.func, ast.Attribute): value = self.visit(node.func.value) if isinstance(value, ops.Symbol): @@ -455,27 +562,43 @@ def visit_Call(self, node): attr = self.visit(node.func.id) if attr == 'sum': - res = self.build(ops.TensorSumOp(*values)).results[0] + if isinstance(values[0], MemrefType): + values[0] = self.build(ops.ToTensorOp, values[0]).results[0] + if isinstance(node.func.value, ast.Name): + self.symbol_table.insert(node.func.value.id, values[0]) + res = self.build(ops.TensorSumOp, values[0]).results[0] return res elif attr == 'transpose': - indices = [self.build(ops.TensorIndexLabelOp()).results[0] for d in values[0].type.shape] + indices = [self.build(ops.TensorIndexLabelOp).results[0] for d in values[0].type.shape] transpose_indices = [indices[1], indices[0]] format = None if isinstance(values[0].type, TensorType) or isinstance(values[0].type, MemrefType): format = DENSE elif isinstance(values[0].type, types.TASparseTensorType): format = values[0].type.format - res = self.build(ops.TensorTransposeOp(values[0], indices, transpose_indices, format)).results[0] + if isinstance(values[0], MemrefType): + values[0] = self.build(ops.ToTensorOp, values[0]).results[0] + if isinstance(node.func.value, ast.Name): + self.symbol_table.insert(node.func.value.id, values[0]) + res = self.build(ops.TensorTransposeOp, values[0], indices, transpose_indices, format, res_val).results[0] return res elif attr == 'einsum': - res = lower_einsum_expresion(node.args, node.keywords, self) + res = lower_einsum_expresion(node.args, node.keywords, self, res_val) return res elif attr == 'multiply': - indices = [self.build(ops.TensorIndexLabelOp()).results[0] for d in values[0].type.shape] + indices = [self.build(ops.TensorIndexLabelOp).results[0] for d in values[0].type.shape] for arg in node.args: values.append(self.visit(arg)) res_format = self.sp_elw_mult_conversions[values[0].type.format][values[1].type.format] - res = self.build(ops.TensorMulOp(values[0], values[1], indices, res_format)).results[0] + if isinstance(values[0], MemrefType): + values[0] = self.build(ops.ToTensorOp, values[0]).results[0] + if isinstance(node.func.value, ast.Name): + self.symbol_table.insert(node.func.value.id, values[0]) + if isinstance(values[1], MemrefType): + values[1] = self.build(ops.ToTensorOp, values[1]).results[0] + if isinstance(node.func.value, ast.Name): + self.symbol_table.insert(node.func.value.id, values[1]) + res = self.build(ops.TensorMulOp, values[0], values[1], indices, res_format, res_val).results[0] return res @@ -494,9 +617,9 @@ def visit_Constant(self, node): val = None if isinstance(node.value, int): - val = self.build(ops.ConstantOp(node.value)) + val = self.build(ops.ConstantOp, node.value) elif isinstance(node.value, float): - val = self.build(ops.ConstantOp(node.value)) + val = self.build(ops.ConstantOp, node.value) if val: return val.results[0] @@ -505,7 +628,7 @@ def visit_Return(self, node): # self.build(ops.TensorPrintOp(ret)) if not isinstance(ret, list): ret = [ret] - retOp = self.build(ops.ReturnOp(ret)) + retOp = self.build(ops.ReturnOp, ret) self.return_type = [ret.type for ret in retOp.results] @@ -731,7 +854,6 @@ def visit_Assign(self, node): # We do not support multiple targets currently if isinstance(node.targets[0], ast.Subscript) and isinstance(node.targets[0].ctx, ast.Store) and (isinstance(node.targets[0].slice, ast.Tuple)) : mask = node.targets[0].slice self.mask = mask - print(slice) v = NewVisitor.visit(self, node.value) if self.tsemantics[v]['shape'] != 1: diff --git a/frontends/numpy-scipy/integration_tests/kernels/test_triangleCount_SandiaLL.py b/frontends/numpy-scipy/integration_tests/kernels/test_triangleCount_SandiaLL.py index 7f4218dd..0712ad93 100644 --- a/frontends/numpy-scipy/integration_tests/kernels/test_triangleCount_SandiaLL.py +++ b/frontends/numpy-scipy/integration_tests/kernels/test_triangleCount_SandiaLL.py @@ -13,7 +13,6 @@ def run_comet_with_jit(L0): C = ((L0 @ L0) * L0).sum() return C -@pytest.mark.skip(reason="Triangle counting currently fails") def test_triangleCount_SandiaLL(data_tc_path): A = sp.sparse.csr_array(sp.io.mmread(data_tc_path)) L0 = sp.sparse.csr_array(sp.sparse.tril(A, format='csr')) diff --git a/frontends/numpy-scipy/integration_tests/opts/test_fusion.py b/frontends/numpy-scipy/integration_tests/opts/test_fusion.py index c24009d1..0aab0058 100644 --- a/frontends/numpy-scipy/integration_tests/opts/test_fusion.py +++ b/frontends/numpy-scipy/integration_tests/opts/test_fusion.py @@ -17,13 +17,11 @@ def run_comet_with_jit(B,C,D): return A -@pytest.mark.skip("Fusing more than one iterations is not currently supported") def test_fusion(data_rank2_path): B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) C = np.full([B.shape[1], 4], 1.2, dtype=float) D = np.full([4, 4], 3.4, dtype=float) - A = np.full([B.shape[0], 4], 0.0, dtype=float) - T = np.full([B.shape[0], 4], 0.0, dtype=float) + expected_result = run_numpy(B,C,D) result_with_jit = run_comet_with_jit(B,C,D) if sp.sparse.issparse(expected_result): diff --git a/lib/Dialect/TensorAlgebra/Transforms/LinalgTransforms.cpp b/lib/Dialect/TensorAlgebra/Transforms/LinalgTransforms.cpp index e9853379..aee201b9 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/LinalgTransforms.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/LinalgTransforms.cpp @@ -775,7 +775,6 @@ struct OptDenseTranspose : public ConversionPattern rewriter.replaceAllUsesWith(op->getResult(0), forAll->getResult(0)); rewriter.eraseOp(op); - //module.dump(); return success(); } @@ -796,7 +795,7 @@ namespace comet_debug() << "OptDenseTransposePass : public PassWrapper\n"; func::FuncOp func = getOperation(); ConversionTarget target(getContext()); - target.addLegalDialect(); + target.addLegalDialect(); RewritePatternSet patterns(&getContext()); patterns.insert(&getContext(), tile_size, seperate_tiles); From 8ca914e79bd5bc3fa665c8e6d7447d7bfa73dc56 Mon Sep 17 00:00:00 2001 From: Polykarpos Thomadakis Date: Thu, 24 Apr 2025 15:21:50 -0700 Subject: [PATCH 11/11] [COMETPY] 1. Fixed small incosistenices in output formats 2. Made `flags` argument of `comet.compile`optional 3. Added tests/exampes for new python constructs supported 4. Added "in_place" and "return" versions of several tests that will store results in an input argument or return them normally. --- frontends/numpy-scipy/cometpy/comet.py | 25 ++++------- .../compound_exps/test_Dense_chain_mult.py | 31 +++++++------ .../test_Dense_eltwise_dTranspose.py | 37 +++++++++------- .../test_Dense_mult_dTranspose.py | 36 +++++++++------- .../test_dTranspose_eltwise_CSR.py | 43 +++++++++++++------ .../test_dTranspose_eltwise_Dense.py | 39 +++++++++-------- .../compound_exps/test_dTranspose_mult_CSR.py | 41 ++++++++++++------ .../test_dTranspose_mult_Dense.py | 38 +++++++++------- .../python_constructs/test_jacobi1D.py | 30 +++++++++++++ .../python_constructs/test_jacobi2D.py | 36 ++++++++++++++++ .../python_constructs/test_mixed.py | 33 ++++++++++++++ 11 files changed, 271 insertions(+), 118 deletions(-) create mode 100644 frontends/numpy-scipy/integration_tests/python_constructs/test_jacobi1D.py create mode 100644 frontends/numpy-scipy/integration_tests/python_constructs/test_jacobi2D.py create mode 100644 frontends/numpy-scipy/integration_tests/python_constructs/test_mixed.py diff --git a/frontends/numpy-scipy/cometpy/comet.py b/frontends/numpy-scipy/cometpy/comet.py index 92478109..8a7ce114 100644 --- a/frontends/numpy-scipy/cometpy/comet.py +++ b/frontends/numpy-scipy/cometpy/comet.py @@ -167,7 +167,6 @@ def lower_einsum_expresion(args, kwargs, visitor, res_val = None): res_symbol_indices.append(indice) res_format = visitor.sp_matmult_conversions[lhs.type.format][rhs.type.format] - print('EINSUM', actual_res_val) if semiring: res = visitor.build(ops.TensorMatMultOp, lhs, rhs, lhs_indices, rhs_indices, res_symbol_indices, res_format, actual_res_val, 1.0, 0.0, mask, masktype, semiring).results[0] else: @@ -283,7 +282,7 @@ def __init__(self,inputs): DENSE : COO }, DENSE: { - CSR : DENSE, + CSR : CSR, COO : DENSE, DENSE : DENSE, } @@ -376,8 +375,6 @@ def visit_Assign(self, node): ip = self.get_insertion_point() ip_to_insert = self.insertion_point[0] while mem.block != ip_to_insert.block: - print(mem.block) - print(ip_to_insert) ip_to_insert = ip_to_insert.block self.set_insertion_point(ip_to_insert) memref = self.build(ops.ToMemrefOp, mem) @@ -458,21 +455,21 @@ def visit_BinOp(self, node, res_val = None): if isinstance(left.type, ShapedType): indices = [self.build(ops.TensorIndexLabelOp).results[0] for d in left.type.shape] res_format = self.sp_elw_add_sub_conversions[left.type.format][right.type.format] - val = self.build(ops.TensorAddOp, left, right, indices, res_format, res_val) + val = self.build(ops.TensorAddOp, left, right, indices, res_format, res_val if res_val and res_format == res_val.type.format else None) else: val = self.build(ops.AddOp, left, right) elif op == self.SUBOP: if isinstance(left.type, ShapedType): indices = [self.build(ops.TensorIndexLabelOp).results[0] for d in left.type.shape] res_format = self.sp_elw_add_sub_conversions[left.type.format][right.type.format] - val = self.build(ops.TensorSubOp, left, right, indices, res_format, res_val) + val = self.build(ops.TensorSubOp, left, right, indices, res_format, res_val if res_val and res_format == res_val.type.format else None) else: val = self.build(ops.SubOp, left, right) elif op == self.MULOP: if isinstance(left.type, ShapedType): indices = [self.build(ops.TensorIndexLabelOp).results[0] for d in left.type.shape] res_format = self.sp_elw_mult_conversions[left.type.format][right.type.format] - val = self.build(ops.TensorMulOp, left, right, indices, res_format, res_val) + val = self.build(ops.TensorMulOp, left, right, indices, res_format, res_val if res_val and res_format == res_val.type.format else None) else: val = self.build(ops.MulOp, left, right) elif op == self.MATMULOP: @@ -486,7 +483,7 @@ def visit_BinOp(self, node, res_val = None): elif len(lhs_indices) == 1 and len(rhs_indices) == 2: res_indices = [rhs_indices[1]] res_format = self.sp_matmult_conversions[left.type.format][right.type.format] - val = self.build(ops.TensorMatMultOp, left, right, lhs_indices, rhs_indices, res_indices, res_format, res_val) + val = self.build(ops.TensorMatMultOp, left, right, lhs_indices, rhs_indices, res_indices, res_format, res_val if res_val and res_format == res_val.type.format else None) if res_format != DENSE: self.need_opt_comp_workspace = True if val: @@ -571,16 +568,12 @@ def visit_Call(self, node, res_val = None): elif attr == 'transpose': indices = [self.build(ops.TensorIndexLabelOp).results[0] for d in values[0].type.shape] transpose_indices = [indices[1], indices[0]] - format = None - if isinstance(values[0].type, TensorType) or isinstance(values[0].type, MemrefType): - format = DENSE - elif isinstance(values[0].type, types.TASparseTensorType): - format = values[0].type.format + res_format = self.sp_mattr_conversions[values[0].type.format] if isinstance(values[0], MemrefType): values[0] = self.build(ops.ToTensorOp, values[0]).results[0] if isinstance(node.func.value, ast.Name): self.symbol_table.insert(node.func.value.id, values[0]) - res = self.build(ops.TensorTransposeOp, values[0], indices, transpose_indices, format, res_val).results[0] + res = self.build(ops.TensorTransposeOp, values[0], indices, transpose_indices, res_format, res_val if res_val and res_format == res_val.type.format else None).results[0] return res elif attr == 'einsum': res = lower_einsum_expresion(node.args, node.keywords, self, res_val) @@ -735,7 +728,7 @@ def __init__(self,inputs): DENSE : COO }, DENSE: { - CSR : DENSE, + CSR : CSR, COO : DENSE, DENSE : DENSE, } @@ -1442,7 +1435,7 @@ def visit_Einsum_Call(self, node: ast.Call): #Wrapper function. The input function (in the form of an object) is passed an arguement to this function. -def compile(flags, target:str = "cpu", with_jit=True): +def compile(flags=None, target:str = "cpu", with_jit=True): def innerfunc(func): def wrapper(*pos_args, **kwargs): diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_chain_mult.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_chain_mult.py index 52b6dcbe..08b0e7f7 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_chain_mult.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_chain_mult.py @@ -2,26 +2,33 @@ import numpy as np @comet.compile(flags="") -def comet_chain_multiply(A,B,C,D): - # res = comet.einsum('ij,jk,kl,lm->im', A, C, B, C.transpose()) @ D - res = A @ C @ B @ C.transpose() @ D +def comet_chain_multiply_in_place(A,B,C,D,E): + E[:] = A @ C @ B @ C.transpose() @ D - return res +@comet.compile(flags="") +def comet_chain_multiply_return(A,B,C,D): + return A @ C @ B @ C.transpose() @ D def numpy_chain_multiply(A,B,C,D): - # res = np.einsum('ij,jk,kl,lm->im', A, C, B, C.transpose()) @ D - res = A @ C @ B @ C.transpose() @ D - + return A @ C @ B @ C.transpose() @ D - return res -def test_Dense_chain_mult(): +def init(): A = np.full([2,4], 1.0) B = np.full([4,4], 1.0) C = np.full([4,4], 1.0) D = np.full([4,6], 1.0) + Enp = numpy_chain_multiply(A,B,C,D) + + return A, B, C, D, Enp - Dn = numpy_chain_multiply(A,B,C,D) - Dc = comet_chain_multiply(A,B,C,D) +def test_Dense_chain_mult_in_place(): + A, B, C, D, Enp = init() + Ecp = np.full([2,6], 0.0) + comet_chain_multiply_in_place(A,B,C,D,Ecp) + np.testing.assert_almost_equal(Ecp,Enp) - np.testing.assert_almost_equal(Dn,Dc) \ No newline at end of file +def test_Dense_chain_mult_return(): + A, B, C, D, Enp = init() + Ecp = comet_chain_multiply_return(A,B,C,D) + np.testing.assert_almost_equal(Ecp,Enp) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_eltwise_dTranspose.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_eltwise_dTranspose.py index 56ec0fda..a2eecbb7 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_eltwise_dTranspose.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_eltwise_dTranspose.py @@ -1,26 +1,31 @@ -import time import numpy as np -import scipy as sp from cometpy import comet def run_numpy(A,B): - C =B * A.transpose() + return B * A.transpose() - return C +@comet.compile() +def run_comet_return(A,B): + return B * A.transpose() -@comet.compile(flags=None) -def run_comet_with_jit(A,B): - C = B * A.transpose() +@comet.compile() +def run_comet_in_place(A,B,C): + C[:] = B * A.transpose() - return C -def test_Dense_eltwise_dTranspose(): +def init(): A = np.full([4, 4], 3.2, dtype=float) B = np.full([4, 4], 2.0, dtype=float) - C = np.full([4, 4], 0.0, dtype=float) - expected_result = run_numpy(A, B) - result_with_jit = run_comet_with_jit(A, B) - if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() - np.testing.assert_almost_equal(result_with_jit, expected_result) + Cnp = run_numpy(A, B) + return A, B, Cnp + +def test_Dense_eltwise_dTranspose_in_place(): + A, B, Cnp = init() + Ccp = np.full([4, 4], 0.0, dtype=float) + run_comet_in_place(A, B, Ccp) + np.testing.assert_almost_equal(Cnp, Ccp) + +def test_Dense_eltwise_dTranspose_return(): + A, B, Cnp = init() + Ccp = run_comet_return(A, B) + np.testing.assert_almost_equal(Cnp, Ccp) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_mult_dTranspose.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_mult_dTranspose.py index 75833abc..c5ee2f33 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_mult_dTranspose.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_Dense_mult_dTranspose.py @@ -1,6 +1,4 @@ -import time import numpy as np -import scipy as sp from cometpy import comet def run_numpy(A,B): @@ -8,19 +6,27 @@ def run_numpy(A,B): return C -@comet.compile(flags=None) -def run_comet_with_jit(A,B): - C = A @ B.transpose() +@comet.compile() +def run_comet_return(A,B): + return A @ B.transpose() - return C +@comet.compile() +def run_comet_in_place(A,B,C): + C[:] = A @ B.transpose() -def test_Dense_mult_dTranspose(): - B = np.full([5, 5], 3.2, dtype=float) +def init(): A = np.full([5, 5], 2.3, dtype=float) - C = np.full([5, 5], 0.0, dtype=float) - expected_result = run_numpy(A,B) - result_with_jit = run_comet_with_jit(A,B) - if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() - np.testing.assert_almost_equal(result_with_jit, expected_result) + B = np.full([5, 5], 3.2, dtype=float) + Cnp = run_numpy(A, B) + return A, B, Cnp + +def test_Dense_mult_dTranspose_return(): + A, B, Cnp = init() + Ccp = run_comet_return(A,B) + np.testing.assert_almost_equal(Ccp, Cnp) + +def test_Dense_mult_dTranspose_in_place(): + A, B, Cnp = init() + Ccp = np.full([5, 5], 0.0, dtype=float) + run_comet_in_place(A,B, Ccp) + np.testing.assert_almost_equal(Ccp, Cnp) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_CSR.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_CSR.py index 1db73d56..2c8eaf4e 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_CSR.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_CSR.py @@ -1,26 +1,43 @@ -import time import numpy as np import scipy as sp from cometpy import comet +import pytest def run_numpy(A, B): C = A.transpose() * B return C -@comet.compile(flags=None) -def run_comet_with_jit(A, B): - C = A.transpose() * B +@comet.compile() +def run_comet_return(A, B): + return A.transpose() * B - return C +@comet.compile() +def run_comet_in_place(A, B, C): + C[:] = A.transpose() * B -def test_dTranspose_eltwise_CSR(data_rank2_path): +def init(data_rank2_path): B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) A = np.full([B.shape[1], B.shape[0]], 3.2, dtype=float) - expected_result = run_numpy(A, B) - result_with_jit = run_comet_with_jit(A, B) - if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - if sp.sparse.issparse(result_with_jit): - result_with_jit = result_with_jit.todense() - np.testing.assert_almost_equal(result_with_jit, expected_result) + Cnp = run_numpy(A, B) + return A, B, Cnp + +def test_dTranspose_eltwise_CSR_return(data_rank2_path): + A, B, Cnp = init(data_rank2_path) + Ccp = run_comet_return(A, B) + if sp.sparse.issparse(Cnp): + Cnp = Cnp.todense() + if sp.sparse.issparse(Ccp): + Ccp = Ccp.todense() + np.testing.assert_almost_equal(Ccp, Cnp) + +@pytest.mark.skip('In place mutation of sparse matrices is not supported yet') +def test_dTranspose_eltwise_CSR_in_place(data_rank2_path): + A, B, Cnp = init(data_rank2_path) + Ccp = sp.sparse.csr_array((B.shape[0], B.shape[1]), dtype = float) + run_comet_in_place(A, B, Ccp) + if sp.sparse.issparse(Cnp): + Cnp = Cnp.todense() + if sp.sparse.issparse(Ccp): + Ccp = Ccp.todense() + np.testing.assert_almost_equal(Ccp, Cnp) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_Dense.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_Dense.py index 3319257b..3f218e0a 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_Dense.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_eltwise_Dense.py @@ -1,26 +1,31 @@ -import time import numpy as np -import scipy as sp from cometpy import comet -def run_numpy(A, B): - C = A.transpose() * B +def run_numpy(A,B): + return A.transpose() * B - return C +@comet.compile(flags=None) +def run_comet_return(A,B,C): + return A.transpose() * B @comet.compile(flags=None) -def run_comet_with_jit(A, B): - C = A.transpose() * B +def run_comet_in_place(A,B,C): + C[:] = A.transpose() * B - return C -def test_dTranspose_eltwise_Dense(): +def init(): A = np.full([4, 4], 3.2, dtype=float) - B = np.full([4, 4], 2.0, dtype=float) - C = np.full([4, 4], 0.0, dtype=float) - expected_result = run_numpy(A, B) - result_with_jit = run_comet_with_jit(A, B) - if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() - np.testing.assert_almost_equal(result_with_jit, expected_result) + B = np.full([4, 4], 2.0, dtype=float) + Cnp = run_numpy(A, B) + return A, B, Cnp + +def test_dTranspose_eltwise_Dense_in_place(): + A, B, Cnp = init() + Ccp = np.full([4, 4], 0.0, dtype=float) + run_comet_in_place(A, B, Ccp) + np.testing.assert_almost_equal(Ccp, Cnp) + +def test_dTranspose_eltwise_Dense_return(): + A, B, Cnp = init() + Ccp = run_comet_return(A, B) + np.testing.assert_almost_equal(Ccp, Cnp) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_CSR.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_CSR.py index 37aa54c5..53feb82c 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_CSR.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_CSR.py @@ -1,4 +1,3 @@ -import time import numpy as np import scipy as sp from cometpy import comet @@ -8,19 +7,35 @@ def run_numpy(A, B): return C -@comet.compile(flags=None) -def run_comet_with_jit(A, B): - C = A.transpose() @ B +@comet.compile() +def run_comet_return(A, B): + return A.transpose() @ B - return C +@comet.compile() +def run_comet_in_place(A, B, C): + C[:] = A.transpose() @ B -def test_dTranspose_mult_CSR(data_rank2_path): +def init(data_rank2_path): B = sp.sparse.csr_array(sp.io.mmread(data_rank2_path)) A = np.full([B.shape[0], 4], 1.7, dtype=float) - C = np.full([4, B.shape[1]], 0.0, dtype=float) - expected_result = run_numpy(A, B) - result_with_jit = run_comet_with_jit(A, B) - if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() - np.testing.assert_almost_equal(result_with_jit, expected_result) + Cnp = run_numpy(A, B) + return A, B, Cnp + + + +def test_dTranspose_mult_CSR_in_place(data_rank2_path): + A, B, Cnp = init(data_rank2_path) + Ccp = np.full([4, B.shape[1]], 0.0, dtype=float) + run_comet_in_place(A, B, Ccp) + if sp.sparse.issparse(Cnp): + Cnp = Cnp.todense() + Ccp = Ccp.todense() + np.testing.assert_almost_equal(Ccp, Cnp) + +def test_dTranspose_mult_CSR_return(data_rank2_path): + A, B, Cnp = init(data_rank2_path) + Ccp = run_comet_return(A, B) + if sp.sparse.issparse(Cnp): + Cnp = Cnp.todense() + Ccp = Ccp.todense() + np.testing.assert_almost_equal(Ccp, Cnp) diff --git a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_Dense.py b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_Dense.py index 67aeaa22..e47649d5 100644 --- a/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_Dense.py +++ b/frontends/numpy-scipy/integration_tests/compound_exps/test_dTranspose_mult_Dense.py @@ -1,26 +1,32 @@ -import time import numpy as np import scipy as sp from cometpy import comet -def run_numpy(A, B): - C = A.transpose() @ B +def run_numpy(A,B): + return A.transpose() @ B - return C +@comet.compile() +def run_comet_return(A,B): + return A.transpose() @ B -@comet.compile(flags=None) -def run_comet_with_jit(A, B): - C = A.transpose() @ B +@comet.compile() +def run_comet_in_place(A,B,C): + C[:] = A.transpose() @ B - return C -def test_dTranspose_mult_Dense(): +def init(): A = np.full([4, 4], 3.2, dtype=float) B = np.full([4, 4], 1.0, dtype=float) - C = np.full([4, 4], 0.0, dtype=float) - expected_result = run_numpy(A, B) - result_with_jit = run_comet_with_jit(A, B) - if sp.sparse.issparse(expected_result): - expected_result = expected_result.todense() - result_with_jit = result_with_jit.todense() - np.testing.assert_almost_equal(result_with_jit, expected_result) + Cnp = run_numpy(A, B) + return A, B, Cnp + +def test_dTranspose_mult_Dense_in_place(): + A, B, Cnp = init() + Ccp = np.full([4, 4], 0.0, dtype=float) + run_comet_in_place(A, B, Ccp) + np.testing.assert_almost_equal(Ccp, Cnp) + +def test_dTranspose_mult_Dense_return(): + A, B, Cnp = init() + Ccp = run_comet_return(A, B) + np.testing.assert_almost_equal(Ccp, Cnp) diff --git a/frontends/numpy-scipy/integration_tests/python_constructs/test_jacobi1D.py b/frontends/numpy-scipy/integration_tests/python_constructs/test_jacobi1D.py new file mode 100644 index 00000000..e248b60e --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/python_constructs/test_jacobi1D.py @@ -0,0 +1,30 @@ +from cometpy import comet +import numpy as np + + +def jacobi1D_numpy(A, B, N): + for i in range(1, N-1): + B[i] = 0.33333 * (A[i-1] + A[i] + A[i + 1]) + for i in range(1, N-1): + A[i] = 0.33333 * (B[i-1] + B[i] + B[i + 1]); + +@comet.compile() +def jacobi1D_comet(A, B, N): + + for i in range(1, N-1): + B[i] = 0.33333 * (A[i-1] + A[i] + A[i + 1]) + for i in range(1, N-1): + A[i] = 0.33333 * (B[i-1] + B[i] + B[i + 1]); + + +def test_jacobi1D(): + Anp = np.full([8], 1.0, dtype=float) + Bnp = np.full([8], 1.0, dtype=float) + jacobi1D_numpy(Anp, Bnp, Anp.shape[0]) + Acp = np.full([8], 1.0, dtype=float) + Bcp = np.full([8], 1.0, dtype=float) + jacobi1D_comet(Acp, Bcp, Acp.shape[0]) + + np.testing.assert_almost_equal(Acp, Anp) + np.testing.assert_almost_equal(Bcp, Bnp) + diff --git a/frontends/numpy-scipy/integration_tests/python_constructs/test_jacobi2D.py b/frontends/numpy-scipy/integration_tests/python_constructs/test_jacobi2D.py new file mode 100644 index 00000000..df3b5e7a --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/python_constructs/test_jacobi2D.py @@ -0,0 +1,36 @@ +from cometpy import comet +import numpy as np + +def jacobi2D_numpy(A, B, N): + for i in range(1, N-1): + for j in range(1, N-1): + B[i,j] = 0.2 * (A[i, j] + A[i, j-1] + A[i, j + 1] + A[i+1, j] + A[i-1, j]) + + for i in range(1, N-1): + for j in range(1, N-1): + A[i, j] = 0.2 * (B[i, j] + B[i, j-1] + B[i, j + 1] + B[i+1, j] + B[i - 1, j]); + +@comet.compile() +def jacobi2D_comet(A, B, N): + #pragma parallel + for i in range(1, N-1): + #pragma parallel + for j in range(1, N-1): + B[i,j] = 0.2 * (A[i, j] + A[i, j-1] + A[i, j + 1] + A[i+1, j] + A[i-1, j]) + + #pragma parallel + for i in range(1, N-1): + #pragma parallel + for j in range(1, N-1): + A[i, j] = 0.2 * (B[i, j] + B[i, j-1] + B[i, j + 1] + B[i+1, j] + B[i - 1, j]); + +def test_jacobi2D(): + Anp = np.full([8, 8], 1.0, dtype=float) + Bnp = np.full([8, 8], 1.0, dtype=float) + jacobi2D_numpy(Anp, Bnp, Anp.shape[0]) + Acp = np.full([8, 8], 1.0, dtype=float) + Bcp = np.full([8, 8], 1.0, dtype=float) + jacobi2D_comet(Acp, Bcp, Acp.shape[0]) + + np.testing.assert_almost_equal(Acp, Anp) + np.testing.assert_almost_equal(Bcp, Bnp) \ No newline at end of file diff --git a/frontends/numpy-scipy/integration_tests/python_constructs/test_mixed.py b/frontends/numpy-scipy/integration_tests/python_constructs/test_mixed.py new file mode 100644 index 00000000..7d480d47 --- /dev/null +++ b/frontends/numpy-scipy/integration_tests/python_constructs/test_mixed.py @@ -0,0 +1,33 @@ +from cometpy import comet +import numpy as np + + +def mixed_numpy(A, B, N): + for i in range(1, N-1): + B[i] = 0.33333 * (A[i-1] + A[i] + A[i + 1]) + for i in range(1, N-1): + A[i] = 0.33333 * (B[i-1] + B[i] + B[i + 1]); + + return A * B + +@comet.compile() +def mixed_comet(A, B, N): + + for i in range(1, N-1): + B[i] = 0.33333 * (A[i-1] + A[i] + A[i + 1]) + for i in range(1, N-1): + A[i] = 0.33333 * (B[i-1] + B[i] + B[i + 1]); + + return A * B + +def test_mixed(): + Anp = np.full([8], 1.0, dtype=float) + Bnp = np.full([8], 1.0, dtype=float) + mixed_numpy(Anp, Bnp, Anp.shape[0]) + Acp = np.full([8], 1.0, dtype=float) + Bcp = np.full([8], 1.0, dtype=float) + mixed_comet(Acp, Bcp, Acp.shape[0]) + + np.testing.assert_almost_equal(Acp, Anp) + np.testing.assert_almost_equal(Bcp, Bnp) +