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

Migrate glu variant ops to TTNN #10501

Merged
merged 5 commits into from
Jul 24, 2024
Merged

Migrate glu variant ops to TTNN #10501

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
d611243
#10384: Migrate reglu op to ttnn
umadevimcw Jul 20, 2024
a39e669
#10384: Migrate geglu op to ttnn
umadevimcw Jul 20, 2024
d6cd2ce
#10384: Update geglu and reglu in profiling script
umadevimcw Jul 20, 2024
f86227f
#10384: Migrate swiglu op to TTNN
umadevimcw Jul 20, 2024
a5c7cc1
#10384: Migrate glu op to TTNN
mouliraj-mcw Jul 23, 2024
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: 0 additions & 8 deletions docs/source/ttnn/ttnn/dependencies/tt_lib.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -608,14 +608,6 @@ Other Operations

.. autofunction:: tt_lib.tensor.normalize_global

.. autofunction:: tt_lib.tensor.glu

.. autofunction:: tt_lib.tensor.geglu

.. autofunction:: tt_lib.tensor.reglu

.. autofunction:: tt_lib.tensor.swiglu

.. autofunction:: tt_lib.tensor.embeddings

.. autofunction:: tt_lib.tensor.nextafter
Expand Down
5 changes: 4 additions & 1 deletion tests/tt_eager/python_api_testing/sweep_tests/pytests/tt_dnn/test_glu_variants.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,10 @@
@pytest.mark.parametrize("input_mem_config", input_mem_cfgs)
@pytest.mark.parametrize("output_mem_config", output_mem_cfgs)
class TestGLUVariants:
@pytest.mark.parametrize("fn_kind", ["glu", "reglu", "geglu", "swiglu"])
@pytest.mark.parametrize(
"fn_kind",
["glu", "reglu", "geglu", "swiglu"],
)
def test_all_glu_ops(
self,
input_shapes,
Expand Down
8 changes: 4 additions & 4 deletions tests/tt_eager/python_api_testing/sweep_tests/tt_lib_ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2846,7 +2846,7 @@ def unpad_from_tile(
def activation_glu(x, *args, device, dtype, layout, input_mem_config, output_mem_config, **kwargs):
dim = kwargs.get("dim", -1)
t0 = setup_tt_tensor(x, device, layout[0], input_mem_config[0], dtype[0])
t1 = ttl.tensor.glu(t0, dim, output_mem_config=output_mem_config)
t1 = ttnn.glu(t0, dim, memory_config=output_mem_config)

return tt2torch_tensor(t1)

Expand All @@ -2855,7 +2855,7 @@ def activation_glu(x, *args, device, dtype, layout, input_mem_config, output_mem
def activation_geglu(x, *args, device, dtype, layout, input_mem_config, output_mem_config, **kwargs):
dim = kwargs.get("dim", -1)
t0 = setup_tt_tensor(x, device, layout[0], input_mem_config[0], dtype[0])
t1 = ttl.tensor.geglu(t0, dim, output_mem_config=output_mem_config)
t1 = ttnn.geglu(t0, dim, memory_config=output_mem_config)

return tt2torch_tensor(t1)

Expand All @@ -2864,7 +2864,7 @@ def activation_geglu(x, *args, device, dtype, layout, input_mem_config, output_m
def activation_reglu(x, *args, device, dtype, layout, input_mem_config, output_mem_config, **kwargs):
dim = kwargs.get("dim", -1)
t0 = setup_tt_tensor(x, device, layout[0], input_mem_config[0], dtype[0])
t1 = ttl.tensor.reglu(t0, dim, output_mem_config=output_mem_config)
t1 = ttnn.reglu(t0, dim, memory_config=output_mem_config)

return tt2torch_tensor(t1)

Expand All @@ -2873,7 +2873,7 @@ def activation_reglu(x, *args, device, dtype, layout, input_mem_config, output_m
def activation_swiglu(x, *args, device, dtype, layout, input_mem_config, output_mem_config, **kwargs):
dim = kwargs.get("dim", -1)
t0 = setup_tt_tensor(x, device, layout[0], input_mem_config[0], dtype[0])
t1 = ttl.tensor.swiglu(t0, dim, output_mem_config=output_mem_config)
t1 = ttnn.swiglu(t0, dim, memory_config=output_mem_config)

return tt2torch_tensor(t1)

Expand Down
32 changes: 16 additions & 16 deletions tests/ttnn/profiling/ops_for_profiling.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1450,35 +1450,35 @@ def logical_noti(x):


def glu_1(x):
tt_lib.tensor.glu(x, -1)
ttnn.glu(x, -1)


def geglu_1(x):
tt_lib.tensor.geglu(x, -1)
ttnn.geglu(x, -1)


def reglu_1(x):
tt_lib.tensor.reglu(x, -1)
ttnn.reglu(x, -1)


def swiglu_1(x):
tt_lib.tensor.swiglu(x, -1)
ttnn.swiglu(x, -1)


def glu_2(x):
tt_lib.tensor.glu(x, -2)
ttnn.glu(x, -2)


def geglu_2(x):
tt_lib.tensor.geglu(x, -2)
ttnn.geglu(x, -2)


def reglu_2(x):
tt_lib.tensor.reglu(x, -2)
ttnn.reglu(x, -2)


def swiglu_2(x):
tt_lib.tensor.swiglu(x, -2)
ttnn.swiglu(x, -2)


def repeat(x):
Expand Down Expand Up @@ -2192,35 +2192,35 @@ def clone(x):
},
{
"op": glu_1,
"name": "tt_lib.tensor.glu_dim_3",
"name": "ttnn.glu_dim_3",
},
{
"op": geglu_1,
"name": "tt_lib.tensor.geglu_dim_3",
"name": "ttnn.geglu_dim_3",
},
{
"op": reglu_1,
"name": "tt_lib.tensor.reglu_dim_3",
"name": "ttnn.reglu_dim_3",
},
{
"op": swiglu_1,
"name": "tt_lib.tensor.swiglu_dim_3",
"name": "ttnn.swiglu_dim_3",
},
{
"op": glu_2,
"name": "tt_lib.tensor.glu_dim_2",
"name": "ttnn.glu_dim_2",
},
{
"op": geglu_2,
"name": "tt_lib.tensor.geglu_dim_2",
"name": "ttnn.geglu_dim_2",
},
{
"op": reglu_2,
"name": "tt_lib.tensor.reglu_dim_2",
"name": "ttnn.reglu_dim_2",
},
{
"op": swiglu_2,
"name": "tt_lib.tensor.swiglu_dim_2",
"name": "ttnn.swiglu_dim_2",
},
{
"op": repeat,
Expand Down
92 changes: 92 additions & 0 deletions tests/ttnn/unit_tests/operations/test_composite.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -394,3 +394,95 @@ def test_unary_threshold_ttnn(input_shapes, device):

comp_pass = compare_pcc([output_tensor], [golden_tensor])
assert comp_pass


@pytest.mark.parametrize(
"input_shapes",
(
(torch.Size([1, 1, 32, 64])),
(torch.Size([1, 1, 320, 384])),
(torch.Size([1, 3, 320, 384])),
),
)
@pytest.mark.parametrize(
"dim",
[-1, 3],
)
def test_unary_glu_ttnn(input_shapes, dim, device):
in_data, input_tensor = data_gen_with_range(input_shapes, -5, 5, device)
golden_fn = ttnn.get_golden_function(ttnn.glu)

output_tensor = ttnn.glu(input_tensor, dim)
golden_tensor = golden_fn(in_data, dim)

comp_pass = compare_pcc([output_tensor], [golden_tensor])
assert comp_pass


@pytest.mark.parametrize(
"input_shapes",
(
(torch.Size([1, 1, 32, 64])),
(torch.Size([1, 1, 320, 384])),
(torch.Size([1, 3, 320, 384])),
),
)
@pytest.mark.parametrize(
"dim",
[-1, 3],
)
def test_unary_reglu_ttnn(input_shapes, dim, device):
in_data, input_tensor = data_gen_with_range(input_shapes, -5, 5, device)
golden_fn = ttnn.get_golden_function(ttnn.reglu)

output_tensor = ttnn.reglu(input_tensor, dim)
golden_tensor = golden_fn(in_data, dim)

comp_pass = compare_pcc([output_tensor], [golden_tensor])
assert comp_pass


@pytest.mark.parametrize(
"input_shapes",
(
(torch.Size([1, 1, 32, 64])),
(torch.Size([1, 1, 320, 384])),
(torch.Size([1, 3, 320, 384])),
),
)
@pytest.mark.parametrize(
"dim",
[-1, 3],
)
def test_unary_geglu_ttnn(input_shapes, dim, device):
in_data, input_tensor = data_gen_with_range(input_shapes, -5, 5, device)
golden_fn = ttnn.get_golden_function(ttnn.geglu)

output_tensor = ttnn.geglu(input_tensor, dim)
golden_tensor = golden_fn(in_data, dim)

comp_pass = compare_pcc([output_tensor], [golden_tensor])
assert comp_pass


@pytest.mark.parametrize(
"input_shapes",
(
(torch.Size([1, 1, 32, 64])),
(torch.Size([1, 1, 320, 384])),
(torch.Size([1, 3, 320, 384])),
),
)
@pytest.mark.parametrize(
"dim",
[-1, 3],
)
def test_unary_swiglu_ttnn(input_shapes, dim, device):
in_data, input_tensor = data_gen_with_range(input_shapes, -5, 5, device)
golden_fn = ttnn.get_golden_function(ttnn.swiglu)

output_tensor = ttnn.swiglu(input_tensor, dim)
golden_tensor = golden_fn(in_data, dim)

comp_pass = compare_pcc([output_tensor], [golden_tensor])
assert comp_pass
85 changes: 0 additions & 85 deletions ttnn/cpp/ttnn/deprecated/tt_dnn/op_library/composite/composite_ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1879,91 +1879,6 @@ std::vector<Tensor> split_tensor_for_glu(const Tensor& input_a, int32_t dim, con
return t_split;
}

// Gated Linear Unit activation: matmul(split[0],sigmoid(split[1]))
Tensor _glu(
const Tensor& input_a,
int32_t dim /* = -1 */,
const MemoryConfig& output_mem_config /* = operation::DEFAULT_OUTPUT_MEMORY_CONFIG */) {
TT_ASSERT(dim == -1 || dim == 3, "last dim GLU only supported at this time ");
if (dim == -1)
dim = 3;

std::vector<Tensor> ab = split_tensor_for_glu(input_a, dim, output_mem_config);
Tensor sigmoid_b = ttnn::sigmoid(ab[1], output_mem_config);
Tensor glu_result = ttnn::multiply(ab[0], sigmoid_b, std::nullopt, output_mem_config);
return glu_result;
}
Tensor glu(
const Tensor& input_a,
int32_t dim /* = -1 */,
const MemoryConfig& output_mem_config /* = operation::DEFAULT_OUTPUT_MEMORY_CONFIG */) {
return operation::decorate_as_composite(__func__, _glu)(input_a, dim, output_mem_config);
}

// ReLU Gated Linear Unit activation: matmul(split[0],relu(split[1]))
Tensor _reglu(
const Tensor& input_a,
int32_t dim /* = -1 */,
const MemoryConfig& output_mem_config /* = operation::DEFAULT_OUTPUT_MEMORY_CONFIG */) {
TT_ASSERT(dim == -1 || dim == 3, "last dim REGLU only supported at this time ");
if (dim == -1)
dim = 3;
std::vector<Tensor> ab = split_tensor_for_glu(input_a, dim, output_mem_config);
Tensor relu_b = ttnn::relu(ab[1], output_mem_config);
Tensor reglu_result = ttnn::multiply(ab[0], relu_b, std::nullopt, output_mem_config);
return reglu_result;
}
Tensor reglu(
const Tensor& input_a,
int32_t dim /* = -1 */,
const MemoryConfig& output_mem_config /* = operation::DEFAULT_OUTPUT_MEMORY_CONFIG */) {
return operation::decorate_as_composite(__func__, _reglu)(input_a, dim, output_mem_config);
}

// Gaussian Error Gated Linear Unit activation: matmul(split[0],gelu(split[1]))
Tensor _geglu(
const Tensor& input_a,
int32_t dim /* = -1 */,
const MemoryConfig& output_mem_config /* = operation::DEFAULT_OUTPUT_MEMORY_CONFIG */) {
TT_ASSERT(dim == -1 || dim == 3, "last dim GEGLU only supported at this time ");
if (dim == -1)
dim = 3;

std::vector<Tensor> ab = split_tensor_for_glu(input_a, dim, output_mem_config);

constexpr bool fast_appx = true;
Tensor gelu_b = ttnn::gelu(ab[1], fast_appx, output_mem_config);
Tensor geglu_result = ttnn::multiply(ab[0], gelu_b, std::nullopt, output_mem_config);
return geglu_result;
}
Tensor geglu(
const Tensor& input_a,
int32_t dim /* = -1 */,
const MemoryConfig& output_mem_config /* = operation::DEFAULT_OUTPUT_MEMORY_CONFIG */) {
return operation::decorate_as_composite(__func__, _geglu)(input_a, dim, output_mem_config);
}

// Swish Gated Linear Unit activation: matmul(split[0],swish(split[1]))
Tensor _swiglu(
const Tensor& input_a,
int32_t dim /* = -1 */,
const MemoryConfig& output_mem_config /* = operation::DEFAULT_OUTPUT_MEMORY_CONFIG */) {
TT_ASSERT(dim == -1 || dim == 3, "last dim SWIGLU only supported at this time ");
if (dim == -1)
dim = 3;

std::vector<Tensor> ab = split_tensor_for_glu(input_a, dim, output_mem_config);

Tensor swish_b = swish(ab[1], output_mem_config);
Tensor swiglu_result = ttnn::multiply(ab[0], swish_b, std::nullopt, output_mem_config);
return swiglu_result;
}
Tensor swiglu(
const Tensor& input_a,
int32_t dim /* = -1 */,
const MemoryConfig& output_mem_config /* = operation::DEFAULT_OUTPUT_MEMORY_CONFIG */) {
return operation::decorate_as_composite(__func__, _swiglu)(input_a, dim, output_mem_config);
}

// on-device tensor creation with shape and filled with value
Tensor _sfpu_eps(const Shape shape, Layout layout, Device* device, const MemoryConfig& output_mem_config) {
Expand Down
21 changes: 0 additions & 21 deletions ttnn/cpp/ttnn/deprecated/tt_dnn/op_library/composite/composite_ops.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -563,27 +563,6 @@ Tensor logical_ori(
float immediate,
const MemoryConfig& output_mem_config = operation::DEFAULT_OUTPUT_MEMORY_CONFIG);

// Gated Linear Unit activation
Tensor glu(
const Tensor& input_a,
int32_t dim = -1,
const MemoryConfig& output_mem_config = operation::DEFAULT_OUTPUT_MEMORY_CONFIG);
// ReLU based GLU
Tensor reglu(
const Tensor& input_a,
int32_t dim = -1,
const MemoryConfig& output_mem_config = operation::DEFAULT_OUTPUT_MEMORY_CONFIG);
// Gelu based GLU
Tensor geglu(
const Tensor& input_a,
int32_t dim = -1,
const MemoryConfig& output_mem_config = operation::DEFAULT_OUTPUT_MEMORY_CONFIG);
// Swish based GLU
Tensor swiglu(
const Tensor& input_a,
int32_t dim = -1,
const MemoryConfig& output_mem_config = operation::DEFAULT_OUTPUT_MEMORY_CONFIG);

// on-device tensor creation with shape and filled with value
Tensor sfpu_eps(const Shape shape, Layout layout, Device* device, const MemoryConfig& output_mem_config);

Expand Down
Loading
Loading