morok_ir/uop/

core.rs

//! Core UOp struct and fundamental operations.
//!
//! This module contains the [`UOp`] struct definition and its core methods
//! for accessing operation data, dtype, shape, and graph traversal.

use std::collections::{HashMap, HashSet};
use std::hash::{Hash, Hasher};
use std::sync::Arc;

use bon::bon;
use smallvec::SmallVec;

use crate::op::Op;
use crate::pattern::{Matcher, RewriteResult};
use crate::shape;
use crate::types::{AxisType, ConstValue};
use morok_dtype::DType;

/// Matcher for `UOp::substitute` — looks up each node in a substitution map.
///
/// Equivalent to Tinygrad's `_substitute = PatternMatcher([(UPat(tuple(Ops)), lambda ctx,x: ctx.get(x))])`.
struct SubstituteMatcher<'a>(&'a HashMap<UOpKey, Arc<UOp>>);

impl Matcher<()> for SubstituteMatcher<'_> {
    fn rewrite(&self, uop: &Arc<UOp>, _ctx: &mut ()) -> RewriteResult {
        match self.0.get(&UOpKey(uop.clone())) {
            Some(replacement) if !Arc::ptr_eq(uop, replacement) => RewriteResult::Rewritten(replacement.clone()),
            _ => RewriteResult::NoMatch,
        }
    }
}

/// Matcher for `UOp::substitute_gated` — substitution with range-scope gating.
///
/// Equivalent to Tinygrad's `_substitute` + `pm_gate_substitute`:
/// - If a node is in the substitution map, replace it.
/// - If a node's ranges don't overlap with substitution keys, gate (skip subtree).
struct SubstituteGatedMatcher<'a> {
    map: &'a HashMap<UOpKey, Arc<UOp>>,
    range_keys: &'a HashSet<UOpKey>,
}

impl Matcher<()> for SubstituteGatedMatcher<'_> {
    fn rewrite(&self, uop: &Arc<UOp>, _ctx: &mut ()) -> RewriteResult {
        // Direct substitution lookup
        if let Some(replacement) = self.map.get(&UOpKey(uop.clone()))
            && !Arc::ptr_eq(uop, replacement)
        {
            return RewriteResult::Rewritten(replacement.clone());
        }
        // Gate: skip subtrees whose ranges don't overlap with substitution keys
        // Tinygrad (rangeify.py:187): `if not any(r in b.ranges for r in ctx.keys()): raise BottomUpGate()`
        if !uop.in_scope_ranges().iter().any(|r| self.range_keys.contains(r)) {
            return RewriteResult::Gate(uop.clone());
        }
        RewriteResult::NoMatch
    }
}

/// Wrapper for Arc<UOp> that implements Hash and Eq based on stable ID.
///
/// This allows using Arc<UOp> as HashMap keys without implementing
/// Hash/Eq on UOp itself (which would be problematic due to OnceCell fields).
///
/// Note: While UOp contains OnceCell fields, Hash/Eq are based solely on the
/// immutable `id` field, making this safe to use as a HashMap key.
#[allow(clippy::mutable_key_type)]
#[derive(Clone)]
pub struct UOpKey(pub Arc<UOp>);

// Custom Debug impl to show only the UOp ID, avoiding recursive printing
impl std::fmt::Debug for UOpKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "UOpKey(id={})", self.0.id)
    }
}

impl PartialEq for UOpKey {
    fn eq(&self, other: &Self) -> bool {
        self.0.id == other.0.id
    }
}

impl Eq for UOpKey {}

impl Hash for UOpKey {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.0.id.hash(state);
    }
}

/// Micro-operation node in the computation graph.
///
/// UOps form a DAG where operations reference their inputs through the Op enum.
/// Hash consing ensures that structurally identical UOps share the same allocation.
///
/// Shape inference is lazy and cached - computed on first access via `shape()` method.
///
/// Note: Debug uses derive_more with `#[debug(skip)]` on cache fields to prevent
/// stack overflow from recursive Arc<UOp> references in caches.
#[derive(derive_more::Debug)]
pub struct UOp {
    /// Unique stable ID for this UOp instance.
    /// Used for identity-based caching instead of fragile raw pointers.
    pub id: u64,
    pub(crate) op: Op,
    pub(crate) dtype: DType,
    /// Cached shape - computed lazily on first access.
    /// OnceLock provides thread-safe lazy initialization.
    #[debug(skip)]
    pub(crate) shape_cache: std::sync::OnceLock<crate::Result<Option<shape::Shape>>>,
    /// Cached list of RANGE operations in this UOp's graph.
    /// Computed lazily via toposort to collect all RANGE ops.
    #[debug(skip)]
    pub(crate) ranges_cache: std::sync::OnceLock<Vec<Arc<UOp>>>,
    /// Cached set of RANGE operations that are in scope at this UOp.
    /// Unlike ranges_cache which contains ALL ranges in the graph,
    /// this contains only the ranges that are currently "active" (not yet ended).
    /// Computed lazily based on Tinygrad's ranges property.
    /// Uses UOpKey wrapper to enable Hash/Eq based on UOp ID.
    #[debug(skip)]
    pub(crate) in_scope_ranges_cache: std::sync::OnceLock<HashSet<UOpKey>>,
    /// Cached vmin/vmax range analysis values.
    /// Computed lazily via range propagation through the computation graph.
    /// Returns (vmin, vmax) as ConstValue types.
    #[debug(skip)]
    pub(crate) vmin_vmax_cache: std::sync::OnceLock<(ConstValue, ConstValue)>,
    /// Sound vmin/vmax: `None` for ops where range analysis is unsound (LOAD, Pow, etc.).
    /// Used by patterns that must not act on unsound bounds (e.g., vmin_vmax_collapse).
    #[debug(skip)]
    pub(crate) sound_vmin_vmax_cache: std::sync::OnceLock<Option<(ConstValue, ConstValue)>>,
    /// Whether this node or any of its sources is an INDEX op.
    /// Cached O(1) lookup used by `simplify_valid` to skip And chains inside INDEX trees.
    #[debug(skip)]
    pub(crate) has_index_in_sources_cache: std::sync::OnceLock<bool>,
    /// Cached backward slice: IDs of all nodes reachable from this UOp (including self).
    /// O(1) membership test via `backward_slice_ids().contains(&target.id)`.
    #[debug(skip)]
    pub(crate) backward_slice_cache: std::sync::OnceLock<HashSet<u64>>,
    /// Structural content hash — deterministic regardless of allocation order.
    /// Computed at creation time: hash(op_discriminant, dtype, op_data, children_content_hashes).
    /// O(1) per node since children are already created with their content_hash set.
    /// Used for schedule-level caching where UOp IDs are not stable across runs.
    pub content_hash: u64,
    /// Tag for tracking tensor identity through the rangeify pipeline.
    ///
    /// Matches Tinygrad's `UOp.tag` (ops.py:128). Tags are tuples of integer indices
    /// that track which original tensor UOps map to which final kernel outputs.
    /// Tags participate in hash consing — different tag = different UOp.
    ///
    /// Values:
    /// - `None` — untagged (default)
    /// - `Some([])` — empty tag (e.g., RANGE ops)
    /// - `Some([i])` — single index (assigned by add_tags)
    /// - `Some([i, j, ...])` — merged indices (from buffer folding)
    pub tag: Option<SmallVec<[usize; 2]>>,
    /// Optional metadata attached to this UOp.
    ///
    /// Metadata is NOT part of hash consing - attaching metadata creates a new UOp
    /// instance with a different ID. This is used for kernel info (name, opts) after
    /// optimization is complete.
    ///
    /// Uses Arc<dyn Any> to allow attaching any metadata type without circular
    /// dependencies (e.g., schedule::KernelInfo).
    #[debug(skip)]
    pub(crate) metadata: Option<std::sync::Arc<dyn std::any::Any + Send + Sync>>,
}

