More disable custom interp (#1029)

* More disable custom interp

* err

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* materialize traced array

* fix

* fix

* fix

* oneto

* fixup

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* up docs

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: wsmoses

docs/src/api/api.md

@@ -50,5 +50,5 @@ Reactant.addressable_devices
 ## Internal utils
 
 ```@docs
-Reactant.TracedUtils.materialize_traced_array
+ReactantCore.materialize_traced_array
 ```

lib/ReactantCore/src/ReactantCore.jl

@@ -516,4 +516,12 @@ function error_if_any_control_flow(expr)
     end
 end
 
+"""
+    materialize_traced_array(AbstractArray{<:TracedRNumber})::TracedRArray
+
+Given an AbstractArray{TracedRNumber}, return or create an equivalent TracedRArray.
+
+"""
+function materialize_traced_array end
+
 end

src/Reactant.jl

@@ -1,6 +1,7 @@
 module Reactant
 
-using ReactantCore: ReactantCore, @trace, within_compile, MissingTracedValue
+using ReactantCore:
+    ReactantCore, @trace, within_compile, MissingTracedValue, materialize_traced_array
 
 using LinearAlgebra: LinearAlgebra
 using Random: Random, AbstractRNG

src/TracedRNumber.jl

@@ -575,6 +575,31 @@ function TracedUnitRange{T}(r::AbstractUnitRange) where {T<:Real}
 end
 TracedUnitRange(r::AbstractUnitRange) = TracedUnitRange(first(r), last(r))
 
+@inline function Base.length(r::TracedUnitRange{TracedRNumber{T}}) where {T}
+    start, stop = first(r), last(r)
+    a = Base.oneunit(Base.zero(stop) - Base.zero(start))
+    if a isa Signed
+        # Signed are allowed to go negative
+        Ops.select(stop >= start, a + stop - start, a)
+    else
+        Ops.select(stop >= start, a + stop - start, zero(a))
+    end
+end
+
+function Base._reshape(v::TracedUnitRange, dims::Dims{1})
+    Base.require_one_based_indexing(v)
+    len = dims[1]
+    # TODO support errors
+    # len == length(v) || Base._throw_dmrs(length(v), "length", len)
+    return v
+end
+function Base._reshape(parent::TracedUnitRange, dims::Dims)
+    n = length(parent)
+    # TODO support errors
+    # prod(dims) == n || Base._throw_dmrs(n, "size", dims)
+    return Base.__reshape((parent, IndexStyle(parent)), dims)
+end
+
 AbstractUnitRange{T}(r::TracedUnitRange) where {T} = TracedUnitRange{T}(r)
 
 struct TracedStepRangeLen{T,R,S,L} <: AbstractRange{T}

src/TracedUtils.jl

@@ -14,48 +14,47 @@ using ..Reactant:
     OrderedIdDict,
     ReactantPrimitive,
     Ops
-using ReactantCore: MissingTracedValue, is_traced
+using ReactantCore: ReactantCore
+using ReactantCore: MissingTracedValue, is_traced, materialize_traced_array
 using Functors: Functors
 
-"""
-    materialize_traced_array(AbstractArray{<:TracedRNumber})::TracedRArray
+ReactantCore.materialize_traced_array(x::AbstractArray) = x
 
-Given an AbstractArray{TracedRNumber}, return or create an equivalent TracedRArray.
+ReactantCore.materialize_traced_array(x::TracedRArray) = x
 
-"""
-materialize_traced_array(x::AbstractArray) = x
+ReactantCore.materialize_traced_array(x::AnyTracedRArray) = x[axes(x)...]
 
-materialize_traced_array(x::TracedRArray) = x
-
-materialize_traced_array(x::AnyTracedRArray) = x[axes(x)...]
-
-function materialize_traced_array(x::AbstractRange{<:TracedRNumber})
+function ReactantCore.materialize_traced_array(x::AbstractRange)
     return Reactant.aos_to_soa(collect(x))
 end
 
-function materialize_traced_array(x::UnitRange{<:TracedRNumber})
+function ReactantCore.materialize_traced_array(x::Base.OneTo)
+    return Ops.iota(Reactant.unwrapped_eltype(x), [length(x)]; iota_dimension=1)
+end
+
+function ReactantCore.materialize_traced_array(x::UnitRange)
     return Ops.add(
         Ops.iota(Reactant.unwrapped_eltype(x), [length(x)]; iota_dimension=1),
         Ops.fill(first(x), [length(x)]),
     )
 end
 
-function materialize_traced_array(x::SubArray{TracedRNumber{T}}) where {T}
+function ReactantCore.materialize_traced_array(x::SubArray)
     z = SubArray(materialize_traced_array(parent(x)), x.indices)
     return z[axes(z)...]
 end
 
-function materialize_traced_array(x::Base.ReshapedArray{TracedRNumber{T}}) where {T}
+function ReactantCore.materialize_traced_array(x::Base.ReshapedArray)
     return Ops.reshape(materialize_traced_array(parent(x)), size(x)...)
 end
 
-function materialize_traced_array(
-    x::PermutedDimsArray{TracedRNumber{T},N,perm}
-) where {T,N,perm}
+function ReactantCore.materialize_traced_array(
+    x::PermutedDimsArray{<:Any,<:Any,perm}
+) where {perm}
     return permutedims(materialize_traced_array(parent(x)), perm)
 end
 
-function materialize_traced_array(x::AbstractArray{TracedRNumber{T}}) where {T}
+function ReactantCore.materialize_traced_array(x::AbstractArray{TracedRNumber{T}}) where {T}
     return Reactant.aos_to_soa(x)
 end
 

src/stdlibs/LinearAlgebra.jl

@@ -12,34 +12,37 @@ using ..Reactant:
     Ops,
     MLIR
 
-using ..TracedUtils: TracedUtils, get_mlir_data, materialize_traced_array, set_mlir_data!
+using ReactantCore: ReactantCore
+using ReactantCore: materialize_traced_array
+
+using ..TracedUtils: TracedUtils, get_mlir_data, set_mlir_data!
 
 using LinearAlgebra
 
 # Various Wrapper Arrays defined in LinearAlgebra
-function TracedUtils.materialize_traced_array(
+function ReactantCore.materialize_traced_array(
     x::Transpose{TracedRNumber{T},<:AnyTracedRArray}
 ) where {T}
-    px = TracedUtils.materialize_traced_array(parent(x))
+    px = materialize_traced_array(parent(x))
     A = ndims(px) == 1 ? reshape(px, :, 1) : px
     return permutedims(A, (2, 1))
 end
 
-function TracedUtils.materialize_traced_array(
+function ReactantCore.materialize_traced_array(
     x::Adjoint{TracedRNumber{T},<:AnyTracedRArray}
 ) where {T}
     return Ops.conj(
         materialize_traced_array(transpose(materialize_traced_array(parent(x))))
     )
 end
 
-function TracedUtils.materialize_traced_array(
+function ReactantCore.materialize_traced_array(
     x::Diagonal{TracedRNumber{T},<:AnyTracedRVector}
 ) where {T}
     return diagm(materialize_traced_array(parent(x)))
 end
 
-function TracedUtils.materialize_traced_array(
+function ReactantCore.materialize_traced_array(
     x::Tridiagonal{TracedRNumber{T},<:AnyTracedRVector}
 ) where {T}
     return diagm(-1 => x.dl, 0 => x.d, 1 => x.du)
@@ -48,22 +51,22 @@ end
 for (AT, comp) in ((:LowerTriangular, "GE"), (:UpperTriangular, "LE"))
     uAT = Symbol(:Unit, AT)
     @eval begin
-        function TracedUtils.materialize_traced_array(
+        function ReactantCore.materialize_traced_array(
             x::$(AT){TracedRNumber{T},<:AnyTracedRMatrix}
         ) where {T}
             m, n = size(x)
-            px = TracedUtils.materialize_traced_array(parent(x))
+            px = materialize_traced_array(parent(x))
             row_idxs = Ops.iota(Int, [m, n]; iota_dimension=1)
             col_idxs = Ops.iota(Int, [m, n]; iota_dimension=2)
             indicator = Ops.compare(row_idxs, col_idxs; comparison_direction=$(comp))
             return Ops.select(indicator, px, zero(px))
         end
 
-        function TracedUtils.materialize_traced_array(
+        function ReactantCore.materialize_traced_array(
             x::$(uAT){TracedRNumber{T},<:AnyTracedRMatrix}
         ) where {T}
             m, n = size(x)
-            px = TracedUtils.materialize_traced_array(parent(x))
+            px = materialize_traced_array(parent(x))
             row_idxs = Ops.iota(Int, [m, n]; iota_dimension=1)
             col_idxs = Ops.iota(Int, [m, n]; iota_dimension=2)
             nondiag_indicator = Ops.compare(row_idxs, col_idxs; comparison_direction="NE")
@@ -73,7 +76,7 @@ for (AT, comp) in ((:LowerTriangular, "GE"), (:UpperTriangular, "LE"))
     end
 end
 
-function TracedUtils.materialize_traced_array(
+function ReactantCore.materialize_traced_array(
     x::Symmetric{TracedRNumber{T},<:AnyTracedRMatrix}
 ) where {T}
     m, n = size(x)

src/utils.jl

@@ -179,7 +179,9 @@ function should_rewrite_call(@nospecialize(ft))
         ft <: typeof(Adapt.adapt_structure) ||
         ft <: typeof(Core.is_top_bit_set) ||
         ft <: typeof(Base.setindex_widen_up_to) ||
-        ft <: typeof(Base.typejoin)
+        ft <: typeof(Base.typejoin) ||
+        ft <: typeof(Base.argtype_decl) ||
+        ft <: typeof(Base.arg_decl_parts)
         return false
     end
 
@@ -188,7 +190,12 @@ function should_rewrite_call(@nospecialize(ft))
 end
 
 # by default, same as `should_rewrite_call`
-should_rewrite_invoke(@nospecialize(ft), @nospecialize(args)) = should_rewrite_call(ft)
+function should_rewrite_invoke(@nospecialize(ft), @nospecialize(args))
+    if ft <: typeof(repeat) && args == Tuple{String,Int64}
+        return false
+    end
+    return should_rewrite_call(ft)
+end
 
 # Avoid recursively interpreting into methods we define explicitly
 # as overloads, which we assume should handle the entirety of the
@@ -582,7 +589,16 @@ function call_with_reactant_generator(
         safe_print("ir", ir)
     end
 
-    if !is_reactant_method(mi::Core.MethodInstance) || guaranteed_error
+    mi = mi::Core.MethodInstance
+
+    if !(
+        is_reactant_method(mi) || (
+            mi.def.sig isa DataType &&
+            !should_rewrite_invoke(
+                mi.def.sig.parameters[1], Tuple{mi.def.sig.parameters[2:end]...}
+            )
+        )
+    ) || guaranteed_error
         ir, any_changed = rewrite_insts!(ir, interp, guaranteed_error)
     end