[AIMIGRAPHX-1215] add runtime symbol resolution op

src/CMakeLists.txt

@@ -220,6 +220,7 @@ register_migraphx_ops(
     elu
     equal
     erf
+    eval_expr_from_shape
     exp
     fill
     fixed_pad

src/include/migraphx/op/eval_expr_from_shape.hpp

@@ -0,0 +1,129 @@
+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2026 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.
+ */
+#ifndef MIGRAPHX_GUARD_OPERATORS_EVAL_EXPR_FROM_SHAPE_HPP
+#define MIGRAPHX_GUARD_OPERATORS_EVAL_EXPR_FROM_SHAPE_HPP
+
+#include <migraphx/argument.hpp>
+#include <migraphx/check_shapes.hpp>
+#include <migraphx/config.hpp>
+#include <migraphx/context.hpp>
+#include <migraphx/shape.hpp>
+#include <migraphx/sym.hpp>
+#include <migraphx/zip_view.hpp>
+#include <algorithm>
+#include <unordered_map>
+#include <unordered_set>
+#include <vector>
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+namespace op {
+
+struct eval_expr_from_shape
+{
+    std::vector<sym::expr> expressions{};
+    std::vector<shape> input_shapes{};
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.expressions, "expressions"), f(self.input_shapes, "input_shapes"));
+    }
+
+    std::string name() const { return "eval_expr_from_shape"; }
+
+    shape compute_shape(const std::vector<shape>& inputs) const
+    {
+        check_shapes{inputs, *this, true}.has_at_least(1);
+
+        std::unordered_set<sym::expr> missing;
+        for(const auto& expression : expressions)
+        {
+            auto variables = sym::find_variables(expression);
+            missing.merge(variables);
+        }
+
+        for(const auto& input : inputs)
+        {
+            if(not input.symbolic())
+                continue;
+            for(const auto& d : input.dyn_dims())
+                if(d.sym_expr.name() == "variable")
+                    missing.erase(sym::as_symbol(d.sym_expr));
+        }
+        if(not missing.empty())
+            MIGRAPHX_THROW("EVAL_EXPR_FROM_SHAPE: Symbol '" + missing.begin()->to_string() +
+                           "' is not a direct input dimension");
+
+        return shape{shape::int64_type, {expressions.size()}};
+    }
+
+    void finalize(context&, const shape&, const std::vector<shape>& inputs)
+    {
+        input_shapes = inputs;
+    }
+
+    argument compute(const shape&, std::vector<argument> args) const
+    {
+        if(input_shapes.empty() or input_shapes.size() != args.size())
+            MIGRAPHX_THROW("EVAL_EXPR_FROM_SHAPE: input shapes not captured; op was not finalized");
+
+        std::unordered_map<sym::expr, std::size_t> values;
+        for(auto&& [input_shape, arg] : views::zip(input_shapes, args))
+        {
+            const auto& lens = arg.get_shape().lens();
+            if(input_shape.ndim() != lens.size())
+                MIGRAPHX_THROW("EVAL_EXPR_FROM_SHAPE: Runtime input rank does not match its "
+                               "symbolic shape");
+            if(not input_shape.symbolic())
+                continue;
+            const auto& dims = input_shape.dyn_dims();
+            for(auto&& [dim, len] : views::zip(dims, lens))
+            {
+                if(dim.sym_expr.name() != "variable")
+                    continue;
+                auto variable = sym::as_symbol(dim.sym_expr);
+                auto result   = values.emplace(variable, len);
+                if(not result.second and result.first->second != len)
+                    MIGRAPHX_THROW("EVAL_EXPR_FROM_SHAPE: Repeated symbol has inconsistent runtime "
+                                   "dimensions");
+            }
+        }
+
+        argument result{shape{shape::int64_type, {expressions.size()}}};
+        result.visit([&](auto output) {
+            std::transform(expressions.begin(),
+                           expressions.end(),
+                           output.begin(),
+                           [&](const auto& e) { return e.eval_uint(values); });
+        });
+        return result;
+    }
+};
+
+} // namespace op
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx
+
+#endif

src/include/migraphx/sym.hpp

@@ -36,6 +36,7 @@
 #include <string>
 #include <string_view>
 #include <unordered_map>
+#include <unordered_set>
 #include <utility>
 #include <variant>
 #include <vector>
@@ -254,6 +255,8 @@ MIGRAPHX_EXPORT expr var(std::string name, interval constraint, std::set<scalar>
 MIGRAPHX_EXPORT expr as_symbol(const expr& e, int max_depth = -1);
 MIGRAPHX_EXPORT bool same_symbol(const expr& a, const expr& b);
 
+// Find distinct variables as metadata-free symbols.
+MIGRAPHX_EXPORT std::unordered_set<expr> find_variables(const expr& e);
 // Whether dividend is evenly divisible by divisor (integral operands only).
 MIGRAPHX_EXPORT bool is_divisible(const expr& dividend, const expr& divisor);
 

src/instruction.cpp

@@ -379,6 +379,9 @@ bool instruction::can_eval() const
         return true;
     if(not is_context_free(op))
         return false;
+    // Finalize-dependent ops cannot be evaluated during pre-finalization constant folding.
+    if(has_finalize(op))
+        return false;
 #if MIGRAPHX_HAS_PMR
     std::array<char, 1024> storage;
     std::pmr::monotonic_buffer_resource resource{storage.data(), storage.size()};
@@ -393,7 +396,7 @@ bool instruction::can_eval() const
         bool evaluable = false;
         if(ins.name() == "@literal")
             evaluable = true;
-        else if(is_context_free(ins.get_operator()))
+        else if(is_context_free(ins.get_operator()) and not has_finalize(ins.get_operator()))
             evaluable = std::all_of(
                 ins.inputs().begin(), ins.inputs().end(), [&](auto arg) { return self(*arg); });
         cache.emplace(&ins, evaluable);

src/sym.cpp

@@ -1623,6 +1623,24 @@ bool same_symbol(const expr& a, const expr& b)
            });
 }
 