/// Hash implementation for UOp based on content (dtype + op).
///
/// This enables content-based hashing for cross-run caching. The hash traverses
/// the DAG structure since Op contains Arc<UOp> children that also get hashed.
/// Cache fields are intentionally skipped - they don't affect semantic identity.
impl Hash for UOp {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.dtype.hash(state);
        self.op.hash(state);
        // Intentionally skip: id, caches, metadata
    }
}

impl UOp {
    /// Get the operation.
    pub fn op(&self) -> &Op {
        &self.op
    }

    /// Get the data type.
    pub fn dtype(&self) -> DType {
        self.dtype.clone()
    }

    /// Get the tag.
    pub fn tag(&self) -> &Option<SmallVec<[usize; 2]>> {
        &self.tag
    }

    /// Create a new UOp with the given tag (Tinygrad: `rtag()`).
    /// Returns self unchanged if tag is already equal.
    pub fn rtag(self: &Arc<Self>, tag: Option<SmallVec<[usize; 2]>>) -> Arc<Self> {
        if self.tag == tag {
            return self.clone();
        }
        Self::new_tagged(self.op.clone(), self.dtype.clone(), tag)
    }

    /// Create a new UOp with the given tag set.
    pub fn with_tag(self: &Arc<Self>, tag: SmallVec<[usize; 2]>) -> Arc<Self> {
        self.rtag(Some(tag))
    }

    /// Check if this UOp has a concrete buffer identity in the graph.
    ///
    /// Returns true for BUFFER or RESHAPE/MULTI chains leading to BUFFER.
    /// These are already contiguous by definition, so wrapping in CONTIGUOUS is a no-op.
    ///
    /// Based on Tinygrad's `UOp.has_buffer_identity()` (ops.py:616-619).
    pub fn has_buffer_identity(&self) -> bool {
        match &self.op {
            Op::Reshape { src, .. } => src.has_buffer_identity(),
            Op::Buffer { .. } => true,
            _ => false,
        }
    }

    /// Get pointer dtype components if this UOp has a Ptr dtype.
    ///
    /// Returns `(base, addrspace, size)` for Ptr types, None otherwise.
    /// This simplifies pattern matching on pointer types.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use morok_ir::UOp;
    /// # use morok_dtype::{DType, AddrSpace, DeviceSpec};
    /// let buffer = UOp::new_buffer(DeviceSpec::Cpu, 10, DType::Float32);
    /// if let Some((base, addrspace, size)) = buffer.ptrdtype() {
    ///     assert_eq!(*base, DType::Float32);
    ///     assert_eq!(addrspace, AddrSpace::Global);
    /// }
    /// ```
    pub fn ptrdtype(&self) -> Option<(&DType, morok_dtype::AddrSpace, Option<usize>)> {
        match &self.dtype {
            DType::Ptr { base, addrspace, size, .. } => Some((base.as_ref(), *addrspace, *size)),
            _ => None,
        }
    }

    /// Create a copy of this UOp with a different dtype.
    ///
    /// If the dtype is unchanged, returns self (clone of Arc).
    /// This is the Rust equivalent of Tinygrad's `buf.replace(dtype=x)`.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use std::sync::Arc;
    /// # use morok_ir::UOp;
    /// # use morok_dtype::DType;
    /// let int_const = UOp::const_(DType::Int32, morok_ir::ConstValue::Int(5));
    /// let float_const = int_const.with_dtype(DType::Float32);
    /// assert_eq!(float_const.dtype(), DType::Float32);
    /// ```
    pub fn with_dtype(self: &Arc<Self>, dtype: DType) -> Arc<Self> {
        if self.dtype == dtype {
            return self.clone();
        }
        Self::new(self.op.clone(), dtype)
    }

    /// Walk through AFTER nodes to get the passthrough value.
    ///
    /// This is the Rust equivalent of Tinygrad's `.or_after()` pattern.
    /// Recursively unwraps AFTER nodes to find the underlying value.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// // Given: AFTER(AFTER(value, [dep1]), [dep2])
    /// // Returns: value
    /// let inner = wrapped.unwrap_after();
    /// ```
    pub fn unwrap_after(self: &Arc<Self>) -> Arc<Self> {
        match self.op() {
            Op::After { passthrough, .. } => passthrough.unwrap_after(),
            _ => self.clone(),
        }
    }

    /// Walk through CAST nodes to get the inner value.
    ///
    /// This is the Rust equivalent of Tinygrad's `.or_casted()` pattern.
    /// Recursively unwraps CAST nodes to find the underlying value.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// // Given: CAST(CAST(value, dtype1), dtype2)
    /// // Returns: value
    /// let inner = casted.unwrap_cast();
    /// ```
    pub fn unwrap_cast(self: &Arc<Self>) -> Arc<Self> {
        match self.op() {
            Op::Cast { src, .. } => src.unwrap_cast(),
            _ => self.clone(),
        }
    }

