Regenerate MLIR Bindings (#1543)

Co-authored-by: enzyme-ci-bot[bot] <78882869+enzyme-ci-bot[bot]@users.noreply.github.com>

Author: github-actions[bot]

src/mlir/Dialects/EnzymeXLA.jl

@@ -244,6 +244,107 @@ function ml_gelu(
     )
 end
 
+"""
+`lapack_gemqrt`
+
+This operation is modeled after LAPACK\'s *GEMQR routines.
+"""
+function lapack_gemqrt(
+    V::Value,
+    T::Value,
+    C::Value;
+    output::IR.Type,
+    side,
+    transpose=nothing,
+    location=Location(),
+)
+    op_ty_results = IR.Type[output,]
+    operands = Value[V, T, C]
+    owned_regions = Region[]
+    successors = Block[]
+    attributes = NamedAttribute[namedattribute("side", side),]
+    !isnothing(transpose) && push!(attributes, namedattribute("transpose", transpose))
+
+    return create_operation(
+        "enzymexla.lapack.gemqrt",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=op_ty_results,
+        result_inference=false,
+    )
+end
+
+"""
+`lapack_geqrf`
+
+This operation computes the QR factorization of a matrix using Householder 
+reflections. Mathematically, it decomposes A into the product of an 
+orthogonal matrix Q and an upper triangular matrix R, such that A = QR.
+
+This operation is modeled after LAPACK\'s *GEQRF routines, which returns the 
+result in the QR packed format.
+"""
+function lapack_geqrf(
+    input::Value; output::IR.Type, tau::IR.Type, info::IR.Type, location=Location()
+)
+    op_ty_results = IR.Type[output, tau, info]
+    operands = Value[input,]
+    owned_regions = Region[]
+    successors = Block[]
+    attributes = NamedAttribute[]
+
+    return create_operation(
+        "enzymexla.lapack.geqrf",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=op_ty_results,
+        result_inference=false,
+    )
+end
+
+"""
+`lapack_geqrt`
+
+This operation computes the QR factorization of a matrix using Householder 
+reflections. Mathematically, it decomposes A into the product of an 
+orthogonal matrix Q and an upper triangular matrix R, such that A = QR.
+
+This operation is modeled after LAPACK\'s *GEQRT routines, which returns the 
+result in the QR CompactWY format.
+"""
+function lapack_geqrt(
+    input::Value;
+    output::IR.Type,
+    T::IR.Type,
+    info::IR.Type,
+    blocksize=nothing,
+    location=Location(),
+)
+    op_ty_results = IR.Type[output, T, info]
+    operands = Value[input,]
+    owned_regions = Region[]
+    successors = Block[]
+    attributes = NamedAttribute[]
+    !isnothing(blocksize) && push!(attributes, namedattribute("blocksize", blocksize))
+
+    return create_operation(
+        "enzymexla.lapack.geqrt",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=op_ty_results,
+        result_inference=false,
+    )
+end
+
 function get_stream(; result::IR.Type, location=Location())
     op_ty_results = IR.Type[result,]
     operands = Value[]
@@ -270,6 +371,8 @@ function jit_call(
     backend_config=nothing,
     operand_layouts=nothing,
     result_layouts=nothing,
+    arg_attrs=nothing,
+    res_attrs=nothing,
     output_operand_aliases=nothing,
     xla_side_effect_free=nothing,
     location=Location(),
@@ -285,6 +388,8 @@ function jit_call(
         push!(attributes, namedattribute("operand_layouts", operand_layouts))
     !isnothing(result_layouts) &&
         push!(attributes, namedattribute("result_layouts", result_layouts))
+    !isnothing(arg_attrs) && push!(attributes, namedattribute("arg_attrs", arg_attrs))
+    !isnothing(res_attrs) && push!(attributes, namedattribute("res_attrs", res_attrs))
     !isnothing(output_operand_aliases) &&
         push!(attributes, namedattribute("output_operand_aliases", output_operand_aliases))
     !isnothing(xla_side_effect_free) &&
@@ -316,6 +421,8 @@ function kernel_call(
     backend_config=nothing,
     operand_layouts=nothing,
     result_layouts=nothing,
+    arg_attrs=nothing,
+    res_attrs=nothing,
     output_operand_aliases=nothing,
     xla_side_effect_free=nothing,
     location=Location(),
@@ -331,6 +438,8 @@ function kernel_call(
         push!(attributes, namedattribute("operand_layouts", operand_layouts))
     !isnothing(result_layouts) &&
         push!(attributes, namedattribute("result_layouts", result_layouts))
+    !isnothing(arg_attrs) && push!(attributes, namedattribute("arg_attrs", arg_attrs))
+    !isnothing(res_attrs) && push!(attributes, namedattribute("res_attrs", res_attrs))
     !isnothing(output_operand_aliases) &&
         push!(attributes, namedattribute("output_operand_aliases", output_operand_aliases))
     !isnothing(xla_side_effect_free) &&
@@ -457,6 +566,63 @@ function noop(blockDims::Vector{Value}; location=Location())
     )
 end
 
+"""
+`lapack_orgqr`
+
+This operation is modeled after LAPACK\'s *ORGQR/*UNGQR routines.
+"""
+function lapack_orgqr(input::Value, tau::Value; output::IR.Type, location=Location())
+    op_ty_results = IR.Type[output,]
+    operands = Value[input, tau]
+    owned_regions = Region[]
+    successors = Block[]
+    attributes = NamedAttribute[]
+
+    return create_operation(
+        "enzymexla.lapack.orgqr",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=op_ty_results,
+        result_inference=false,
+    )
+end
+
+"""
+`lapack_ormqr`
+
+This operation is modeled after LAPACK\'s *ORMQR routines.
+"""
+function lapack_ormqr(
+    A::Value,
+    tau::Value,
+    C::Value;
+    output::IR.Type,
+    side,
+    transpose=nothing,
+    location=Location(),
+)
+    op_ty_results = IR.Type[output,]
+    operands = Value[A, tau, C]
+    owned_regions = Region[]
+    successors = Block[]
+    attributes = NamedAttribute[namedattribute("side", side),]
+    !isnothing(transpose) && push!(attributes, namedattribute("transpose", transpose))
+
+    return create_operation(
+        "enzymexla.lapack.ormqr",
+        location;
+        operands,
+        owned_regions,
+        successors,
+        attributes,
+        results=op_ty_results,
+        result_inference=false,
+    )
+end
+
 function pointer2memref(source::Value; result::IR.Type, location=Location())
     op_ty_results = IR.Type[result,]
     operands = Value[source,]
@@ -495,14 +661,29 @@ function polygeist_yield(; location=Location())
     )
 end
 
+"""
+`linalg_qr`
+
+This operation computes the QR factorization of a matrix using Householder 
+reflections. Mathematically, it decomposes A into the product of an 
+orthogonal (unitary if complex) matrix Q and an upper triangular matrix R, 
+such that A = QR.
+
+If A has size m x n and m > n, Q is an m x n isometric matrix. If m < n, R
+will be a m x n trapezoidal matrix.
+
+This operation is modeled after the mathematical formulation of the QR 
+factorization, and not after LAPACK\'s compact formats.
+"""
 function linalg_qr(
-    input::Value; output::IR.Type, tau::IR.Type, info::IR.Type, location=Location()
+    input::Value; Q::IR.Type, R::IR.Type, algorithm=nothing, location=Location()
 )
-    op_ty_results = IR.Type[output, tau, info]
+    op_ty_results = IR.Type[Q, R]
     operands = Value[input,]
     owned_regions = Region[]
     successors = Block[]
     attributes = NamedAttribute[]
+    !isnothing(algorithm) && push!(attributes, namedattribute("algorithm", algorithm))
 
     return create_operation(
         "enzymexla.linalg.qr",
@@ -642,12 +823,16 @@ function wrap(
     )
 end
 
-function xla_wrapper(inputs::Vector{Value}; fn, location=Location())
+function xla_wrapper(
+    inputs::Vector{Value}; fn, arg_attrs=nothing, res_attrs=nothing, location=Location()
+)
     op_ty_results = IR.Type[]
     operands = Value[inputs...,]
     owned_regions = Region[]
     successors = Block[]
     attributes = NamedAttribute[namedattribute("fn", fn),]
+    !isnothing(arg_attrs) && push!(attributes, namedattribute("arg_attrs", arg_attrs))
+    !isnothing(res_attrs) && push!(attributes, namedattribute("res_attrs", res_attrs))
 
     return create_operation(
         "enzymexla.xla_wrapper",

src/mlir/Dialects/MosaicGPU.jl

@@ -404,7 +404,7 @@ end
 Schedules `tcgen05.mma` instructions that perform the following matrix
 multiply and accumulate:
 
-  accumulator = a * b + accumulator
+  accumulator += a * b
 
 This operation supports larger inputs than the PTX-level MMA instruction
 and will schedule as many PTX-level MMA instructions as needed to
@@ -417,8 +417,6 @@ The inputs should have the following shapes:
 where `s == swizzle / element_bytewidth` and `m` is specified according to
 https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-matrix-shape.
 
-The output has an identical shape and type as the input accumulator.
-
 The `accumulator`, `a` and `b` matrices need to be provided as 2-dimensional
 memrefs. The `accumulator` is always in TMEM and `b` is always in SMEM.
 `a` can be in TMEM or SMEM. `a` and `b` must have the same element
@@ -427,7 +425,7 @@ type and when `a` is in TMEM only F16 or BF16 are supported.
 `a_scale` and `b_scale` are optional scaling matrices that reside in TMEM.
 When set the operation is defined as:
 
-  accumulator = (a * a_scale) * (b * b_scale) + accumulator
+  accumulator += (a * a_scale) * (b * b_scale)
 
 `accumulate` is a boolean that indicates whether to perform the accumulate
 step.
@@ -439,7 +437,6 @@ function tcgen05_mma(
     accumulate::Value,
     a_scale=nothing::Union{Nothing,Value};
     b_scale=nothing::Union{Nothing,Value},
-    result_0=nothing::Union{Nothing,IR.Type},
     collective=nothing,
     location=Location(),
 )
@@ -463,7 +460,6 @@ function tcgen05_mma(
             1
         end,
     ]))
-    !isnothing(result_0) && push!(op_ty_results, result_0)
     !isnothing(collective) && push!(attributes, namedattribute("collective", collective))
 
     return create_operation(
@@ -473,8 +469,8 @@ function tcgen05_mma(
         owned_regions,
         successors,
         attributes,
-        results=(length(op_ty_results) == 0 ? nothing : op_ty_results),
-        result_inference=(length(op_ty_results) == 0 ? true : false),
+        results=op_ty_results,
+        result_inference=false,
     )
 end