Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Implement OpModel interface for TTNN.MultiplyOp #2349

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+392 −111
Open

Implement OpModel interface for TTNN.MultiplyOp #2349

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion include/ttmlir/Dialect/TTNN/IR/TTNNOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -573,7 +573,10 @@ def TTNN_MinimumOp : TTNN_ElementwiseBinaryOp<"minimum"> {
}];
}

def TTNN_MultiplyOp : TTNN_ElementwiseBinaryOp<"multiply"> {
def TTNN_MultiplyOp : TTNN_ElementwiseBinaryOp<"multiply",
[DeclareOpInterfaceMethods<TTNN_OpModelInterface, ["getOpConstraints", "getOpRuntime"]>]
> {

let summary = "Eltwise multiply.";
let description = [{
Eltwise multiply operation.
Expand Down
23 changes: 23 additions & 0 deletions include/ttmlir/OpModel/TTNN/TTNNOpModel.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -143,5 +143,28 @@ llvm::Expected<size_t> getOpRuntime(llvm::ArrayRef<int64_t> inputShapeA,
bool transposeA, bool transposeB);
}; // namespace MatmulOpInterface

//===----------------------------------------------------------------------===//
// MultiplyOp
//===----------------------------------------------------------------------===//

namespace MultiplyOpInterface {
llvm::Expected<std::tuple<size_t, size_t, size_t>>
getOpConstraints(llvm::ArrayRef<int64_t> inputShapeA,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutA,
llvm::ArrayRef<int64_t> inputShapeB,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutB,
llvm::ArrayRef<int64_t> outputShape,
mlir::tt::ttnn::TTNNLayoutAttr outputLayout);

llvm::Expected<size_t>
getOpRuntime(llvm::ArrayRef<int64_t> inputShapeA,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutA,
llvm::ArrayRef<int64_t> inputShapeB,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutB,
llvm::ArrayRef<int64_t> outputShape,
mlir::tt::ttnn::TTNNLayoutAttr outputLayout);

}; // namespace MultiplyOpInterface

} // namespace mlir::tt::op_model::ttnn
#endif // TTMLIR_OPMODEL_TTNN_TTNNOPMODEL_H
43 changes: 43 additions & 0 deletions lib/Dialect/TTNN/IR/TTNNOpModelInterface.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -283,4 +283,47 @@ MatmulOp::getOpRuntime(const std::vector<TTNNLayoutAttr> &inputs,
false, false);
}

//===----------------------------------------------------------------------===//
// MultiplyOp - TTNN Op Model Interface
//===----------------------------------------------------------------------===//

llvm::Expected<std::tuple<size_t, size_t, size_t>>
MultiplyOp::getOpConstraints(const std::vector<TTNNLayoutAttr> &inputs,
const TTNNLayoutAttr &output) {
assert(inputs.size() == 2);

const auto inputShapeA =
mlir::cast<RankedTensorType>(getOperand(0).getType()).getShape();
const auto inputShapeB =
mlir::cast<RankedTensorType>(getOperand(1).getType()).getShape();

const auto outputShape =
mlir::cast<RankedTensorType>(getResult(0).getType()).getShape();

llvm::Expected<bool> check = detail::checkDeviceWorkerGrid(getOperation());
if (!check) {
return check.takeError();
}

return op_model::ttnn::MultiplyOpInterface::getOpConstraints(
inputShapeA, inputs[0], inputShapeB, inputs[1], outputShape, output);
}

llvm::Expected<size_t>
MultiplyOp::getOpRuntime(const std::vector<TTNNLayoutAttr> &inputs,
const TTNNLayoutAttr &output) {
assert(inputs.size() == 2);

const auto inputShapeA =
mlir::cast<RankedTensorType>(getOperand(0).getType()).getShape();
const auto inputShapeB =
mlir::cast<RankedTensorType>(getOperand(1).getType()).getShape();

const auto outputShape =
mlir::cast<RankedTensorType>(getResult(0).getType()).getShape();

return op_model::ttnn::MultiplyOpInterface::getOpRuntime(
inputShapeA, inputs[0], inputShapeB, inputs[1], outputShape, output);
}

} // namespace mlir::tt::ttnn
157 changes: 107 additions & 50 deletions lib/OpModel/TTNN/TTNNOpModelLib.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,74 @@ Device::getDeviceConstraints(mlir::tt::GridAttr workerGrid) {
return true;
}