    /// Get the buffer from a STORE operation (via its INDEX child).
    ///
    /// STORE operations reference the buffer indirectly through an INDEX node.
    /// This helper extracts the buffer from `STORE.index.buffer`.
    ///
    /// Returns `None` if:
    /// - This is not a STORE operation
    /// - The STORE's index is not an INDEX operation
    pub fn store_buffer(&self) -> Option<&Arc<UOp>> {
        match self.op() {
            Op::Store { index, .. } => match index.op() {
                Op::Index { buffer, .. } => Some(buffer),
                _ => None,
            },
            _ => None,
        }
    }

    /// Get the buffer from a LOAD operation.
    ///
    /// Returns `None` if this is not a LOAD operation.
    pub fn load_buffer(&self) -> Option<Arc<UOp>> {
        match self.op() {
            Op::Load { buffer, .. } => Some(buffer.clone()),
            _ => None,
        }
    }

    /// Store a value at this INDEX node.
    ///
    /// Convenience method for `self.store(value)`.
    /// Matches Tinygrad's `idx.store(val)` pattern.
    ///
    /// # Panics
    ///
    /// Debug-asserts that self is an INDEX operation.
    pub fn store_value(self: &Arc<Self>, value: Arc<Self>) -> Arc<Self> {
        debug_assert!(matches!(self.op(), Op::Index { .. }), "store_value requires INDEX");
        self.store(value)
    }

    /// Alias for `with_sources()`.
    ///
    /// Creates a new UOp with the same operation type and dtype, but with
    /// the provided sources replacing the original ones.
    pub fn with_src(self: &Arc<Self>, new_srcs: Vec<Arc<Self>>) -> Arc<Self> {
        self.with_sources(new_srcs)
    }

    /// Get the shape of this UOp.
    ///
    /// Shape is computed lazily on first access and cached.
    /// Returns Ok(None) if shape cannot be determined (e.g., for control flow ops).
    /// Returns Err if there is a shape mismatch error.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use morok_ir::{UOp, ConstValue};
    /// # use morok_dtype::DType;
    /// let scalar = UOp::const_(DType::Float32, ConstValue::Float(1.0));
    /// assert_eq!(scalar.shape().unwrap().as_ref().map(|s| s.len()), Some(0)); // Scalar has empty shape
    /// ```
    pub fn shape(self: &Arc<Self>) -> crate::Result<Option<&shape::Shape>> {
        use crate::uop::cached_property::CachedProperty;
        use crate::uop::properties::ShapeProperty;
        match ShapeProperty::get(self) {
            Ok(opt) => Ok(opt.as_ref()),
            Err(e) => Err(e.clone()),
        }
    }

    /// Get the minimum possible value of this UOp.
    ///
    /// Returns the minimum value based on range analysis.
    /// Computed lazily on first access and cached.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use morok_ir::{UOp, ConstValue};
    /// # use morok_dtype::DType;
    /// let five = UOp::const_(DType::Int32, ConstValue::Int(5));
    /// assert_eq!(five.vmin(), &ConstValue::Int(5));
    /// ```
    pub fn vmin(self: &Arc<Self>) -> &ConstValue {
        use crate::uop::cached_property::CachedProperty;
        use crate::uop::properties::VminVmaxProperty;
        &VminVmaxProperty::get(self).0
    }

    /// Get the maximum possible value of this UOp.
    ///
    /// Returns the maximum value based on range analysis.
    /// Computed lazily on first access and cached.
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use morok_ir::{UOp, ConstValue};
    /// # use morok_dtype::DType;
    /// let five = UOp::const_(DType::Int32, ConstValue::Int(5));
    /// assert_eq!(five.vmax(), &ConstValue::Int(5));
    /// ```
    pub fn vmax(self: &Arc<Self>) -> &ConstValue {
        use crate::uop::cached_property::CachedProperty;
        use crate::uop::properties::VminVmaxProperty;
        &VminVmaxProperty::get(self).1
    }

    /// Extract device specification from this UOp graph.
    ///
    /// Traverses the graph to find Op::Device nodes following Tinygrad's
    /// `_device` recursive property (ops.py:585-599):
    /// - DEVICE: returns the DeviceSpec directly
    /// - BUFFER: returns device from the device child
    /// - COPY: returns device from the device child (target device)
    /// - Otherwise: searches children recursively
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use morok_ir::UOp;
    /// # use morok_dtype::{DType, DeviceSpec};
    /// let buffer = UOp::new_buffer(DeviceSpec::Cpu, 10, DType::Float32);
    /// assert_eq!(buffer.device_spec(), Some(DeviceSpec::Cpu));
    /// ```
    pub fn device_spec(&self) -> Option<morok_dtype::DeviceSpec> {
        match self.op() {
            Op::Device(spec) => Some(spec.clone()),
            Op::Buffer { device, .. } => {
                if let Op::Device(spec) = device.op() {
                    Some(spec.clone())
                } else {
                    None
                }
            }
            Op::Param { device: Some(device), .. } => {
                if let Op::Device(spec) = device.op() {
                    Some(spec.clone())
                } else {
                    None
                }
            }
            Op::Param { device: None, .. } => None,
            Op::Copy { device, .. } => {
                if let Op::Device(spec) = device.op() {
                    Some(spec.clone())
                } else {
                    None
                }
            }
            _ => {
                // Search children for device
                for child in self.op().children() {
                    if let Some(spec) = child.device_spec() {
                        return Some(spec);
                    }
                }
                None
            }
        }
    }

