remove compute optimizer from lte (learning on the edge) builds (#15637)

### Description
Removing compute optimizer from on device training builds.

### Motivation and Context
1. mitigate android build failures
2. reduce binary size

Since only CPU EP is enabled for LTE builds, we can optimize the models
offline.

onnxruntime/core/optimizer/compute_optimizer/shared_utils.cc

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 
 #include <onnx/defs/attr_proto_util.h>
 #include "core/graph/graph_utils.h"

onnxruntime/core/optimizer/compute_optimizer/shared_utils.h

@@ -3,7 +3,7 @@
 
 // The optimization here ideally is applicable to both training and inferencing,
 // while so far we mainly validate on training during cooking the optimization.
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 #pragma once
 
 #include <initializer_list>

onnxruntime/core/optimizer/compute_optimizer/upstream_gather.cc

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 
 #include <onnx/defs/attr_proto_util.h>
 #include "core/graph/graph_utils.h"

onnxruntime/core/optimizer/compute_optimizer/upstream_gather.h

@@ -3,7 +3,7 @@
 
 // The optimization here ideally applies to both training and inferencing,
 // while so far we mainly validate training during cooking the optimization.
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 #pragma once
 
 #include "core/optimizer/compute_optimizer/upstream_transformer_base.h"

onnxruntime/core/optimizer/compute_optimizer/upstream_gather_actors.cc

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 
 #include <onnx/defs/attr_proto_util.h>
 #include "core/graph/graph_utils.h"

onnxruntime/core/optimizer/compute_optimizer/upstream_gather_actors.h

@@ -3,7 +3,7 @@
 
 // The optimization here ideally applies to both training and inferencing,
 // while so far we mainly validate training during cooking the optimization.
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 #pragma once
 
 #include "core/optimizer/compute_optimizer/shared_utils.h"

onnxruntime/core/optimizer/compute_optimizer/upstream_reshape.cc

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 
 #include "core/framework/tensorprotoutils.h"
 #include "core/graph/graph_utils.h"

onnxruntime/core/optimizer/compute_optimizer/upstream_reshape.h

@@ -3,7 +3,7 @@
 
 // The optimization here ideally applies to both training and inference,
 // while so far we mainly validate training during cooking the optimization.
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 #pragma once
 
 #include "core/optimizer/compute_optimizer/upstream_transformer_base.h"

onnxruntime/core/optimizer/compute_optimizer/upstream_reshape_actors.cc

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 
 #include "core/optimizer/utils.h"
 #include "core/optimizer/compute_optimizer/upstream_reshape_actors.h"

onnxruntime/core/optimizer/compute_optimizer/upstream_reshape_actors.h

@@ -3,7 +3,7 @@
 
 // The optimization here ideally applies to both training and inference,
 // while so far we mainly validate training during cooking the optimization.
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 #pragma once
 
 #include "core/optimizer/compute_optimizer/shared_utils.h"

onnxruntime/core/optimizer/compute_optimizer/upstream_transformer_base.cc

@@ -1,7 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 
 #include <onnx/defs/attr_proto_util.h>
 #include "core/common/safeint.h"

onnxruntime/core/optimizer/compute_optimizer/upstream_transformer_base.h

@@ -3,7 +3,7 @@
 
 // The optimization here ideally applies to both training and inferencing,
 // while so far we mainly validate training during cooking the optimization.
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 #pragma once
 
 #include "core/optimizer/graph_transformer.h"

onnxruntime/core/optimizer/graph_transformer_utils.cc

@@ -70,12 +70,10 @@
 #include "core/optimizer/slice_elimination.h"
 #include "core/optimizer/transpose_optimizer/ort_transpose_optimizer.h"
 #include "core/optimizer/unsqueeze_elimination.h"
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
 #include "orttraining/core/optimizer/bias_softmax_dropout_fusion.h"
 #include "orttraining/core/optimizer/bitmask_dropout_replacement.h"
 #include "orttraining/core/optimizer/sce_loss_grad_bias_fusion.h"
-#endif
-#ifdef ENABLE_TRAINING
 #include "orttraining/core/optimizer/memory_optimizer.h"
 #endif
 
@@ -299,7 +297,7 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
       transformers.emplace_back(std::make_unique<BiasGeluFusion>(cpu_cuda_dml_rocm_eps));
       transformers.emplace_back(std::make_unique<BiasSoftmaxFusion>(cpu_cuda_rocm_eps));
       transformers.emplace_back(std::make_unique<BiasDropoutFusion>(cuda_rocm_eps));
-#ifdef ENABLE_TRAINING_CORE
+#ifdef ENABLE_TRAINING
       transformers.emplace_back(std::make_unique<BitmaskDropoutReplacement>(cuda_rocm_eps));
       transformers.emplace_back(std::make_unique<BiasSoftmaxDropoutFusion>(cuda_rocm_eps));
       transformers.emplace_back(std::make_unique<SceLossGradBiasFusion>(cpu_cuda_rocm_eps));

orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.cc

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#ifdef ENABLE_TRAINING
 #include <onnx/defs/attr_proto_util.h>
 
 #include "orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h"
@@ -223,3 +224,5 @@ Status InsertGatherBeforeSceLoss::ApplyImpl(Graph& graph, bool& modified, int /*
 }
 
 }  // namespace onnxruntime
+
+#endif

orttraining/orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#ifdef ENABLE_TRAINING
 #pragma once
 
 #include "core/optimizer/graph_transformer.h"
@@ -75,3 +76,5 @@ class InsertGatherBeforeSceLoss : public GraphTransformer {
 };
 
 }  // namespace onnxruntime
+
+#endif

orttraining/orttraining/core/optimizer/graph_transformer_utils.cc

@@ -10,8 +10,6 @@
 #include "core/optimizer/bias_softmax_fusion.h"
 #include "core/optimizer/cast_elimination.h"
 #include "core/optimizer/common_subexpression_elimination.h"
-#include "core/optimizer/compute_optimizer/upstream_gather.h"
-#include "core/optimizer/compute_optimizer/upstream_reshape.h"
 #include "core/optimizer/concat_slice_elimination.h"
 #include "core/optimizer/constant_folding.h"
 #include "core/optimizer/constant_sharing.h"
@@ -53,7 +51,6 @@
 #include "orttraining/core/framework/distributed_run_context.h"
 #include "orttraining/core/optimizer/batchnorm_replacement.h"
 #include "orttraining/core/optimizer/bitmask_dropout_replacement.h"
-#include "orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h"
 #include "orttraining/core/optimizer/concat_replacement.h"
 #include "orttraining/core/optimizer/graph_transformer_registry.h"
 #include "orttraining/core/optimizer/insert_output_rewriter.h"
@@ -63,6 +60,13 @@
 #include "orttraining/core/optimizer/transformer_layer_recompute.h"
 #include "orttraining/core/optimizer/qdq_fusion.h"
 
+// Only enabled in full training build. Not in on device training builds
+#ifdef ENABLE_TRAINING
+#include "core/optimizer/compute_optimizer/upstream_gather.h"
+#include "core/optimizer/compute_optimizer/upstream_reshape.h"
+#include "orttraining/core/optimizer/compute_optimizer/sceloss_compute_optimization.h"
+#endif
+
 namespace onnxruntime {
 namespace training {
 namespace transformer_utils {
@@ -162,6 +166,7 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
                                                                      cuda_execution_provider));
       }
 
+#ifdef ENABLE_TRAINING
       if (config.enable_compute_optimizer) {
         transformers.emplace_back(std::make_unique<UpStreamGatherGraphTransformer>(compatible_eps));
         if (config.enable_label_sparsity_optimization) {
@@ -169,6 +174,7 @@ std::vector<std::unique_ptr<GraphTransformer>> GeneratePreTrainingTransformers(
           transformers.emplace_back(std::make_unique<InsertGatherBeforeSceLoss>(compatible_eps));
         }
       }
+#endif
 
     } break;
 

orttraining/orttraining/test/optimizer/compute_optimizer_test.cc

@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+// Only enabled in training full build, not in on device training build.
+#ifdef ENABLE_TRAINING
 #include <random>
 #include "core/graph/onnx_protobuf.h"
 
@@ -426,3 +428,5 @@ TEST(ComputeOptimizerTests, InsertGatherBeforeSceLoss_MlmBertE2E) {
 
 }  // namespace test
 }  // namespace onnxruntime
+
+#endif