+std::unordered_set<expr> find_variables(const expr& e)
+{
+    std::unordered_set<expr> visited;
+    std::unordered_set<expr> result;
+    fix([&](auto self, const expr& x) {
+        if(x.empty() or not visited.insert(x).second)
+            return;
+        if(x.name() == "variable")
+        {
+            result.insert(as_symbol(x));
+            return;
+        }
+        for(const auto& c : x.children())
+            self(c);
+    })(e);
+    return result;
+}
+
 [[maybe_unused]] static bool has_float_literal(const expr& e)
 {
     if(e.empty())

src/targets/ref/lowering.cpp

@@ -183,6 +183,11 @@ struct ref_op
     {
         return op.compute(output_shape, args);
     }
+    void
+    finalize(migraphx::context& ctx, const shape& output_shape, const std::vector<shape>& inputs)
+    {
+        op.finalize(ctx, output_shape, inputs);
+    }
     value to_value() const
     {
         value v;

test/instruction.cpp

@@ -28,6 +28,39 @@
 #include "test.hpp"
 #include "rob.hpp"
 
+struct can_eval_finalize_passthrough
+{
+    std::string name() const { return "can_eval_finalize_passthrough"; }
+
+    migraphx::shape compute_shape(const std::vector<migraphx::shape>& inputs) const
+    {
+        return inputs.at(0);
+    }
+
+    migraphx::argument compute(const migraphx::shape&,
+                               const std::vector<migraphx::argument>& args) const
+    {
+        return args.at(0);
+    }
+
+    void finalize(migraphx::context&, const migraphx::shape&, const std::vector<migraphx::shape>&)
+    {
+    }
+};
+
+TEST_CASE(can_eval_rejects_finalize_op)
+{
+    migraphx::module m;
+    auto one       = m.add_literal(1);
+    auto evaluable = m.add_instruction(migraphx::make_op("identity"), one);
+    auto finalized = m.add_instruction(can_eval_finalize_passthrough{}, one);
+    auto dependent = m.add_instruction(migraphx::make_op("identity"), finalized);
+
+    EXPECT(evaluable->can_eval());
+    EXPECT(not finalized->can_eval());
+    EXPECT(not dependent->can_eval());
+}
+
 TEST_CASE(check_undefined)
 {
     migraphx::module m;

test/op_shape_test.cpp

@@ -5585,6 +5585,61 @@ TEST_CASE(slice_dyn_nonfixed_keeps_other_optimals)
                  input);
 }
 
+TEST_CASE(eval_expr_from_shape_shape)
+{
+    auto n = var("n", {1, 16});
+    auto h = var("h", {1, 32});
+    auto w = var("w", {1, 32});
+    migraphx::shape input{migraphx::shape::float_type, {dd{n}, dd{lit(3)}, dd{h}, dd{w}}};
+    expect_shape(migraphx::shape{migraphx::shape::int64_type, {3}},
+                 migraphx::make_op("eval_expr_from_shape",
+                                   {{"expressions",
+                                     migraphx::value::array{migraphx::to_value(n),
+                                                            migraphx::to_value(h / lit(2)),
+                                                            migraphx::to_value(w / lit(2))}}}),
+                 input);
+}
+
+TEST_CASE(eval_expr_from_shape_missing_symbol)
+{
+    auto m = var("m", {1, 16});
+    auto n = var("n", {1, 16});
+    migraphx::shape input{migraphx::shape::float_type, {dd{n}, dd{lit(3)}}};
+    throws_shape(
+        migraphx::make_op("eval_expr_from_shape",
+                          {{"expressions", migraphx::value::array{migraphx::to_value(m)}}}),
+        input);
+}
+
+TEST_CASE(eval_expr_from_shape_multi_input)
+{
+    auto m = var("m", {1, 16});
+    auto n = var("n", {1, 16});
+    migraphx::shape a{migraphx::shape::float_type, {dd{m}, dd{lit(3)}}};
+    migraphx::shape b{migraphx::shape::float_type, {dd{lit(2)}, dd{n}}};
+    expect_shape(migraphx::shape{migraphx::shape::int64_type, {2}},
+                 migraphx::make_op(
+                     "eval_expr_from_shape",
+                     {{"expressions",
+                       migraphx::value::array{migraphx::to_value(m + n), migraphx::to_value(m)}}}),
+                 a,
+                 b);
+}
+
+TEST_CASE(eval_expr_from_shape_missing_symbol_multi_input)
+{
+    auto m = var("m", {1, 16});
+    auto n = var("n", {1, 16});
+    auto k = var("k", {1, 16});
+    migraphx::shape a{migraphx::shape::float_type, {dd{m}, dd{lit(3)}}};
+    migraphx::shape b{migraphx::shape::float_type, {dd{lit(2)}, dd{n}}};
+    throws_shape(
+        migraphx::make_op("eval_expr_from_shape",
+                          {{"expressions", migraphx::value::array{migraphx::to_value(m + k)}}}),
+        a,
+        b);
+}
+
 TEST_CASE(slice_sym)
 {
     auto n                                      = var("n", {1, 8});