    /// Get the base UOp by walking through movement operations.
    ///
    /// Movement operations (RESHAPE, PERMUTE, EXPAND, etc.) are views that don't
    /// change the underlying data. This method recursively walks through these
    /// operations to find the actual buffer or computation that owns the data.
    ///
    /// Based on Tinygrad's `base` property (ops.py:524-527).
    ///
    /// # Examples
    ///
    /// ```rust
    /// # use morok_ir::{UOp, SInt, shape::Shape};
    /// # use morok_dtype::DType;
    /// # use morok_dtype::DeviceSpec;
    /// let buffer = UOp::new_buffer(DeviceSpec::Cpu, 10, DType::Float32);
    /// let shape = Shape::from_iter([SInt::Const(2), SInt::Const(5)]);
    /// let reshaped = buffer.try_reshape(&shape).unwrap();
    ///
    /// // base() walks through RESHAPE to get the original BUFFER
    /// assert!(std::sync::Arc::ptr_eq(&reshaped.base(), &buffer));
    /// ```
    pub fn base(self: &Arc<Self>) -> Arc<Self> {
        match &self.op {
            // Movement operations - recursively get base of source
            Op::Reshape { src, .. }
            | Op::Permute { src, .. }
            | Op::Expand { src, .. }
            | Op::Pad { src, .. }
            | Op::Shrink { src, .. }
            | Op::Flip { src, .. }
            | Op::Multi { src, .. } => src.base(),
            // All other operations are their own base
            _ => self.clone(),
        }
    }

    /// Get the underlying buffer UOp, walking through AFTER/MSELECT/MSTACK chains.
    ///
    /// Based on Tinygrad's `buf_uop` property (ops.py:601-606).
    /// This recursively unwraps AFTER chains to find the actual buffer.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// use morok_ir::UOp;
    ///
    /// // AFTER wrapping a buffer
    /// let buffer = UOp::new_buffer(...);
    /// let after = buffer.after(deps);
    ///
    /// // buf_uop() walks through AFTER to get the underlying buffer
    /// assert!(Arc::ptr_eq(&after.buf_uop(), &buffer));
    /// ```
    pub fn buf_uop(self: &Arc<Self>) -> Arc<Self> {
        match self.op() {
            Op::Buffer { .. } => self.clone(),
            Op::MSelect { buffer, .. } => buffer.buf_uop(),
            Op::MStack { buffers } if !buffers.is_empty() => buffers[0].buf_uop(),
            Op::After { passthrough, .. } => passthrough.buf_uop(),
            _ => {
                // For other ops, check if base is AFTER
                let base = self.base();
                if matches!(base.op(), Op::After { .. }) { base.buf_uop() } else { self.clone() }
            }
        }
    }

    /// Topological sort of the computation graph.
    ///
    /// Returns nodes in an order where all dependencies come before their dependents.
    pub fn toposort(self: &Arc<Self>) -> Vec<Arc<Self>> {
        let mut visited = HashSet::new();
        let mut result = Vec::new();
        let mut stack = vec![(self.clone(), false)];

        while let Some((node, processed)) = stack.pop() {
            let ptr = Arc::as_ptr(&node);

            if visited.contains(&ptr) {
                continue;
            }

            if processed {
                visited.insert(ptr);
                result.push(node);
            } else {
                stack.push((node.clone(), true));

                // Use for_each_child for zero-allocation traversal
                let mut children = Vec::new();
                node.op.map_child(|child| {
                    if !visited.contains(&Arc::as_ptr(child)) {
                        children.push(child.clone());
                    }
                });

                // Push in reverse order for proper traversal
                for child in children.into_iter().rev() {
                    stack.push((child, false));
                }
            }
        }

        result
    }

    /// Topological sort with gate function (filtered toposort).
    ///
    /// Only traverses nodes for which `gate(node)` returns true.
    /// Nodes for which gate returns false are excluded from the
    /// traversal entirely (along with their ancestors).
    ///
    /// This is a key optimization for cached property computation,
    /// allowing us to skip nodes that already have a property cached.
    ///
    /// # Performance
    ///
    /// For a graph with 10,000 nodes where 9,900 already have a cached property:
    /// - **Full toposort**: 10,000 nodes visited
    /// - **Filtered toposort**: 100 nodes visited
    /// - **Speedup**: 100x
    ///
    /// # Example
    ///
    /// ```ignore
    /// // Only process nodes that don't have shape cached
    /// let uncached = uop.toposort_filtered(|node| {
    ///     node.shape_cache.get().is_none()
    /// });
    /// ```
    pub fn toposort_filtered<F>(self: &Arc<Self>, gate: F) -> Vec<Arc<Self>>
    where
        F: Fn(&Arc<UOp>) -> bool,
    {
        let mut visited = HashSet::new();
        let mut result = Vec::new();
        let mut stack = vec![(self.clone(), false)];

        while let Some((node, processed)) = stack.pop() {
            let ptr = Arc::as_ptr(&node);

            if visited.contains(&ptr) {
                continue;
            }

            if processed {
                visited.insert(ptr);
                result.push(node);
            } else {
                // Key optimization: only traverse nodes that pass the gate
                if gate(&node) {
                    stack.push((node.clone(), true));

                    let mut children = Vec::new();
                    node.op.map_child(|child| {
                        if !visited.contains(&Arc::as_ptr(child)) {
                            children.push(child.clone());
                        }
                    });

                    // Push in reverse order for proper traversal
                    for child in children.into_iter().rev() {
                        stack.push((child, false));
                    }
                }
            }
        }

        result
    }

    /// Check if any node in the backward slice satisfies a predicate.
    ///
    /// Early-exit DFS — returns `true` as soon as a matching node is found,
    /// without building the full toposort Vec. Use this instead of
    /// `toposort().iter().any(pred)` when you only need an existential check.
    pub fn any_in_subtree<F>(self: &Arc<Self>, pred: F) -> bool
    where
        F: Fn(&Arc<UOp>) -> bool,
    {
        let mut visited = HashSet::new();
        let mut stack = vec![self.clone()];
        while let Some(node) = stack.pop() {
            if !visited.insert(Arc::as_ptr(&node)) {
                continue;
            }
            if pred(&node) {
                return true;
            }
            node.op.map_child(|child| {
                if !visited.contains(&Arc::as_ptr(child)) {
                    stack.push(child.clone());
                }
            });
        }
        false
    }

    /// Collect all nodes in the backward slice that match a predicate.
    ///
    /// DFS collecting matches — cheaper than `toposort().iter().filter(pred).collect()`
    /// when you don't need topological ordering.
    pub fn collect_in_subtree<F>(self: &Arc<Self>, pred: F) -> Vec<Arc<UOp>>
    where
        F: Fn(&Arc<UOp>) -> bool,
    {
        let mut visited = HashSet::new();
        let mut stack = vec![self.clone()];
        let mut result = Vec::new();
        while let Some(node) = stack.pop() {
            if !visited.insert(Arc::as_ptr(&node)) {
                continue;
            }
            if pred(&node) {
                result.push(node.clone());
            }
            node.op.map_child(|child| {
                if !visited.contains(&Arc::as_ptr(child)) {
                    stack.push(child.clone());
                }
            });
        }
        result
    }

