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

[lower_to_mlir] Lower reduce op from tt-forge to mlir #78

Merged
merged 1 commit into from
Aug 19, 2024
Merged

[lower_to_mlir] Lower reduce op from tt-forge to mlir #78

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
8 changes: 8 additions & 0 deletions pybuda/csrc/passes/lower_to_mlir.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,12 @@ class MLIRGenerator
return builder_.getSI32IntegerAttr(arg);
} else if constexpr (std::is_same_v<T, float>) {
return builder_.getF32FloatAttr(arg);
} else if constexpr (std::is_same_v<T, std::vector<int>>) {
llvm::SmallVector<mlir::Attribute> attributes;
for (auto& element : arg) {
attributes.push_back(builder_.getI32IntegerAttr(element));
}
return builder_.getArrayAttr(attributes);
} else {
// If type not handled, throw an exception or handle it appropriately
throw std::runtime_error("Unhandled attribute type");
Expand Down Expand Up @@ -445,6 +451,8 @@ class MLIRGenerator
lowering_handler_map["relu"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::ReluOp>;
lowering_handler_map["matmul"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::MatmulOp>;
lowering_handler_map["softmax"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::SoftmaxOp>;
lowering_handler_map["reduce_sum"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::SumOp>;
lowering_handler_map["reduce_avg"] = &MLIRGenerator::emit_mlir_ttforge_op<mlir::tt::ttir::MeanOp>;
}
};
}
Expand Down
21 changes: 21 additions & 0 deletions pybuda/pybuda/op/eval/common.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -191,6 +191,27 @@ def calculate_pcc(a, b):

return pcc

# Calculates pcc between golden and calculated tensors. If calculated pcc is >= than pcc threshold, returns True
def compare_with_golden_pcc(golden: Union[torch.Tensor, tf.Tensor, tf.Variable], calculated: torch.Tensor, pcc, rtol=None, atol=None):
pcc_value = 0
if not (pcc is None or golden.flatten().size() == (1,)): # PCC for single values doesn't work
nvukobratTT marked this conversation as resolved.
Show resolved Hide resolved
pcc_value = calculate_pcc(golden, calculated)
if pcc_value >= pcc :
logger.trace("PCC is correct")
logger.trace("Golden: (shape = {}", golden.shape)
logger.trace(golden)
logger.trace("Calculated: (shape = {}", calculated.shape)
logger.trace(calculated)
return True
else:
logger.error("Tensor mismatch")
return False
else:
# For scalar values, we can't calculate PCC, but we can compare golden and calculated values using relative and absolute tolerances
golden = golden.flatten()[0]
calculated = calculated.flatten()[0]
return torch.allclose(golden, calculated, atol=atol, rtol=rtol)

def compare_tensor_to_golden(name: str, golden: Union[torch.Tensor, tf.Tensor, tf.Variable], calculated: torch.Tensor, is_buda=False, rtol=None, atol=None, pcc=None, warning_only=False, relative_atol = None, verify_cfg = None):
# Convert golden to pytorch tensor for comparisons
if isinstance(golden, (tf.Tensor, tf.Variable)):
Expand Down
10 changes: 6 additions & 4 deletions pybuda/pybuda/op/reduce.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,8 @@
def ReduceSum(
name: str,
operandA: Tensor,
dim: int) -> Tensor:
dim: int,
keep_dim: bool = True) -> Tensor:
"""
Reduce by summing along the given dimension

Expand All @@ -32,12 +33,13 @@ def ReduceSum(
# if dim < 0:
# dim += 4

return op("reduce_sum", name, operandA, attrs=(dim,)).get_tensor()
return op("reduce_sum", name, operandA, attrs=(dim,), dim_arg=[dim], keep_dim= keep_dim).get_tensor()

def ReduceAvg(
name: str,
operandA: Tensor,
dim: int) -> Tensor:
dim: int,
keep_dim: bool = True) -> Tensor:
"""
Reduce by averaging along the given dimension

Expand All @@ -62,7 +64,7 @@ def ReduceAvg(
# if dim < 0:
# dim += 4

return op("reduce_avg", name, operandA, attrs=(dim,)).get_tensor()
return op("reduce_avg", name, operandA, attrs=(dim,), dim_arg=[dim], keep_dim= keep_dim).get_tensor()

def GroupedReduceAvg(
name: str,
Expand Down
3 changes: 2 additions & 1 deletion pybuda/test/mlir/mnist/test_inference.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
import torch
from .utils import *
import pybuda
from pybuda.op.eval.common import compare_with_golden_pcc

def test_mnist_inference():
inputs = [torch.rand(1, 784)]
Expand All @@ -16,4 +17,4 @@ def test_mnist_inference():
co_out = compiled_model(*[i.to("tt") for i in inputs])

co_out = [co.to("cpu") for co in co_out]
assert [torch.allclose(fo, co) for fo, co in zip(fw_out, co_out)]
assert [compare_with_golden_pcc(golden=fo, calculated=co, pcc=0.99) for fo, co in zip(fw_out, co_out)]
58 changes: 52 additions & 6 deletions pybuda/test/mlir/test_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,12 @@
import os
import pytest

import pytest
import torch
from torch import nn

import pybuda
from pybuda.op.eval.common import compare_with_golden_pcc

def test_add():
class Add(nn.Module):
Expand All @@ -27,7 +29,7 @@ def forward(self, a, b):
co_out = compiled_model(*inputs)

co_out = [co.to("cpu") for co in co_out]
assert [torch.allclose(fo, co) for fo, co in zip(fw_out, co_out)]
assert [compare_with_golden_pcc(golden=fo, calculated=co, pcc=0.99) for fo, co in zip(fw_out, co_out)]


def test_subtract():
Expand All @@ -47,7 +49,7 @@ def forward(self, a, b):
co_out = compiled_model(*inputs)

co_out = [co.to("cpu") for co in co_out]
assert [torch.allclose(fo, co) for fo, co in zip(fw_out, co_out)]
assert [compare_with_golden_pcc(golden=fo, calculated=co, pcc=0.99) for fo, co in zip(fw_out, co_out)]


def test_multiply():
Expand All @@ -67,7 +69,7 @@ def forward(self, a, b):
co_out = compiled_model(*inputs)

co_out = [co.to("cpu") for co in co_out]
assert [torch.allclose(fo, co) for fo, co in zip(fw_out, co_out)]
assert [compare_with_golden_pcc(golden=fo, calculated=co, pcc=0.99) for fo, co in zip(fw_out, co_out)]


def test_relu():
Expand All @@ -88,7 +90,7 @@ def forward(self, a):
co_out = compiled_model(*inputs)

co_out = [co.to("cpu") for co in co_out]
assert [torch.allclose(fo, co) for fo, co in zip(fw_out, co_out)]
assert [compare_with_golden_pcc(golden=fo, calculated=co, pcc=0.99) for fo, co in zip(fw_out, co_out)]

@pytest.mark.skip(reason="This is not ready yet")
def test_linear():
Expand All @@ -109,7 +111,7 @@ def forward(self, a):
co_out = compiled_model(*inputs)

co_out = [co.to("cpu") for co in co_out]
assert [torch.allclose(fo, co) for fo, co in zip(fw_out, co_out)]
assert [compare_with_golden_pcc(golden=fo, calculated=co, pcc=0.99) for fo, co in zip(fw_out, co_out)]


def test_softmax():
Expand All @@ -130,4 +132,48 @@ def forward(self, a):
co_out = compiled_model(*inputs)

co_out = [co.to("cpu") for co in co_out]
assert [torch.allclose(fo, co) for fo, co in zip(fw_out, co_out)]
assert [compare_with_golden_pcc(golden=fo, calculated=co, pcc=0.99) for fo, co in zip(fw_out, co_out)]

@pytest.mark.parametrize("input_shape", [(1,32,32), (1,64,64), (1,128,128,128)], ids=["32","64","128"])
@pytest.mark.parametrize("dim", [-1,-2], ids=["-1","-2"])
def test_reduce_sum(input_shape, dim):
class ReduceSum(nn.Module):
def __init__(self):
super().__init__()

def forward(self, a):
# reduce is supported on tt-metal only with keepdim=True
return torch.sum(a, dim=dim, keepdim=True)

inputs = [torch.rand(input_shape)]

framework_model = ReduceSum()
fw_out = framework_model(*inputs)

compiled_model = pybuda.compile(framework_model, sample_inputs=inputs)
co_out = compiled_model(*inputs)

co_out = [co.to("cpu") for co in co_out]
assert compare_with_golden_pcc(golden=fw_out, calculated=co_out[0], pcc=0.99)
Copy link
Contributor

Choose a reason for hiding this comment

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

Should we use this new function for verifying the output in other tests also?

Copy link
Contributor

Choose a reason for hiding this comment

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

I agree, can we update the rest of the test/mlir?

Also, please update pybuda/test/mlir/mnist/test_inference.py. At the moment, it has the similar issues as other ops. However, the current implementation doesn't detect it as [False] interpreted as valid assert case. In sum, code should look something like this:

def test_mnist_inference():
    inputs = [torch.rand(1, 784)]

    framework_model = MNISTLinear()
    fw_out = framework_model(*inputs)

    compiled_model = pybuda.compile(framework_model, sample_inputs=inputs)
    co_out = compiled_model(*[i.to("tt") for i in inputs])
    co_out = co_out[0].to("cpu")

    torch.allclose(fw_out, co_out) # This will do a proper failure. Should be replaced with compare_with_golden_pcc for valid verification.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I tried running this test but it fails due to data type error in ttnn: ttnn.linear: Tensor must be of type {DataType::BFLOAT16, DataType::BFLOAT8_B, DataType::BFLOAT4_B}, but got DataType::FLOAT32


@pytest.mark.parametrize("input_shape", [(1,32,32), (1,64,64), (1,128,128,128)], ids=["32","64","128"])
@pytest.mark.parametrize("dim", [-1,-2], ids=["-1","-2"])
def test_reduce_mean(input_shape, dim):
class ReduceMean(nn.Module):
def __init__(self):
super().__init__()

def forward(self, a):
# reduce is supported on tt-metal only with keepdim=True
return torch.mean(a, dim=1, keepdim=True)
nvukobratTT marked this conversation as resolved.
Show resolved Hide resolved

inputs = [torch.rand(1, 32, 32)]

framework_model = ReduceMean()
fw_out = framework_model(*inputs)

compiled_model = pybuda.compile(framework_model, sample_inputs=inputs)
co_out = compiled_model(*inputs)

co_out = [co.to("cpu") for co in co_out]
assert compare_with_golden_pcc(golden=fw_out, calculated=co_out[0], pcc=0.99)