MaxGekk commented on code in PR #59225:
URL: https://github.com/apache/spark/pull/59225#discussion_r4201484817
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala:
##########
@@ -610,6 +616,90 @@ object WholeStageCodegenExec {
// Reset generation time of Java source code.
// Visible for testing
def resetCodeGenTime(): Unit = _codeGenTime.set(0L)
+
+ /**
+ * A trial compile's outcome for `WholeStageCodegenExec.chooseCode`: the
bytecode size of each
+ * method, None where they are unknown, or the failure; and the reports and
metric updates the
+ * compile held back for the code kept.
+ */
+ private[execution] case class TrialResult(
+ methodSizes: Try[Option[Map[String, Int]]],
+ report: () => Unit)
+
+ /**
+ * Compiles `code` as a trial (`CodeCompiler.trial`). A hit in the compile
cache produces no
+ * method sizes, so the class is then dropped and compiled again. A hit
happens where the class
+ * is cached while the decision for its code is not: the decision was
evicted, or the same body
+ * was compiled earlier without being decided - with the split off, whose
code is the recorded
+ * code without its marks, or under another limit. Should another thread
load the class again in
+ * between, the second compile is a hit as well, and the sizes are unknown.
+ */
+ private[execution] def trialCompile(
+ code: CodeAndComment,
+ failureExpected: Boolean): TrialResult = {
+ def once(): (Try[Map[String, Int]], CodeCompiler.TrialCompile) = {
+ val (result, trial) =
CodeCompiler.trial(failureExpected)(CodeGenerator.compile(code))
+ (result.map(_ => trial.methodSizes), trial)
+ }
+ val (sizes, trial) = once() match {
+ case (Success(m), _) if m.isEmpty =>
+ CodeGenerator.invalidateCompiled(code)
+ once()
+ case result => result
+ }
+ TrialResult(sizes.map(m => Option.when(m.nonEmpty)(m)), () =>
trial.report())
+ }
+
+ /** The methods a class runs once, whose being interpreted costs nothing per
row. */
+ private val RunOnce = Set("<init>", "<clinit>", "init")
Review Comment:
Done. `CodeCompiler.runOnceMethods` is `private[sql]`, and
`WholeStageCodegenExec` uses it in place of `RunOnce`.
`CodeGenerator.compileCacheScope(code)` returns the class loader and backend
part of the cache key, which both `cacheKey` and `decisionKey` build from. I
left `init_N_M` as it is in both, as you noted it's harmless today.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1271,49 +1759,159 @@ class CodegenContext extends Logging {
DEFAULT_JVM_HUGE_METHOD_LIMIT / 2 / (MAX_BYTES_PER_SPLIT_CALL + 2 *
arguments.length))
/**
- * The calls to the functions `splitExpressions` split out, grouped into
functions of their own
- * where there are too many for the calling method to be JIT-compiled, level
by level.
+ * The calls to the functions `splitExpressions` split out, grouped by
`groupSplitCalls` where
+ * there are too many for the calling method to be JIT-compiled. Every call
takes `arguments`, so
+ * the groups are of `SPLIT_CALLS_PER_GROUP` calls, and the calls stay in
the calling method while
+ * they fit in `splitCallsPerMethod`.
+ */
+ private def groupFunctionCalls(
+ calls: Seq[String],
+ funcName: String,
+ arguments: Seq[(String, String)],
+ returnType: String,
+ makeSplitFunction: String => String,
+ foldFunctions: Seq[String] => String): Seq[String] = {
+ groupSplitCalls(calls.map(SplitCall(_, arguments)), funcName, returnType,
makeSplitFunction,
+ foldFunctions, callsLeft = splitCallsPerMethod(arguments)).map(_.call)
+ }
+
+ /** A call to a split function, and the (type, name) arguments it passes. */
+ private case class SplitCall(call: String, arguments: Seq[(String, String)])
+
+ /**
+ * Adds the function `name`, taking `arguments` and running `body`, and
returns a call to it.
+ * Where `audited`, its arguments are kept under testing for
`assertSplitCallsWithin`: those of a
+ * whole stage split, computed from what its code reads where it was
generated, which an operator
+ * moving that code into a method of its own must take. The other splits'
arguments are their
+ * callers' own, which may be locals of the method they are made in.
+ */
+ private def splitCall(
+ name: String,
+ arguments: Seq[(String, String)],
+ returnType: String,
+ body: String,
+ audited: Boolean = false): SplitCall = {
+ if (Utils.isTesting) {
+ // Passing global variables to the split method is dangerous, as any
mutating to it is
+ // ignored and may lead to unexpected behavior.
+ arguments.foreach { case (_, argument) =>
+ assert(!mutableStateNames.contains(argument),
+ s"split function argument $argument cannot be a global variable.")
+ }
+ }
+ val argDefinitionString = arguments.map { case (t, argument) => s"$t
$argument" }.mkString(", ")
+ val code =
+ s"""
+ |private $returnType $name($argDefinitionString) {
+ | $body
+ |}
+ """.stripMargin
+ val call = s"${addNewFunction(name,
code)}(${arguments.map(_._2).mkString(", ")})"
+ if (INPUT_ROW != null && arguments.exists(_._2 == INPUT_ROW)) {
+ functionsTakingRow += name
+ }
+ if (audited && Utils.isTesting) {
+ splitCallArguments(name) = arguments.map(_._2)
+ }
+ SplitCall(call, arguments)
+ }
+
+ /**
+ * The calls to split functions, grouped into functions of their own, level
by level, until at
+ * most `callsLeft` remain in the calling method.
*
* The calls are what is left in the calling method, one per block of code,
and a `CASE WHEN`
* of a thousand branches has hundreds of blocks: their calls alone take the
caller past the
* 8000 bytes HotSpot compiles
(`CodeGenerator.DEFAULT_JVM_HUGE_METHOD_LIMIT`), so the method
* every row goes through runs interpreted. A call and its fold compile to
at most about 20 bytes
* plus 2 for each argument - the invocation, a load per argument, and the
fold's store, test
- * and branch - so the calls are left as they are while they fit in half of
that limit, the rest
- * being the caller's own code, and grouped otherwise. The groups are small,
of
- * `CodeGenerator.SPLIT_CALLS_PER_GROUP` calls, so that C2 can inline the
split functions into
- * each group within its inlining budget; a group of hundreds of calls stays
compilable but has
- * most of its calls left out of inlining, and runs slower. A group is
folded with
- * `foldFunctions` and wrapped with `makeSplitFunction`, which is what
- * `generateInnerClassesFunctionCalls` does with the functions of one inner
class, so every
- * caller of `splitExpressions` supports it.
- */
- private def groupFunctionCalls(
- calls: Seq[String],
+ * and branch - so outside whole stage codegen the calls are left as they
are while they fit in
+ * half of that limit, the rest being the caller's own code
(`splitCallsPerMethod`). Under whole
+ * stage codegen the caller is the operator's method, which holds every
expression's calls, so
+ * at most `SPLIT_CALLS_PER_GROUP` are left there for each.
+ *
+ * The groups are small, of `CodeGenerator.SPLIT_CALLS_PER_GROUP` calls, so
that C2 can inline
+ * the split functions into each group within its inlining budget; a group
of hundreds of calls
+ * stays compilable but has most of its calls left out of inlining, and runs
slower. A group
+ * takes the union of what its calls take, and is closed when one more call
would take it past
+ * those calls or past the parameters a method can have; where no group can
hold two calls, the
+ * calls are left as they are. A group is folded with `foldFunctions` and
wrapped with
+ * `makeSplitFunction`, which is what `generateInnerClassesFunctionCalls`
does with the functions
+ * of one inner class, so every caller of `splitExpressions` supports it.
+ */
+ private def groupSplitCalls(
+ calls: Seq[SplitCall],
funcName: String,
- arguments: Seq[(String, String)],
returnType: String,
makeSplitFunction: String => String,
foldFunctions: Seq[String] => String,
- level: Int = 0): Seq[String] = {
- val callsPerMethod = splitCallsPerMethod(arguments)
- if (calls.length <= callsPerMethod) {
+ callsLeft: Int,
+ level: Int = 0,
+ audited: Boolean = false): Seq[SplitCall] = {
+ if (calls.length <= callsLeft) {
calls
} else {
- val argDefinitionString = arguments.map { case (t, name) => s"$t $name"
}.mkString(", ")
- val argInvocationString = arguments.map(_._2).mkString(", ")
- val groupCalls = calls.grouped(SPLIT_CALLS_PER_GROUP).zipWithIndex.map {
case (group, i) =>
- val name = s"${funcName}_group${level}_$i"
- val code =
- s"""
- |private $returnType $name($argDefinitionString) {
- | ${makeSplitFunction(foldFunctions(group))}
- |}
- """.stripMargin
- s"${addNewFunction(name, code)}($argInvocationString)"
- }.toSeq
- groupFunctionCalls(groupCalls, funcName, arguments, returnType,
makeSplitFunction,
- foldFunctions, level + 1)
+ val groups = mutable.ArrayBuffer.empty[mutable.ArrayBuffer[SplitCall]]
+ var taken = Seq.empty[(String, String)]
+ calls.foreach { c =>
+ val union = (taken ++ c.arguments).distinctBy(_._2)
+ if (groups.isEmpty || groups.last.length == SPLIT_CALLS_PER_GROUP ||
+ !isValidParamLength(paramSlots(union))) {
+ groups += mutable.ArrayBuffer(c)
+ taken = c.arguments
+ } else {
+ groups.last += c
+ taken = union
+ }
+ }
+ if (groups.length == calls.length) {
+ calls
+ } else {
+ val grouped = groups.zipWithIndex.map { case (group, i) =>
+ splitCall(s"${funcName}_group${level}_$i",
+ group.flatMap(_.arguments).distinctBy(_._2).toSeq, returnType,
+ makeSplitFunction(foldFunctions(group.map(_.call).toSeq)), audited)
+ }.toSeq
+ groupSplitCalls(grouped, funcName, returnType, makeSplitFunction,
foldFunctions,
+ callsLeft, level + 1, audited)
+ }
+ }
+ }
+
+ /** The parameter slots of a method taking `arguments`, `this` included. */
+ private def paramSlots(arguments: Seq[(String, String)]): Int =
Review Comment:
Done. `paramSlots` is gone. `wholeStageSplitArguments` checks the values
with `calculateParamLengthFromExprValues`, and the group check uses a new
`calculateParamLengthFromTypes`, which both of the existing calculators share
for the per-type count. `addNewFunctionTaking` adds a function and registers it
in `functionsTakingRow` where it takes `INPUT_ROW`. `CommonExprSlots.build` and
`splitCall` both use it.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -210,9 +339,10 @@ class CodegenContext extends Logging {
/**
* The slots a `CommonExpressionRef` reads: the value and its nullness, plus
the flag saying
- * whether this row has computed them yet, and the definition to compute
them from.
+ * whether this row has computed them yet, and the definition, with its id,
to compute them from.
*/
case class CommonExprSlots(
+ id: Long,
Review Comment:
Removed.
##########
sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/expressions/codegen/CodeGenerator.scala:
##########
@@ -1823,6 +2444,96 @@ object CodeGenerator extends Logging {
// method which is going to be called by the outer class instead of the many
small ones
final val MERGE_SPLIT_METHODS_THRESHOLD = 3
+ //
---------------------------------------------------------------------------------------------
Review Comment:
Moved: the whole group, header included, is now at the end of `object
CodeGenerator`. The JVM-limit constants are one run again, as on master.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