    /// Count unique nodes in the DAG rooted at this UOp.
    ///
    /// Much cheaper than `toposort().len()` — no result Vec, no ordering.
    /// Uses pointer-based visited set for O(1) identity checks.
    pub fn node_count(self: &Arc<Self>) -> usize {
        let mut visited = HashSet::new();
        let mut stack = vec![self.clone()];
        while let Some(node) = stack.pop() {
            if !visited.insert(Arc::as_ptr(&node)) {
                continue;
            }
            node.op.map_child(|child| {
                if !visited.contains(&Arc::as_ptr(child)) {
                    stack.push(child.clone());
                }
            });
        }
        visited.len()
    }

    /// O(1) cached check: does this node or any of its sources contain an INDEX op?
    ///
    /// Computed lazily and cached. Each node checks itself and its direct sources'
    /// cached values, so the total cost across the graph is O(N).
    pub fn has_index_in_sources(self: &Arc<Self>) -> bool {
        *self.has_index_in_sources_cache.get_or_init(|| {
            if matches!(self.op, Op::Index { .. }) {
                return true;
            }
            let mut result = false;
            self.op.map_child(|child| {
                if child.has_index_in_sources() {
                    result = true;
                }
            });
            result
        })
    }

    /// Render this UOp and its sources as a compact ASCII tree.
    ///
    /// Shared nodes (appearing multiple times due to hash-consing) are shown
    /// as back-references: `[id] → (see above)`
    ///
    /// # Example Output
    ///
    /// ```text
    /// [42] STORE : Void
    /// ├── [10] PARAM(0) : Ptr<Float32> shape=[4]
    /// ├── [35] INDEX : Ptr<Float32> shape=[4]
    /// │   ├── [10] → (see above)
    /// │   └── [30] RANGE(0, Reduce) : Index
    /// │       └── [5] CONST(Int(4)) : Index
    /// └── [40] REDUCE(Add) : Float32 shape=[]
    ///     └── [35] → (see above)
    /// ```
    pub fn tree(self: &Arc<Self>) -> String {
        crate::uop::tree::render_tree_compact(self)
    }

    /// Render this UOp and its sources as a full ASCII tree.
    ///
    /// Shared nodes are expanded every time they appear (verbose but complete).
    /// Use this when you need to see the full subtree at every occurrence.
    pub fn tree_full(self: &Arc<Self>) -> String {
        crate::uop::tree::render_tree_full(self)
    }

    /// Get all RANGE operations in this UOp's computation graph.
    ///
    /// Lazily computed and cached. Useful for rangeify pass to track loop variables.
    pub fn ranges(self: &Arc<Self>) -> &Vec<Arc<Self>> {
        use crate::uop::cached_property::CachedProperty;
        use crate::uop::properties::RangesProperty;
        RangesProperty::get(self)
    }

    /// Get the RANGE operations that are in scope at this UOp.
    ///
    /// Returns only the ranges that are currently "active" (not yet ended).
    /// This is computed by:
    /// 1. Merging ranges from all source operations
    /// 2. Removing ranges that are ended by this operation
    /// 3. Adding self if this is a RANGE operation
    ///
    /// Based on Tinygrad's `ranges` property (ops.py:318-320) and
    /// `_ranges` recursive property (ops.py:302-315).
    ///
    /// # Returns
    ///
    /// A HashSet of RANGE UOps that are in scope at this point in the graph.
    /// The result is cached for performance.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// use morok_ir::{UOp, AxisType};
    ///
    /// // A simple computation inside a range
    /// let range = UOp::range(end, 0, AxisType::Loop);
    /// let value = UOp::const_(...);
    /// let end_op = value.end(vec![range.clone()]);
    ///
    /// // Value has range in scope
    /// assert!(value.in_scope_ranges().contains(&range));
    ///
    /// // After END, range is no longer in scope
    /// assert!(!end_op.in_scope_ranges().contains(&range));
    /// ```
    #[allow(clippy::mutable_key_type)]
    pub fn in_scope_ranges(self: &Arc<Self>) -> &HashSet<UOpKey> {
        use crate::uop::cached_property::CachedProperty;
        use crate::uop::properties::InScopeRangesProperty;
        InScopeRangesProperty::get(self)
    }

    /// Check if all in-scope ranges at this UOp have the given AxisType.
    ///
    /// Returns true if the in-scope ranges set is empty or all ranges
    /// match the specified axis type.
    ///
    /// # Use Cases
    ///
    /// - `all_in_scope_ranges_are(AxisType::Outer)` - Used in split_store
    ///   to determine if we're at a kernel boundary
    ///
    /// # Examples
    ///
    /// ```ignore
    /// use morok_ir::{UOp, AxisType};
    ///
    /// // At kernel boundary: only OUTER ranges in scope
    /// assert!(uop.all_in_scope_ranges_are(AxisType::Outer));
    ///
    /// // Inside kernel: has non-OUTER ranges
    /// assert!(!uop.all_in_scope_ranges_are(AxisType::Outer));
    /// ```
    #[allow(clippy::mutable_key_type)]
    pub fn all_in_scope_ranges_are(self: &Arc<Self>, axis_type: AxisType) -> bool {
        use crate::Op;

        let ranges = self.in_scope_ranges();

        // Empty scope means we're at the top level (treat as all OUTER)
        if ranges.is_empty() {
            return true;
        }

        ranges.iter().all(|r| match r.0.op() {
            Op::Range { axis_type: at, .. } => *at == axis_type,
            _ => false, // Should never happen
        })
    }

    /// Check if any in-scope range is NOT of the given AxisType.
    ///
    /// Inverse of `all_in_scope_ranges_are`. Useful for Tinygrad-style
    /// filtering: "skip if any range is not OUTER".
    ///
    /// # Examples
    ///
    /// ```ignore
    /// use morok_ir::{UOp, AxisType};
    ///
    /// // Has non-OUTER ranges: should skip in split_store
    /// if uop.has_non_outer_ranges() {
    ///     return None;  // Don't split here
    /// }
    /// ```
    pub fn has_non_outer_ranges(self: &Arc<Self>) -> bool {
        !self.all_in_scope_ranges_are(AxisType::Outer)
    }

