treat .= as syntactic sugar for broadcast! (#17510)

* treat .= as syntactic sugar for broadcast!

* tests

* optimized .= assignment of scalars and vector copies

* .= documentation

* fix show of .= ops

* .-= tests

* NEWS for .=

NEWS.md

@@ -10,8 +10,11 @@ New language features
  * Generators and comprehensions support filtering using `if` ([#550]) and nested
  iteration using multiple `for` keywords ([#4867]).
 
- * Broadcasting syntax: ``f.(args...)`` is equivalent to ``broadcast(f, args...)`` ([#15032]),
+ * Fused broadcasting syntax: ``f.(args...)`` is equivalent to ``broadcast(f, args...)`` ([#15032]),
  and nested `f.(g.(args...))` calls are fused into a single `broadcast` loop ([#17300]).
+ Similarly, the syntax `x .= ...` is equivalent to a `broadcast!(identity, x, ...)`
+ call and fuses with nested "dot" calls; also, `x .+= y` and similar is now
+ equivalent to `x .= x .+ y`, rather than `=` ([#17510]).
 
  * Macro expander functions are now generic, so macros can have multiple definitions
  (e.g. for different numbers of arguments, or optional arguments) ([#8846], [#9627]).
@@ -357,3 +360,4 @@ Deprecated or removed
 [#17393]: https://github.com/JuliaLang/julia/issues/17393
 [#17402]: https://github.com/JuliaLang/julia/issues/17402
 [#17404]: https://github.com/JuliaLang/julia/issues/17404
+[#17510]: https://github.com/JuliaLang/julia/issues/17510

base/broadcast.jl

@@ -15,6 +15,15 @@ export broadcast_getindex, broadcast_setindex!
 broadcast(f) = f()
 broadcast(f, x::Number...) = f(x...)
 
+# special cases for "X .= ..." (broadcast!) assignments
+broadcast!(::typeof(identity), X::AbstractArray, x::Number) = fill!(X, x)
+broadcast!(f, X::AbstractArray) = fill!(X, f())
+broadcast!(f, X::AbstractArray, x::Number...) = fill!(X, f(x...))
+function broadcast!{T,S,N}(::typeof(identity), x::AbstractArray{T,N}, y::AbstractArray{S,N})
+ check_broadcast_shape(size(x), size(y))
+ copy!(x, y)
+end
+
 ## Calculate the broadcast shape of the arguments, or error if incompatible
 # array inputs
 broadcast_shape() = ()

base/show.jl

@@ -408,8 +408,10 @@ show_unquoted(io::IO, ex, ::Int,::Int) = show(io, ex)
 const indent_width = 4
 const quoted_syms = Set{Symbol}([:(:),:(::),:(:=),:(=),:(==),:(!=),:(===),:(!==),:(=>),:(>=),:(<=)])
 const uni_ops = Set{Symbol}([:(+), :(-), :(!), :(¬), :(~), :(<:), :(>:), :(√), :(∛), :(∜)])
-const expr_infix_wide = Set{Symbol}([:(=), :(+=), :(-=), :(*=), :(/=), :(\=), :(&=),
- :(|=), :($=), :(>>>=), :(>>=), :(<<=), :(&&), :(||), :(<:), :(=>), :(÷=)])
+const expr_infix_wide = Set{Symbol}([
+ :(=), :(+=), :(-=), :(*=), :(/=), :(\=), :(^=), :(&=), :(|=), :(÷=), :(%=), :(>>>=), :(>>=), :(<<=),
+ :(.=), :(.+=), :(.-=), :(.*=), :(./=), :(.\=), :(.^=), :(.&=), :(.|=), :(.÷=), :(.%=), :(.>>>=), :(.>>=), :(.<<=),
+ :(&&), :(||), :(<:), :(=>), :($=)])
 const expr_infix = Set{Symbol}([:(:), :(->), Symbol("::")])
 const expr_infix_any = union(expr_infix, expr_infix_wide)
 const all_ops = union(quoted_syms, uni_ops, expr_infix_any)

doc/manual/arrays.rst

@@ -566,7 +566,7 @@ function elementwise:
  1.71056 0.847604
  1.73659 0.873631
 
-Elementwise operators such as ``.+`` and ``.*`` perform broadcasting if necessary. There is also a :func:`broadcast!` function to specify an explicit destination, and :func:`broadcast_getindex` and :func:`broadcast_setindex!` that broadcast the indices before indexing. Moreover, ``f.(args...)`` is equivalent to ``broadcast(f, args...)``, providing a convenient syntax to broadcast any function (:ref:`man-dot-vectorizing`:.).
+Elementwise operators such as ``.+`` and ``.*`` perform broadcasting if necessary. There is also a :func:`broadcast!` function to specify an explicit destination, and :func:`broadcast_getindex` and :func:`broadcast_setindex!` that broadcast the indices before indexing. Moreover, ``f.(args...)`` is equivalent to ``broadcast(f, args...)``, providing a convenient syntax to broadcast any function (:ref:`man-dot-vectorizing`:).
 
 Implementation
 --------------

doc/manual/functions.rst

@@ -652,9 +652,20 @@ the fusion stops as soon as a "non-dot" function is encountered; for example,
 in ``sin.(sort(cos.(X)))`` the ``sin`` and ``cos`` loops cannot be merged
 because of the intervening ``sort`` function.
 
+Finally, the maximum efficiency is typically achieved when the output
+array of a vectorized operation is *pre-allocated*, so that repeated
+calls do not allocate new arrays over and over again for the results
+(:ref:`man-preallocation`:). A convenient syntax for this is
+``X .= ...``, which is equivalent to ``broadcast!(identity, X, ...)``
+except that, as above, the ``broadcast!`` loop is fused with any nested
+"dot" calls. For example, ``X .= sin.(Y)`` is equivalent to
+``broadcast!(sin, X, Y)``, overwriting ``X`` with ``sin.(Y)`` in-place.
+
 (In future versions of Julia, operators like ``.*`` will also be handled with
 the same mechanism: they will be equivalent to ``broadcast`` calls and
-will be fused with other nested "dot" calls.)
+will be fused with other nested "dot" calls. ``x .+= y`` is equivalent
+to ``x .= x .+ y`` and will eventually result in a fused in-place assignment.
+Similarly for ``.*=`` etcetera.)
 
 Further Reading
 ---------------

doc/manual/performance-tips.rst

@@ -944,7 +944,10 @@ above, we could have passed a :class:`SubArray` rather than an :class:`Array`,
 had we so desired.
 
 Taken to its extreme, pre-allocation can make your code uglier, so
-performance measurements and some judgment may be required.
+performance measurements and some judgment may be required. However,
+for "vectorized" (element-wise) functions, the convenient syntax
+``x .= f.(y)`` can be used for in-place operations with fused loops
+and no temporary arrays (:ref:`dot-vectorizing`).
 
 
 Avoid string interpolation for I/O

src/julia-syntax.scm

@@ -1418,12 +1418,12 @@
  `(call ,(cadr e) ,(expand-forms a) ,(expand-forms b))))))
 
 ;; convert `a+=b` to `a=a+b`
-(define (expand-update-operator- op lhs rhs declT)
+(define (expand-update-operator- op op= lhs rhs declT)
  (let ((e (remove-argument-side-effects lhs)))
  `(block ,@(cdr e)
  ,(if (null? declT)
- `(= ,(car e) (call ,op ,(car e) ,rhs))
- `(= ,(car e) (call ,op (:: ,(car e) ,(car declT)) ,rhs))))))
+ `(,op= ,(car e) (call ,op ,(car e) ,rhs))
+ `(,op= ,(car e) (call ,op (:: ,(car e) ,(car declT)) ,rhs))))))
 
 (define (partially-expand-ref e)
  (let ((a (cadr e))
@@ -1443,31 +1443,32 @@
  ,@(append stmts stuff)
  (call getindex ,arr ,@new-idxs))))))
 
-(define (expand-update-operator op lhs rhs . declT)
+(define (expand-update-operator op op= lhs rhs . declT)
  (cond ((and (pair? lhs) (eq? (car lhs) 'ref))
  ;; expand indexing inside op= first, to remove "end" and ":"
  (let* ((ex (partially-expand-ref lhs))
  (stmts (butlast (cdr ex)))
  (refex (last (cdr ex)))
  (nuref `(ref ,(caddr refex) ,@(cdddr refex))))
  `(block ,@stmts
- ,(expand-update-operator- op nuref rhs declT))))
+ ,(expand-update-operator- op op= nuref rhs declT))))
  ((and (pair? lhs) (eq? (car lhs) '|::|))
  ;; (+= (:: x T) rhs)
  (let ((e (remove-argument-side-effects (cadr lhs)))
  (T (caddr lhs)))
  `(block ,@(cdr e)
- ,(expand-update-operator op (car e) rhs T))))
+ ,(expand-update-operator op op= (car e) rhs T))))
  (else
- (expand-update-operator- op lhs rhs declT))))
+ (expand-update-operator- op op= lhs rhs declT))))
 
 (define (lower-update-op e)
  (expand-forms
- (expand-update-operator
- (let ((str (string (car e))))
- (symbol (string.sub str 0 (- (length str) 1))))
- (cadr e)
- (caddr e))))
+ (let ((str (string (car e))))
+ (expand-update-operator
+ (symbol (string.sub str 0 (- (length str) 1)))
+ (if (= (string.char str 0) #\.) '.= '=)
+ (cadr e)
+ (caddr e)))))
 
 (define (expand-and e)
  (let ((e (cdr (flatten-ex '&& e))))
@@ -1546,11 +1547,9 @@
  (cadr expr) ;; eta reduce `x->f(x)` => `f`
  `(-> ,argname (block ,@splat ,expr)))))
 
-(define (getfield-field? x) ; whether x from (|.| f x) is a getfield call
- (or (eq? (car x) 'quote) (eq? (car x) 'inert) (eq? (car x) '$)))
-
-;; fuse nested calls to f.(args...) into a single broadcast call
-(define (expand-fuse-broadcast f args)
+; fuse nested calls to expr == f.(args...) into a single broadcast call,
+; or a broadcast! call if lhs is non-null.
+(define (expand-fuse-broadcast lhs rhs)
  (define (fuse? e) (and (pair? e) (eq? (car e) 'fuse)))
  (define (anyfuse? exprs)
  (if (null? exprs) #f (if (fuse? (car exprs)) #t (anyfuse? (cdr exprs)))))
@@ -1594,72 +1593,83 @@
  oldarg))
  fargs args)))
  (let ,fbody ,@(reverse (fuse-lets fargs args '()))))))
- (define (make-fuse f args) ; check for nested (fuse f args) exprs and combine
- (define (split-kwargs args) ; return (cons keyword-args positional-args) extracted from args
- (define (sk args kwargs pargs)
- (if (null? args)
- (cons kwargs pargs)
- (if (kwarg? (car args))
- (sk (cdr args) (cons (car args) kwargs) pargs)
- (sk (cdr args) kwargs (cons (car args) pargs)))))
- (if (has-parameters? args)
- (sk (reverse (cdr args)) (cdar args) '())
- (sk (reverse args) '() '())))
- (define (dot-to-fuse e) ; convert e == (. f (tuple args)) to (fuse f args)
- (if (and (pair? e) (eq? (car e) '|.|) (not (getfield-field? (caddr e))))
- (make-fuse (cadr e) (cdaddr e))
- e))
- (let* ((kws.args (split-kwargs args))
- (kws (car kws.args))
- (args (cdr kws.args)) ; fusing occurs on positional args only
- (args_ (map dot-to-fuse args)))
- (if (anyfuse? args_)
- `(fuse ,(fuse-funcs (to-lambda f args kws) args_) ,(fuse-args args_))
- `(fuse ,(to-lambda f args kws) ,args_))))
+ (define (dot-to-fuse e) ; convert e == (. f (tuple args)) to (fuse f args)
+ (define (make-fuse f args) ; check for nested (fuse f args) exprs and combine
+ (define (split-kwargs args) ; return (cons keyword-args positional-args) extracted from args
+ (define (sk args kwargs pargs)
+ (if (null? args)
+ (cons kwargs pargs)
+ (if (kwarg? (car args))
+ (sk (cdr args) (cons (car args) kwargs) pargs)
+ (sk (cdr args) kwargs (cons (car args) pargs)))))
+ (if (has-parameters? args)
+ (sk (reverse (cdr args)) (cdar args) '())
+ (sk (reverse args) '() '())))
+ (let* ((kws.args (split-kwargs args))
+ (kws (car kws.args))
+ (args (cdr kws.args)) ; fusing occurs on positional args only
+ (args_ (map dot-to-fuse args)))
+ (if (anyfuse? args_)
+ `(fuse ,(fuse-funcs (to-lambda f args kws) args_) ,(fuse-args args_))
+ `(fuse ,(to-lambda f args kws) ,args_))))
+ (if (and (pair? e) (eq? (car e) '|.|))
+ (let ((f (cadr e)) (x (caddr e)))
+ (if (or (eq? (car x) 'quote) (eq? (car x) 'inert) (eq? (car x) '$))
+ `(call (core getfield) ,f ,x)
+ (make-fuse f (cdr x))))
+ e))
  ; given e == (fuse lambda args), compress the argument list by removing (pure)
  ; duplicates in args, inlining literals, and moving any varargs to the end:
  (define (compress-fuse e)
  (define (findfarg arg args fargs) ; for arg in args, return corresponding farg
  (if (eq? arg (car args))
  (car fargs)
  (findfarg arg (cdr args) (cdr fargs))))
- (let ((f (cadr e))
- (args (caddr e)))
- (define (cf old-fargs old-args new-fargs new-args renames varfarg vararg)
- (if (null? old-args)
- (let ((nfargs (if (null? varfarg) new-fargs (cons varfarg new-fargs)))
- (nargs (if (null? vararg) new-args (cons vararg new-args))))
- `(fuse (-> (tuple ,@(reverse nfargs)) ,(replace-vars (caddr f) renames))
- ,(reverse nargs)))
- (let ((farg (car old-fargs)) (arg (car old-args)))
- (cond
- ((and (vararg? farg) (vararg? arg)) ; arg... must be the last argument
- (if (null? varfarg)
- (cf (cdr old-fargs) (cdr old-args)
- new-fargs new-args renames farg arg)
- (if (eq? (cadr vararg) (cadr arg))
+ (if (fuse? e)
+ (let ((f (cadr e))
+ (args (caddr e)))
+ (define (cf old-fargs old-args new-fargs new-args renames varfarg vararg)
+ (if (null? old-args)
+ (let ((nfargs (if (null? varfarg) new-fargs (cons varfarg new-fargs)))
+ (nargs (if (null? vararg) new-args (cons vararg new-args))))
+ `(fuse (-> (tuple ,@(reverse nfargs)) ,(replace-vars (caddr f) renames))
+ ,(reverse nargs)))
+ (let ((farg (car old-fargs)) (arg (car old-args)))
+ (cond
+ ((and (vararg? farg) (vararg? arg)) ; arg... must be the last argument
+ (if (null? varfarg)
  (cf (cdr old-fargs) (cdr old-args)
- new-fargs new-args (cons (cons (cadr farg) (cadr varfarg)) renames)
- varfarg vararg)
- (error "multiple splatted args cannot be fused into a single broadcast"))))
- ((number? arg) ; inline numeric literals
- (cf (cdr old-fargs) (cdr old-args)
- new-fargs new-args
- (cons (cons farg arg) renames)
- varfarg vararg))
- ((and (symbol? arg) (memq arg new-args)) ; combine duplicate args
- ; (note: calling memq for every arg is O(length(args)^2) ...
- ; ... would be better to replace with a hash table if args is long)
- (cf (cdr old-fargs) (cdr old-args)
- new-fargs new-args
- (cons (cons farg (findfarg arg new-args new-fargs)) renames)
- varfarg vararg))
- (else
- (cf (cdr old-fargs) (cdr old-args)
- (cons farg new-fargs) (cons arg new-args) renames varfarg vararg))))))
- (cf (cdadr f) args '() '() '() '() '())))
- (let ((e (compress-fuse (make-fuse f args)))) ; an expression '(fuse func args)
- (expand-forms `(call broadcast ,(from-lambda (cadr e)) ,@(caddr e)))))
+ new-fargs new-args renames farg arg)
+ (if (eq? (cadr vararg) (cadr arg))
+ (cf (cdr old-fargs) (cdr old-args)
+ new-fargs new-args (cons (cons (cadr farg) (cadr varfarg)) renames)
+ varfarg vararg)
+ (error "multiple splatted args cannot be fused into a single broadcast"))))
+ ((number? arg) ; inline numeric literals
+ (cf (cdr old-fargs) (cdr old-args)
+ new-fargs new-args
+ (cons (cons farg arg) renames)
+ varfarg vararg))
+ ((and (symbol? arg) (memq arg new-args)) ; combine duplicate args
+ ; (note: calling memq for every arg is O(length(args)^2) ...
+ ; ... would be better to replace with a hash table if args is long)
+ (cf (cdr old-fargs) (cdr old-args)
+ new-fargs new-args
+ (cons (cons farg (findfarg arg new-args new-fargs)) renames)
+ varfarg vararg))
+ (else
+ (cf (cdr old-fargs) (cdr old-args)
+ (cons farg new-fargs) (cons arg new-args) renames varfarg vararg))))))
+ (cf (cdadr f) args '() '() '() '() '()))
+ e)) ; (not (fuse? e))
+ (let ((e (compress-fuse (dot-to-fuse rhs)))) ; an expression '(fuse func args) if expr is a dot call
+ (if (fuse? e)
+ (if (null? lhs)
+ (expand-forms `(call broadcast ,(from-lambda (cadr e)) ,@(caddr e)))
+ (expand-forms `(call broadcast! ,(from-lambda (cadr e)) ,lhs ,@(caddr e))))
+ (if (null? lhs)
+ (expand-forms e)
+ (expand-forms `(call broadcast! identity ,lhs ,e))))))
 
 ;; table mapping expression head to a function expanding that form
 (define expand-table
@@ -1697,13 +1707,11 @@
 
  '|.|
  (lambda (e) ; e = (|.| f x)
- (let ((f (cadr e))
- (x (caddr e)))
- (if (getfield-field? x)
- `(call (core getfield) ,(expand-forms f) ,(expand-forms x))
- ; otherwise, came from f.(args...) --> broadcast(f, args...),
- ; where we want to fuse with any nested broadcast calls.
- (expand-fuse-broadcast f (cdr x)))))
+ (expand-fuse-broadcast '() e))
+
+ '.=
+ (lambda (e)
+ (expand-fuse-broadcast (cadr e) (caddr e)))
 
  '|<:| syntactic-op-to-call
  '|>:| syntactic-op-to-call
@@ -2008,11 +2016,16 @@
  '%= lower-update-op
  '.%= lower-update-op
  '|\|=| lower-update-op
+ '|.\|=| lower-update-op
  '&= lower-update-op
+ '.&= lower-update-op
  '$= lower-update-op
  '<<= lower-update-op
+ '.<<= lower-update-op
  '>>= lower-update-op
+ '.>>= lower-update-op
  '>>>= lower-update-op
+ '.>>>= lower-update-op
 
  ':
  (lambda (e)

test/broadcast.jl

@@ -248,6 +248,25 @@ let x = [1:4;]
  @test sin.(f17300kw.(x, y=1)) == sin.(f17300kw.(x; y=1)) == sin.(x .+ 1)
 end
 
+# PR #17510: Fused in-place assignment
+let x = [1:4;], y = x
+ y .= 2:5
+ @test y === x == [2:5;]
+ y .= factorial.(x)
+ @test y === x == [2,6,24,120]
+ y .= 7
+ @test y === x == [7,7,7,7]
+ y .= factorial.(3)
+ @test y === x == [6,6,6,6]
+ f17510() = 9
+ y .= f17510.()
+ @test y === x == [9,9,9,9]
+ y .-= 1
+ @test y === x == [8,8,8,8]
+ y .-= 1:4
+ @test y === x == [7,6,5,4]
+end
+
 # PR 16988
 @test Base.promote_op(+, Bool) === Int
 @test isa(broadcast(+, [true]), Array{Int,1})

test/show.jl

@@ -557,3 +557,7 @@ end
 @test repr(:(x for x in y if aa for z in w if bb)) == ":(x for x = y if aa for z = w if bb)"
 @test repr(:([x for x = y])) == ":([x for x = y])"
 @test repr(:([x for x = y if z])) == ":([x for x = y if z])"
+
+for op in (:(.=), :(.+=), :(.&=))
+ @test repr(parse("x $op y")) == ":(x $op y)"
+end