Add basic Trilu op support (#1509)

supports upper and lower trilu
not able to support dynamic k values
not able to support negative k values

src/onnx/parse_trilu.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <migraphx/onnx/op_parser.hpp>
+#include <migraphx/onnx/checks.hpp>
+#include <migraphx/ranges.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/make_op.hpp>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace onnx {
+
+struct parse_trilu : op_parser<parse_trilu>
+{
+    std::vector<op_desc> operators() const { return {{"Trilu"}}; }
+
+    instruction_ref parse(const op_desc&,
+                          const onnx_parser&,
+                          const onnx_parser::node_info& info,
+                          std::vector<instruction_ref> args) const
+    {
+        auto input_shape = args[0]->get_shape();
+        assert(input_shape.ndim() >= 2);
+        auto input_lens = input_shape.lens();
+
+        size_t num_rows = *(input_lens.rbegin() + 1);
+        size_t num_cols = input_lens.back();
+        int k           = 0;
+        bool upper      = true;
+
+        if(args.size() > 1)
+        {
+            auto arg_k = args[1]->eval();
+            check_arg_empty(arg_k, "PARSE_TRILU: dynamic k not supported");
+            k = arg_k.at<int>();
+        }
+
+        if(k < 0)
+            MIGRAPHX_THROW("PARSE_TRILU: negative k values not supported");
+
+        if(contains(info.attributes, "upper"))
+        {
+            upper = static_cast<bool>(info.attributes.at("upper").i());
+        }
+
+        shape::type_t output_type = args[0]->get_shape().type();
+
+        // when creating the mask, if upper == 1,
+        // the inner triangle will have values set to 0
+        std::vector<bool> mask_mat(num_rows * num_cols, upper);
+        for(size_t i = 0; i < num_rows; i++)
+        {
+            for(size_t j = 0; j < std::min(k, static_cast<int>(num_cols)); j++)
+            {
+                mask_mat[i * num_cols + j] = not upper;
+            }
+            k++;
+        }
+
+        auto mask = info.add_literal(
+            migraphx::literal{migraphx::shape{output_type, {num_rows, num_cols}}, mask_mat});
+
+        return info.add_broadcastable_binary_op("mul", mask, args[0]);
+    }
+};
+
+} // namespace onnx
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

test/onnx/gen_onnx.py

@@ -6771,6 +6771,92 @@ def transpose_gather_test():
     return ([td, ti, node], [x, i], [y])
 
 
+@onnx_test()
+def trilu_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 4])
+
+    node = onnx.helper.make_node(
+        'Trilu',
+        inputs=['x'],
+        outputs=['y'],
+    )
+    return ([node], [x], [y])
+
+
+@onnx_test()
+def trilu_batch_diff_k_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [2, 2, 3])
+    k = np.array([2])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [2, 2, 3])
+    k_tensor = helper.make_tensor(name='k',
+                                  data_type=TensorProto.INT64,
+                                  dims=k.shape,
+                                  vals=k.astype(np.int64))
+
+    node = onnx.helper.make_node(
+        'Trilu',
+        inputs=['x', 'k'],
+        outputs=['y'],
+    )
+    return ([node], [x], [y], [k_tensor])
+
+
+@onnx_test()
+def trilu_lower_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 4])
+
+    node = onnx.helper.make_node('Trilu', inputs=['x'], outputs=['y'], upper=0)
+    return ([node], [x], [y])
+
+
+@onnx_test()
+def trilu_neg_k_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4])
+    k = np.array([-1])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 4])
+    k_tensor = helper.make_tensor(name='k',
+                                  data_type=TensorProto.INT64,
+                                  dims=k.shape,
+                                  vals=k.astype(np.int64))
+
+    node = onnx.helper.make_node('Trilu', inputs=['x', 'k'], outputs=['y'])
+    return ([node], [x], [y], [k_tensor])
+
+
+@onnx_test()
+def trilu_out_k_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [3, 4])
+    k = np.array([5])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [3, 4])
+    k_tensor = helper.make_tensor(name='k',
+                                  data_type=TensorProto.INT64,
+                                  dims=k.shape,
+                                  vals=k.astype(np.int64))
+
+    node = onnx.helper.make_node('Trilu', inputs=['x', 'k'], outputs=['y'])
+    return ([node], [x], [y], [k_tensor])
+
+
+@onnx_test()
+def trilu_row_one_test():
+    x = helper.make_tensor_value_info('x', TensorProto.FLOAT, [1, 4])
+    k = np.array([1])
+    y = helper.make_tensor_value_info('y', TensorProto.FLOAT, [1, 4])
+    k_tensor = helper.make_tensor(name='k',
+                                  data_type=TensorProto.INT64,
+                                  dims=k.shape,
+                                  vals=k.astype(np.int64))
+
+    node = onnx.helper.make_node(
+        'Trilu',
+        inputs=['x', 'k'],
+        outputs=['y'],
+    )
+    return ([node], [x], [y], [k_tensor])
+
+
 @onnx_test()
 def undefined_test():
     x = helper.make_tensor_value_info('0', TensorProto.FLOAT, [2, 3, 4, 5])