    /// Build a consumer map for this UOp's computation graph.
    ///
    /// Returns a HashMap where each UOp maps to the list of UOps that consume it.
    /// Useful for reverse traversal and dependency analysis.
    #[allow(clippy::mutable_key_type)]
    pub fn get_consumer_map(self: &Arc<Self>) -> HashMap<UOpKey, Vec<Arc<Self>>> {
        let mut consumer_map: HashMap<UOpKey, Vec<Arc<Self>>> = HashMap::new();

        for node in self.toposort() {
            node.op.map_child(|child| {
                consumer_map.entry(UOpKey(child.clone())).or_default().push(node.clone());
            });
        }

        consumer_map
    }

    /// Reverse topological sort of the computation graph.
    ///
    /// Returns nodes in bottom-up order (leaves first, root last).
    /// Requires a consumer map to traverse from leaves to roots.
    #[allow(clippy::mutable_key_type)]
    pub fn reverse_toposort(self: &Arc<Self>, consumer_map: &HashMap<UOpKey, Vec<Arc<Self>>>) -> Vec<Arc<Self>> {
        let mut visited = HashMap::new(); // Use HashMap to track visited by ID
        let mut result = Vec::new();
        let mut stack = vec![(self.clone(), false)];

        while let Some((node, processed)) = stack.pop() {
            if visited.contains_key(&node.id) {
                continue;
            }

            if processed {
                visited.insert(node.id, ());
                result.push(node);
            } else {
                stack.push((node.clone(), true));

                // Visit consumers (nodes that depend on this node)
                if let Some(consumers) = consumer_map.get(&UOpKey(node.clone())) {
                    for consumer in consumers {
                        if !visited.contains_key(&consumer.id) {
                            stack.push((consumer.clone(), false));
                        }
                    }
                }
            }
        }

        result
    }

    /// Replace UOps in the computation graph according to a substitution map.
    ///
    /// Delegates to `graph_rewrite_bottom_up` with a wildcard pattern that looks up
    /// each node in the map — exactly like Tinygrad's `substitute`. The rewrite engine
    /// provides O(n) memoization via its result cache.
    #[allow(clippy::mutable_key_type)]
    pub fn substitute(self: &Arc<Self>, map: &HashMap<UOpKey, Arc<Self>>) -> Arc<Self> {
        if map.is_empty() {
            return self.clone();
        }
        let matcher = SubstituteMatcher(map);
        crate::rewrite::graph_rewrite_bottom_up(&matcher, self.clone(), &mut ())
    }

    /// Replace UOps with range-gated substitution (Tinygrad: `extra_pm=pm_gate_substitute`).
    ///
    /// Like `substitute`, but skips subtrees whose `in_scope_ranges()` don't contain
    /// any of the substitution keys. This prevents substituting ranges in subexpressions
    /// that don't reference them, matching Tinygrad's `gate_substitute` behavior.
    #[allow(clippy::mutable_key_type)]
    pub fn substitute_gated(self: &Arc<Self>, map: &HashMap<UOpKey, Arc<Self>>) -> Arc<Self> {
        if map.is_empty() {
            return self.clone();
        }
        let range_keys: HashSet<UOpKey> = map.keys().cloned().collect();
        let matcher = SubstituteGatedMatcher { map, range_keys: &range_keys };
        crate::rewrite::graph_rewrite_bottom_up(&matcher, self.clone(), &mut ())
    }