//===----------------------------------------------------------------------===//
// Template functions for binary elementwise operations.
//===----------------------------------------------------------------------===//

template <typename OpSymbol>
llvm::Expected<std::tuple<size_t, size_t, size_t>>
getEltwiseBinaryOpConstraints(std::string_view opName, OpSymbol opSymbol,
llvm::ArrayRef<int64_t> inputShapeA,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutA,
llvm::ArrayRef<int64_t> inputShapeB,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutB,
llvm::ArrayRef<int64_t> outputShape,
mlir::tt::ttnn::TTNNLayoutAttr outputLayout) {
auto query = [&](llvm::ArrayRef<int64_t> aShape,
mlir::tt::ttnn::TTNNLayoutAttr aLayout,
llvm::ArrayRef<int64_t> bShape,
mlir::tt::ttnn::TTNNLayoutAttr bLayout,
llvm::ArrayRef<int64_t> outShape,
mlir::tt::ttnn::TTNNLayoutAttr outLayout) {
::tt::tt_metal::v0::IDevice *device =
SingletonDeviceContext::getInstance().getDevice();
const auto [inputSpecA, inputSpecB, outputSpec] =
detail::convertToTensorSpec(device, std::make_tuple(aShape, aLayout),
std::make_tuple(bShape, bLayout),
std::make_tuple(outShape, outLayout));

return ::ttnn::graph::query_op_constraints(
opSymbol, device, inputSpecA, inputSpecB, outputSpec.data_type(),
outputSpec.tensor_layout().get_memory_config());
};

return operation::getOpConstraints(opName, query, inputShapeA, inputLayoutA,
inputShapeB, inputLayoutB, outputShape,
outputLayout);
}

template <typename OpSymbol>
llvm::Expected<size_t>
getEltwiseBinaryOpRuntime(std::string_view opName, OpSymbol opSymbol,
llvm::ArrayRef<int64_t> inputShapeA,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutA,
llvm::ArrayRef<int64_t> inputShapeB,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutB,
llvm::ArrayRef<int64_t> outputShape,
mlir::tt::ttnn::TTNNLayoutAttr outputLayout) {
auto query = [&](llvm::ArrayRef<int64_t> aShape,
mlir::tt::ttnn::TTNNLayoutAttr aLayout,
llvm::ArrayRef<int64_t> bShape,
mlir::tt::ttnn::TTNNLayoutAttr bLayout,
llvm::ArrayRef<int64_t> outShape,
mlir::tt::ttnn::TTNNLayoutAttr outLayout) {
::tt::tt_metal::v0::IDevice *device =
SingletonDeviceContext::getInstance().getDevice();
const auto [inputSpecA, inputSpecB, outputSpec] =
detail::convertToTensorSpec(device, std::make_tuple(aShape, aLayout),
std::make_tuple(bShape, bLayout),
std::make_tuple(outShape, outLayout));

return ::ttnn::graph::query_op_runtime(
opSymbol, device, inputSpecA, inputSpecB, outputSpec.data_type(),
outputSpec.tensor_layout().get_memory_config());
};

return operation::getOpRuntime(opName, query, inputShapeA, inputLayoutA,
inputShapeB, inputLayoutB, outputShape,
outputLayout);
}

