diff --git a/java/config/config.go b/java/config/config.go
index eb71ddb..f4abd37 100644
--- a/java/config/config.go
+++ b/java/config/config.go
@@ -77,7 +77,22 @@
 	pctx.SourcePathVariable("JarArgsCmd", "build/soong/scripts/jar-args.sh")
 	pctx.HostBinToolVariable("SoongZipCmd", "soong_zip")
 	pctx.HostBinToolVariable("MergeZipsCmd", "merge_zips")
-	pctx.HostBinToolVariable("DxCmd", "dx")
+	pctx.VariableFunc("DxCmd", func(config interface{}) (string, error) {
+		dexer := "dx"
+		if config.(android.Config).Getenv("USE_D8") == "true" {
+			dexer = "d8"
+		}
+		if config.(android.Config).UnbundledBuild() {
+			return "prebuilts/build-tools/common/bin/" + dexer, nil
+		} else {
+			path, err := pctx.HostBinToolPath(config, dexer)
+			if err != nil {
+				return "", err
+			}
+			return path.String(), nil
+		}
+	})
+
 	pctx.HostJavaToolVariable("JarjarCmd", "jarjar.jar")
 	pctx.HostJavaToolVariable("DesugarJar", "desugar.jar")
 
diff --git a/java/config/makevars.go b/java/config/makevars.go
index 6b49592..2735242 100644
--- a/java/config/makevars.go
+++ b/java/config/makevars.go
@@ -43,6 +43,8 @@
 	ctx.Strict("JAR_ARGS", "${JarArgsCmd}")
 	ctx.Strict("JAVADOC", "${JavadocCmd}")
 	ctx.Strict("COMMON_JDK_FLAGS", "${CommonJdkFlags}")
+	ctx.Strict("DX", "${DxCmd}")
+	ctx.Strict("DX_COMMAND", "${DxCmd} -JXms16M -JXmx2048M")
 
 	if ctx.Config().IsEnvTrue("RUN_ERROR_PRONE") {
 		ctx.Strict("TARGET_JAVAC", "${ErrorProneCmd}")