    /// Reconstruct this UOp with new sources.
    ///
    /// Creates a new UOp with the same operation type and dtype, but with the provided
    /// sources replacing the original ones. Hash consing ensures that if an identical
    /// UOp already exists, it will be reused.
    ///
    /// This is used by the graph rewrite engine when sources have been rewritten.
    ///
    /// # Panics
    ///
    /// Panics if the number of sources doesn't match the operation's arity.
    ///
    /// # Examples
    ///
    /// ```ignore
    /// // Original: a + b
    /// let add = UOp::add(a.clone(), b.clone());
    ///
    /// // Rewrite sources: a' + b'
    /// let new_add = add.with_sources(vec![a_prime, b_prime]);
    /// ```
    pub fn with_sources(self: &Arc<Self>, new_srcs: Vec<Arc<Self>>) -> Arc<Self> {
        use smallvec::SmallVec;

        // Helper to get nth source
        let src = |n: usize| new_srcs[n].clone();

        let new_op = match &self.op {
            // Nullary operations - no sources
            Op::Const(_)
            | Op::Unique(_)
            | Op::Device(_)
            | Op::Noop
            | Op::Invalid
            | Op::DefineLocal(_)
            | Op::VConst { .. }
            | Op::DefineVar { .. }
            | Op::DefineReg { .. } => {
                assert_eq!(new_srcs.len(), 0, "Nullary op should have no sources");
                return self.clone(); // No sources to replace
            }

            // Unary operations
            Op::Unary(op_type, _) => {
                assert_eq!(new_srcs.len(), 1);
                Op::Unary(*op_type, src(0))
            }

            // Binary operations
            Op::Binary(op_type, _, _) => {
                assert_eq!(new_srcs.len(), 2);
                Op::Binary(*op_type, src(0), src(1))
            }

            // Ternary operations
            Op::Ternary(op_type, _, _, _) => {
                assert_eq!(new_srcs.len(), 3);
                Op::Ternary(*op_type, src(0), src(1), src(2))
            }

            // Type operations
            Op::Cast { dtype, .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::Cast { src: src(0), dtype: dtype.clone() }
            }
            Op::BitCast { dtype, .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::BitCast { src: src(0), dtype: dtype.clone() }
            }

            // Special operations
            Op::MSelect { device_index, .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::MSelect { buffer: src(0), device_index: *device_index }
            }
            Op::Special { name, .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::Special { end: src(0), name: name.clone() }
            }

            // Buffer operations
            Op::Buffer { size, .. } => {
                assert_eq!(new_srcs.len(), 2);
                Op::Buffer { unique: src(0), device: src(1), size: *size }
            }
            Op::Param { slot, size, device } => {
                if device.is_some() {
                    assert_eq!(new_srcs.len(), 1);
                    Op::Param { slot: *slot, size: *size, device: Some(src(0)) }
                } else {
                    assert_eq!(new_srcs.len(), 0);
                    return self.clone();
                }
            }
            Op::BufferView { size, offset, .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::BufferView { buffer: src(0), size: *size, offset: *offset }
            }
            Op::Bufferize { opts, .. } => {
                assert!(!new_srcs.is_empty());
                Op::Bufferize { compute: src(0), ranges: new_srcs[1..].iter().cloned().collect(), opts: opts.clone() }
            }
            Op::Index { gate, .. } => {
                assert!(!new_srcs.is_empty());
                // First source is buffer, rest are indices, last might be gate
                let buffer = src(0);
                let (indices, gate_new) = if gate.is_some() && new_srcs.len() >= 2 {
                    let gate_src = new_srcs.last().unwrap().clone();
                    let indices: SmallVec<[Arc<Self>; 4]> = new_srcs[1..new_srcs.len() - 1].iter().cloned().collect();
                    (indices, Some(gate_src))
                } else {
                    let indices: SmallVec<[Arc<Self>; 4]> = new_srcs[1..].iter().cloned().collect();
                    (indices, None)
                };
                Op::Index { buffer, indices, gate: gate_new }
            }
            Op::PointerIndex { .. } => {
                assert_eq!(new_srcs.len(), 2);
                Op::PointerIndex { ptr: src(0), offset: src(1) }
            }
            Op::Copy { .. } => {
                assert_eq!(new_srcs.len(), 2);
                Op::Copy { src: src(0), device: src(1) }
            }
            Op::MStack { .. } => Op::MStack { buffers: new_srcs.iter().cloned().collect() },

            // Movement operations
            Op::Reshape { .. } => {
                assert_eq!(new_srcs.len(), 2);
                Op::Reshape { src: src(0), new_shape: src(1) }
            }
            Op::Permute { axes, .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::Permute { src: src(0), axes: axes.clone() }
            }
            Op::Expand { .. } => {
                assert_eq!(new_srcs.len(), 2);
                Op::Expand { src: src(0), new_shape: src(1) }
            }
            Op::Pad { .. } => {
                assert_eq!(new_srcs.len(), 3);
                Op::Pad { src: src(0), begin_pads: src(1), end_pads: src(2) }
            }
            Op::Shrink { .. } => {
                assert_eq!(new_srcs.len(), 3);
                Op::Shrink { src: src(0), begins: src(1), ends: src(2) }
            }
            Op::Flip { axes, .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::Flip { src: src(0), axes: axes.clone() }
            }
            Op::Multi { axis, .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::Multi { src: src(0), axis: *axis }
            }

            // Reduction operations
            Op::ReduceAxis { reduce_op, axes, .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::ReduceAxis { src: src(0), reduce_op: *reduce_op, axes: axes.clone() }
            }
            Op::Reduce { reduce_op, .. } => {
                assert!(!new_srcs.is_empty());
                Op::Reduce { src: src(0), ranges: new_srcs[1..].iter().cloned().collect(), reduce_op: *reduce_op }
            }
            Op::AllReduce { reduce_op, .. } => {
                assert_eq!(new_srcs.len(), 2);
                Op::AllReduce { src: src(0), device: src(1), reduce_op: *reduce_op }
            }

            // Control flow operations
            Op::If { .. } => {
                assert!(!new_srcs.is_empty());
                Op::If { condition: src(0), body: new_srcs[1..].iter().cloned().collect() }
            }
            Op::EndIf { .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::EndIf { if_op: src(0) }
            }
            Op::Range { axis_id, axis_type, .. } => {
                assert!(!new_srcs.is_empty());
                Op::Range {
                    end: src(0),
                    axis_id: *axis_id,
                    axis_type: *axis_type,
                    deps: new_srcs[1..].iter().cloned().collect(),
                }
            }
            Op::End { .. } => {
                assert!(!new_srcs.is_empty());
                Op::End { computation: src(0), ranges: new_srcs[1..].iter().cloned().collect() }
            }
            Op::Barrier { .. } => {
                assert!(!new_srcs.is_empty());
                Op::Barrier { src: src(0), deps: new_srcs[1..].iter().cloned().collect() }
            }

            // Vector operations — recompute dtype from new elements when element
            // dtype category changed (e.g. Scalar → Ptr during rewrite reconstruction).
            // Preserving old dtype is wrong when DEFINE_LOCAL → AFTER(Ptr) changes
            // element types from Scalar to Ptr, causing pm_add_loads infinite loops.
            Op::Vectorize { .. } => {
                let elements: SmallVec<[Arc<Self>; 4]> = new_srcs.iter().cloned().collect();
                let elem_dtype = elements[0].dtype();
                let new_dtype = match elem_dtype {
                    DType::Scalar(_) | DType::Ptr { .. } => elem_dtype.vec(elements.len()),
                    _ => self.dtype.clone(),
                };
                return Self::new(Op::Vectorize { elements }, new_dtype);
            }
            Op::Gep { indices, .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::Gep { vector: src(0), indices: indices.clone() }
            }
            Op::Cat { .. } => Op::Cat { sources: new_srcs.iter().cloned().collect() },
            Op::PtrCat { .. } => Op::PtrCat { sources: new_srcs.iter().cloned().collect() },

            // Symbolic/Define operations
            Op::Bind { .. } => {
                assert_eq!(new_srcs.len(), 2);
                Op::Bind { var: src(0), value: src(1) }
            }

            // Advanced operations
            Op::Wmma { metadata, .. } => {
                assert_eq!(new_srcs.len(), 3);
                Op::Wmma { a: src(0), b: src(1), c: src(2), metadata: metadata.clone() }
            }
            Op::Contract { upcast_ranges, .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::Contract { src: src(0), upcast_ranges: upcast_ranges.clone() }
            }
            Op::Unroll { unroll_axes, .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::Unroll { src: src(0), unroll_axes: unroll_axes.clone() }
            }
            Op::Kernel { .. } => {
                assert!(!new_srcs.is_empty());
                Op::Kernel {
                    sources: new_srcs[..new_srcs.len() - 1].iter().cloned().collect(),
                    ast: src(new_srcs.len() - 1),
                }
            }
            Op::Assign { .. } => {
                assert!(new_srcs.len() >= 2 && new_srcs.len() <= 3, "Assign requires 2-3 sources");
                Op::Assign {
                    target: src(0),
                    value: src(1),
                    movement_ops: if new_srcs.len() > 2 { Some(src(2)) } else { None },
                }
            }
            Op::Detach { .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::Detach { src: src(0) }
            }
            Op::Contiguous { opts, .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::Contiguous { src: src(0), opts: opts.clone() }
            }
            Op::ContiguousBackward { .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::ContiguousBackward { src: src(0) }
            }
            Op::After { .. } => {
                assert!(!new_srcs.is_empty());
                let passthrough = src(0);
                // Validate: AFTER passthrough must not be control flow (Tinygrad semantics)
                debug_assert!(
                    !matches!(passthrough.op(), Op::Range { .. } | Op::End { .. }),
                    "reconstruct_sources: AFTER passthrough is {:?} (id={}), violates Tinygrad semantics",
                    passthrough.op(),
                    passthrough.id
                );
                Op::After { passthrough, deps: new_srcs[1..].iter().cloned().collect() }
            }
            Op::Precast { .. } => {
                assert_eq!(new_srcs.len(), 1);
                Op::Precast { src: src(0) }
            }
            Op::Custom { code, .. } => Op::Custom { deps: new_srcs.iter().cloned().collect(), code: code.clone() },
            Op::CustomI { code, .. } => Op::CustomI { deps: new_srcs.iter().cloned().collect(), code: code.clone() },

            // Memory operations
            Op::Load { alt, .. } => {
                // Load has 2-3 sources: buffer, index, and optionally alt
                assert!(new_srcs.len() >= 2 && new_srcs.len() <= 3, "Load requires 2-3 sources");
                let new_alt = if new_srcs.len() == 3 { Some(src(2)) } else { alt.clone() };
                Op::Load { buffer: src(0), index: src(1), alt: new_alt }
            }
            Op::Store { .. } => {
                assert!(new_srcs.len() >= 2, "Store requires at least 2 sources (index, value)");
                Op::Store { index: src(0), value: src(1), ranges: new_srcs[2..].iter().cloned().collect() }
            }

            // Graph organization
            Op::Sink { .. } => Op::Sink { sources: new_srcs.iter().cloned().collect() },
            Op::Group { .. } => Op::Group { sources: new_srcs.iter().cloned().collect() },
        };

        // Preserve original dtype and tag (Tinygrad ops.py:1256: preserves tag through source reconstruction)
        Self::new_tagged(new_op, self.dtype.clone(), self.tag.clone())
    }
}