//===----------------------------------------------------------------------===//
// ReluOp
//===----------------------------------------------------------------------===//
Expand Down Expand Up @@ -272,31 +340,9 @@ AddOpInterface::getOpConstraints(llvm::ArrayRef<int64_t> inputShapeA,
llvm::ArrayRef<int64_t> outputShape,
mlir::tt::ttnn::TTNNLayoutAttr outputLayout) {
#ifdef TTMLIR_ENABLE_OPMODEL
auto addOpQuery = [](llvm::ArrayRef<int64_t> inputShapeA,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutA,
llvm::ArrayRef<int64_t> inputShapeB,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutB,
llvm::ArrayRef<int64_t> outputShape,
mlir::tt::ttnn::TTNNLayoutAttr outputLayout) {
// open device device, will close it at the end of function
::tt::tt_metal::v0::IDevice *device =
SingletonDeviceContext::getInstance().getDevice();

// prepare io specs
const auto [inputSpecA, inputSpecB, outputSpec] =
detail::convertToTensorSpec(device,
std::make_tuple(inputShapeA, inputLayoutA),
std::make_tuple(inputShapeB, inputLayoutB),
std::make_tuple(outputShape, outputLayout));

return ::ttnn::graph::query_op_constraints(
::ttnn::add, device, inputSpecA, inputSpecB, outputSpec.data_type(),
outputSpec.tensor_layout().get_memory_config());
};

return operation::getOpConstraints("AddOpInterface", addOpQuery, inputShapeA,
inputLayoutA, inputShapeB, inputLayoutB,
outputShape, outputLayout);
return getEltwiseBinaryOpConstraints("AddOpInterface", ::ttnn::add,
inputShapeA, inputLayoutA, inputShapeB,
inputLayoutB, outputShape, outputLayout);
#else
return std::make_tuple(0, 0, 0);
#endif // TTMLIR_ENABLE_OPMODEL
Expand All @@ -310,31 +356,9 @@ AddOpInterface::getOpRuntime(llvm::ArrayRef<int64_t> inputShapeA,
llvm::ArrayRef<int64_t> outputShape,
mlir::tt::ttnn::TTNNLayoutAttr outputLayout) {
#ifdef TTMLIR_ENABLE_OPMODEL
auto addOpQuery = [](llvm::ArrayRef<int64_t> inputShapeA,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutA,
llvm::ArrayRef<int64_t> inputShapeB,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutB,
llvm::ArrayRef<int64_t> outputShape,
mlir::tt::ttnn::TTNNLayoutAttr outputLayout) {
// open device device, will close it at the end of function
::tt::tt_metal::v0::IDevice *device =
SingletonDeviceContext::getInstance().getDevice();

// prepare io specs
const auto [inputSpecA, inputSpecB, outputSpec] =
detail::convertToTensorSpec(device,
std::make_tuple(inputShapeA, inputLayoutA),
std::make_tuple(inputShapeB, inputLayoutB),
std::make_tuple(outputShape, outputLayout));

return ::ttnn::graph::query_op_runtime(
::ttnn::add, device, inputSpecA, inputSpecB, outputSpec.data_type(),
outputSpec.tensor_layout().get_memory_config());
};

return operation::getOpRuntime("AddOpInterface", addOpQuery, inputShapeA,
inputLayoutA, inputShapeB, inputLayoutB,
outputShape, outputLayout);
return getEltwiseBinaryOpRuntime("AddOpInterface", ::ttnn::add, inputShapeA,
inputLayoutA, inputShapeB, inputLayoutB,
outputShape, outputLayout);
#else
return llvm::createStringError("Not Implemented");
#endif // TTMLIR_ENABLE_OPMODEL
Expand Down Expand Up @@ -652,4 +676,37 @@ MatmulOpInterface::getOpRuntime(llvm::ArrayRef<int64_t> inputShapeA,
#endif // TTMLIR_ENABLE_OPMODEL
}

llvm::Expected<std::tuple<size_t, size_t, size_t>>
MultiplyOpInterface::getOpConstraints(
llvm::ArrayRef<int64_t> inputShapeA,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutA,
llvm::ArrayRef<int64_t> inputShapeB,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutB,
llvm::ArrayRef<int64_t> outputShape,
mlir::tt::ttnn::TTNNLayoutAttr outputLayout) {
#ifdef TTMLIR_ENABLE_OPMODEL
return getEltwiseBinaryOpConstraints("MultiplyOpInterface", ::ttnn::multiply,
inputShapeA, inputLayoutA, inputShapeB,
inputLayoutB, outputShape, outputLayout);
#else
return std::make_tuple(0, 0, 0);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Might be better to return a string error here as well.

#endif // TTMLIR_ENABLE_OPMODEL
}

llvm::Expected<size_t>
MultiplyOpInterface::getOpRuntime(llvm::ArrayRef<int64_t> inputShapeA,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutA,
llvm::ArrayRef<int64_t> inputShapeB,
mlir::tt::ttnn::TTNNLayoutAttr inputLayoutB,
llvm::ArrayRef<int64_t> outputShape,
mlir::tt::ttnn::TTNNLayoutAttr outputLayout) {
#ifdef TTMLIR_ENABLE_OPMODEL
return getEltwiseBinaryOpRuntime("MultiplyOpInterface", ::ttnn::multiply,
inputShapeA, inputLayoutA, inputShapeB,
inputLayoutB, outputShape, outputLayout);
#else
return llvm::createStringError("Not Implemented");
#endif // TTMLIR_ENABLE_OPMODEL
}

} // namespace mlir::tt::op_model::ttnn
Loading
Loading