diff --git a/frontends/comet_dsl/comet.cpp b/frontends/comet_dsl/comet.cpp index 7ad5da23..e45694b9 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) @@ -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 51991669..96740875 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" @@ -48,13 +49,17 @@ #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/Casting.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 @@ -80,7 +85,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 *********// @@ -159,8 +164,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, @@ -346,7 +352,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 @@ -439,10 +445,35 @@ 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); + 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); + mlir::Value rhs = mlirGen(*binop.getRHS(), lhs_lbls, "", nullptr); if (!rhs) return nullptr; auto location = loc(binop.loc()); @@ -479,44 +510,68 @@ 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); - /// 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())) @@ -857,24 +912,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 { @@ -942,24 +982,32 @@ 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; 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 '+': @@ -967,14 +1015,16 @@ 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, + out_val, 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, + out_val, MaskingAttr); case '*': { @@ -986,7 +1036,8 @@ 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, + 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)); @@ -1003,8 +1054,9 @@ 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, + out_val, MaskingAttr); comet_vdump(tcop); @@ -1231,15 +1283,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 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); + // lhs.getDefiningOp()->setAttr("__beta__", builder.getF64FloatAttr(tens_beta)); return rhs; } @@ -1253,16 +1307,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'"); @@ -1273,13 +1319,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: @@ -1287,7 +1333,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: @@ -1471,7 +1517,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)); @@ -1511,15 +1568,15 @@ 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); + 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"; @@ -1565,7 +1622,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) { @@ -1573,6 +1630,7 @@ namespace { if (isa(var.getDefiningOp())) { + rhs_labels_val.push_back(var); all_labels_val.push_back(var); } else @@ -1632,37 +1690,40 @@ 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()); + 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, all_labels_val, affineMapArrayAttr); - builder.create(loc(transpose.loc()), t, lhs_tensor); + rhs_tensor, rhs_labels_val, lhs_labels_val, affineMapArrayAttr, val); comet_vdump(t); return t; @@ -1940,8 +2001,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... @@ -1966,8 +2026,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 +2168,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(); @@ -2214,7 +2266,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) { @@ -2222,10 +2274,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) @@ -2321,14 +2375,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"; + // } } } @@ -2349,14 +2403,14 @@ 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, + tensors[2], 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 == '-') { @@ -2364,38 +2418,37 @@ 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, + tensors[2], 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) { 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, + tensors[2], MaskingAttr, maskVal); 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) { - 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, tensors[2], MaskingAttr); 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 { @@ -2484,7 +2537,7 @@ 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); return mlir::success(); } 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..8a7ce114 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]) @@ -145,9 +168,9 @@ def lower_einsum_expresion(args, kwargs, visitor): res_format = visitor.sp_matmult_conversions[lhs.type.format][rhs.type.format] 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 +182,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 = { @@ -255,7 +282,7 @@ def __init__(self,inputs): DENSE : COO }, DENSE: { - CSR : DENSE, + CSR : CSR, COO : DENSE, DENSE : DENSE, } @@ -271,34 +298,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 +351,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 +360,122 @@ 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: + 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 if res_val and res_format == res_val.type.format else None) 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 if res_val and res_format == res_val.type.format else None) 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 if res_val and res_format == res_val.type.format else None) 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 +483,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_val and res_format == res_val.type.format else None) if res_format != DENSE: self.need_opt_comp_workspace = True if val: @@ -390,32 +495,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 +528,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 +546,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 +559,39 @@ 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] + 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, 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 = 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 +610,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 +621,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] @@ -612,7 +728,7 @@ def __init__(self,inputs): DENSE : COO }, DENSE: { - CSR : DENSE, + CSR : CSR, COO : DENSE, DENSE : DENSE, } @@ -731,7 +847,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: @@ -1320,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/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/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) + diff --git a/include/comet/Dialect/TensorAlgebra/IR/TAOps.td b/include/comet/Dialect/TensorAlgebra/IR/TAOps.td index 148dc97c..dfe74bba 100644 --- a/include/comet/Dialect/TensorAlgebra/IR/TAOps.td +++ b/include/comet/Dialect/TensorAlgebra/IR/TAOps.td @@ -351,88 +351,19 @@ 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, + Optional:$lhs, OptionalAttr:$MaskType, /// TODO: try to use: DefaultValuedAttr Optional:$mask); /// 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 +371,15 @@ 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, + Optional:$lhs, 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,80 +402,14 @@ 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]> { +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, Optional:$lhs); 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]> { @@ -630,7 +427,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. @@ -639,79 +436,18 @@ 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, + Optional:$lhs, 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. @@ -720,75 +456,14 @@ 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, + Optional:$lhs, 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/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 1fe4086d..4ee45565 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; @@ -125,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(); @@ -235,9 +192,55 @@ 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 = mult_op.getLhs(); + + 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(!shapeT.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; + } + } + assert(dims.size() == shapeT.getNumDynamicDims()); + + if(!lhs_tensor) + { + + 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)); + } + } comet_vdump(rhs1_tensor); comet_vdump(rhs2_tensor); @@ -253,7 +256,6 @@ mlir::LogicalResult generalIndexOperationRewrite( } } - auto indexing_maps = mult_op.getIndexingMaps(); auto semiring = cast(mult_op.getSemiringAttr()).getValue(); auto tensor_type = op->getResultTypes()[0]; @@ -261,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); @@ -482,8 +492,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/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 a75eea87..1ac8c2b2 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 @@ -222,26 +195,25 @@ namespace /// There are tensors for transpose operation: input and output tensors unsigned int tensors_num = 2; - tensorAlgebra::TensorSetOp setOp; - Value lhs; + Value lhs = op.getLhs(); - 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(tensorT.isDynamicDim(perm)) { - setOp = cast(u); - Value dstTensor = u->getOperand(1); - lhs = dstTensor; + auto dim = rewriter.create(loc, inputTensor, perm); + dims.push_back(dim); } } + if(!lhs) + { + lhs = rewriter.create(loc, op.getResult().getType(), dims); + } comet_vdump(lhs); // auto outputMemref = lhs.getDefiningOp()->getOperand(0); @@ -279,16 +251,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 +361,6 @@ namespace rewriter.create(loc, func_name, SmallVector{}, ValueRange(allInputs) ); - rewriter.eraseOp(setOp); rewriter.eraseOp(op); return success(); @@ -551,32 +512,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"; @@ -588,22 +524,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 { @@ -616,26 +596,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 @@ -680,11 +640,11 @@ void LowerTensorAlgebraToSCFPass::runOnOperation() bufferization::BufferizationDialect>(); target.addLegalDialect(); + target.addLegalOp(); 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. @@ -692,8 +652,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/IndexTree/Transforms/KernelFusion.cpp b/lib/Dialect/IndexTree/Transforms/KernelFusion.cpp index d9627067..d2f8c293 100644 --- a/lib/Dialect/IndexTree/Transforms/KernelFusion.cpp +++ b/lib/Dialect/IndexTree/Transforms/KernelFusion.cpp @@ -75,7 +75,7 @@ Value getRealLhsTensor(IndexTreeOp itree, block_args.push_back(arg); } llvm::SmallVector 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) { @@ -833,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. @@ -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)); } } 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/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/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/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); diff --git a/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp b/lib/Dialect/TensorAlgebra/Transforms/Passes.cpp index f7ecdc3c..9dfe3d86 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,13 @@ 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; + auto out_val = op.getLhs(); std::vector MultOpsToRemove; std::vector LTOpsToRemove; @@ -240,11 +243,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()) @@ -254,7 +257,7 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::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++) @@ -505,11 +508,11 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::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"; @@ -518,7 +521,7 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::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); @@ -528,7 +531,7 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::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"; @@ -578,35 +581,40 @@ void FindOptimalTCFactorizationPass::FindOptimalTCFactorization(tensorAlgebra::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, + nullptr, 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 + cast(newRhs1.getDefiningOp()).getLhsMutable().assign(out_val); + 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 01adbe4f..9b1e8930 100644 --- a/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp +++ b/lib/Dialect/TensorAlgebra/Transforms/TCtoTTGTDyn.cpp @@ -189,29 +189,50 @@ 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 = multop.getLhs(); 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, rewriter.getZeroAttr( + shapeT.getElementType())); + if(!lhsTensor) + { + lhsTensor = rewriter.create(loc, shapeT, zero, ValueRange(dims)); + } + + auto lhsTensorType = cast(lhsTensor.getType()); std::vector allShapes{rhs1Tensor, rhs2Tensor, @@ -710,8 +731,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 +741,7 @@ namespace // }); // op->replaceAllUsesWith(switchOp); // rewriter.eraseOp(setnewop); - rewriter.eraseOp(op); + // rewriter.eraseOp(op); return success(); } 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, diff --git a/lib/Dialect/Utils/Utils.cpp b/lib/Dialect/Utils/Utils.cpp index f2a27110..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, - 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)