#[bon]
impl UOp {
    /// Create a modified copy with optional field overrides.
    ///
    /// Enables concise pattern implementations by allowing selective field modification.
    /// Returns `self.clone()` if nothing changed (optimization for hash consing).
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let new_load = load.replace().dtype(new_dtype).src(new_sources).call();
    /// let dtype_only = load.replace().dtype(new_dtype).call();
    /// ```
    #[builder]
    pub fn replace(self: &Arc<Self>, dtype: Option<DType>, src: Option<Vec<Arc<Self>>>) -> Arc<Self> {
        let new_dtype = dtype.unwrap_or_else(|| self.dtype());
        let new_sources = src.unwrap_or_else(|| self.op().sources().to_vec());

        // Short-circuit if nothing changed
        let old_sources = self.op().sources();
        let sources_unchanged = new_sources.len() == old_sources.len()
            && new_sources.iter().zip(old_sources.iter()).all(|(a, b)| Arc::ptr_eq(a, b));

        if new_dtype == self.dtype() && sources_unchanged {
            return self.clone();
        }

        self.with_sources(new_sources).with_dtype(new_dtype)
    }
}

impl Clone for UOp {
    fn clone(&self) -> Self {
        Self {
            id: self.id,
            op: self.op.clone(),
            dtype: self.dtype.clone(),
            content_hash: self.content_hash,
            tag: self.tag.clone(),
            shape_cache: std::sync::OnceLock::new(),
            ranges_cache: std::sync::OnceLock::new(),
            in_scope_ranges_cache: std::sync::OnceLock::new(),
            vmin_vmax_cache: std::sync::OnceLock::new(),
            sound_vmin_vmax_cache: std::sync::OnceLock::new(),
            has_index_in_sources_cache: std::sync::OnceLock::new(),
            backward_slice_cache: std::sync::OnceLock::new(),
            metadata: self.metadata.clone(),
        }
    }
}

/// Trait for converting scalar values into UOps.
///
/// This allows operator overloading to work with mixed scalar/UOp operands.
/// For example: `uop + 5.0` or `5.0 + uop`.
pub trait IntoUOp {
    fn into_uop(self, dtype: DType) -> Arc<UOp>;
}

impl IntoUOp for ConstValue {
    fn into_uop(self, dtype: DType) -> Arc<UOp> {
        UOp::const_(dtype, self)
    }
}

impl IntoUOp for f32 {
    fn into_uop(self, dtype: DType) -> Arc<UOp> {
        UOp::const_(dtype, ConstValue::Float(self as f64))
    }
}

impl IntoUOp for f64 {
    fn into_uop(self, dtype: DType) -> Arc<UOp> {
        UOp::const_(dtype, ConstValue::Float(self))
    }
}

impl IntoUOp for i32 {
    fn into_uop(self, dtype: DType) -> Arc<UOp> {
        UOp::const_(dtype, ConstValue::Int(self as i64))
    }
}

impl IntoUOp for i64 {
    fn into_uop(self, dtype: DType) -> Arc<UOp> {
        UOp::const_(dtype, ConstValue::Int(self))
    }
}

impl IntoUOp for u32 {
    fn into_uop(self, dtype: DType) -> Arc<UOp> {
        UOp::const_(dtype, ConstValue::UInt(self as u64))
    }
}

impl IntoUOp for u64 {
    fn into_uop(self, dtype: DType) -> Arc<UOp> {
        UOp::const_(dtype, ConstValue::UInt(self))
    }
}

impl IntoUOp for bool {
    fn into_uop(self, dtype: DType) -> Arc<UOp> {
        UOp::const_(dtype, ConstValue::Bool(self))
    }
}