test/onnx/onnx_test.cpp

@@ -6552,6 +6552,11 @@ TEST_CASE(transpose_gather_test)
     EXPECT(p.sort() == prog.sort());
 }
 
+TEST_CASE(trilu_neg_k_test)
+{
+    EXPECT(test::throws([&] { migraphx::parse_onnx("trilu_neg_k_test.onnx"); }));
+}
+
 TEST_CASE(undefined_test)
 {
     migraphx::program p;

test/onnx/trilu_batch_diff_k_test.onnx

+trilu_batch_diff_k_test:i
+
+
x
+
k
y"Trilutrilu_batch_diff_k_test*
+
:
B
kZ
+
x
+

+
+
+b
+
y
+

+
+
+B
\ No newline at end of file

test/onnx/trilu_neg_k_test.onnx

+trilu_neg_k_test:c
+
+
x
+
k
y"Trilutrilu_neg_k_test*
:
+
B
kZ
+
x
+

+
+b
+
y
+

+
+B
\ No newline at end of file

test/onnx/trilu_out_k_test.onnx

+trilu_out_k_test:Z
+
+
x
+
k
y"Trilutrilu_out_k_test*
+
:
B
kZ
+
x
+

+
+b
+
y
+

+
+B
\ No newline at end of file

test/onnx/trilu_row_one_test.onnx

+trilu_row_one_test:\
+
+
x
+
k
y"Trilutrilu_row_one_test*
+
:

B
kZ
+
x
+

+

+b
+
y
+

+

+B
\ No newline at end of file

test/onnx/trilu_test.onnx

+
+trilu_test:E
+
+
x
y"Trilu
+trilu_testZ
+
x
+

+
+b
+
y
+

+
+B
\ No newline at end of file

test/onnx/verify_onnx.cpp

@@ -1458,4 +1458,73 @@ TEST_CASE(where_test)
     EXPECT(migraphx::verify_range(result_vector, gold));
 }
 
+std::vector<float> gen_trilu_test(const migraphx::shape& s, const migraphx::program& p)
+{
+    // input data filled with values 1 to nelements
+    std::vector<float> x_data(s.elements());
+    std::iota(x_data.begin(), x_data.end(), 1);
+
+    migraphx::parameter_map pp;
+    pp["x"] = migraphx::argument(s, x_data.data());
+
+    auto result = p.eval(pp).back();
+    std::vector<float> result_vector;
+    result.visit([&](auto output) { result_vector.assign(output.begin(), output.end()); });
+    return result_vector;
+}
+TEST_CASE(trilu_test)
+{
+    migraphx::program p = migraphx::parse_onnx("trilu_test.onnx");
+
+    std::vector<float> result_vector = gen_trilu_test({migraphx::shape::float_type, {3, 4}}, p);
+
+    std::vector<float> gold = {1, 2, 3, 4, 0, 6, 7, 8, 0, 0, 11, 12};
+
+    EXPECT(migraphx::verify_range(result_vector, gold));
+}
+
+TEST_CASE(trilu_batch_diff_k_test)
+{
+    migraphx::program p = migraphx::parse_onnx("trilu_batch_diff_k_test.onnx");
+
+    std::vector<float> result_vector = gen_trilu_test({migraphx::shape::float_type, {2, 2, 3}}, p);
+
+    std::vector<float> gold = {0, 0, 3, 0, 0, 0, 0, 0, 9, 0, 0, 0};
+
+    EXPECT(migraphx::verify_range(result_vector, gold));
+}
+
+TEST_CASE(trilu_lower_test)
+{
+    migraphx::program p = migraphx::parse_onnx("trilu_lower_test.onnx");
+
+    std::vector<float> result_vector = gen_trilu_test({migraphx::shape::float_type, {3, 4}}, p);
+
+    std::vector<float> gold = {0, 0, 0, 0, 5, 0, 0, 0, 9, 10, 0, 0};
+
+    EXPECT(migraphx::verify_range(result_vector, gold));
+}
+
+TEST_CASE(trilu_out_k_test)
+{
+    migraphx::program p = migraphx::parse_onnx("trilu_out_k_test.onnx");
+
+    std::vector<float> result_vector = gen_trilu_test({migraphx::shape::float_type, {3, 4}}, p);
+
+    std::vector<float> gold(12, 0);
+
+    EXPECT(migraphx::verify_range(result_vector, gold));
+}
+
+TEST_CASE(trilu_row_one_test)
+{
+    migraphx::program p = migraphx::parse_onnx("trilu_row_one_test.onnx");
+
+    std::vector<float> result_vector = gen_trilu_test({migraphx::shape::float_type, {1, 4}}, p);
+
+    std::vector<float> gold = {0, 2, 3, 4};
+
+    EXPECT(migraphx::verify_range(result_vector, gold));
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }