1 /* Copyright 2022 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_DTENSOR_MLIR_SPARSE_EXPANDER_H_ 17 #define TENSORFLOW_DTENSOR_MLIR_SPARSE_EXPANDER_H_ 18 19 #include <string> 20 21 #include "absl/container/flat_hash_map.h" 22 #include "mlir/IR/Builders.h" // from @llvm-project 23 #include "mlir/IR/Operation.h" // from @llvm-project 24 #include "mlir/IR/UseDefLists.h" // from @llvm-project 25 #include "tensorflow/core/framework/registration/registration.h" 26 #include "tensorflow/dtensor/cc/dstatus.h" 27 28 namespace tensorflow { 29 namespace dtensor { 30 31 // Base class for handling Sparse expansion of a MLIR TF Operation. 32 // Note that an op will only go through Sparse Expansion only if it has 33 // any sparse input tensors. 34 class SparseExpanderBase { 35 public: ~SparseExpanderBase()36 virtual ~SparseExpanderBase() {} 37 38 // Converts `op` to a Sparse expanded form. Sparse expansion logic is 39 // a function of op type and op's operand type. 40 // Must return the `op` that is expanded as the final return value. 41 // 42 // An op has a SparseTensor operand if the defining op of that operand 43 // is a SparseToDenseOp. 44 virtual StatusOr<mlir::Operation*> ExpandOp(mlir::Operation* op) = 0; 45 }; 46 47 // Computes the Sparse expansion for `op`. 48 Status RunSparseExpansion(mlir::Operation* op, mlir::Operation** output); 49 50 // A registry of sparse SPMD expanders. This map is statically stored and 51 // initialized with all the registered sparse SPMD expanders. 52 class SparseExpanderRegistry { 53 public: 54 ~SparseExpanderRegistry() = default; 55 56 // A singleton available at startup. 57 static SparseExpanderRegistry* Global(); 58 59 // Returns the sparse expansion for the given operation (or nullptr if no 60 // expansion has been registered). 61 SparseExpanderBase* GetSparseExpansionFnForOp(mlir::Operation* op); 62 63 // Registers a sparse expander for the provided opName. 64 InitOnStartupMarker RegisterSparseExpansionFn( 65 std::string opName, std::unique_ptr<SparseExpanderBase> prop); 66 67 private: 68 absl::flat_hash_map<std::string, std::unique_ptr<SparseExpanderBase>> 69 op_to_sparse_expansion_fn_map_; 70 }; 71 72 #define REGISTER_SPARSE(name, op, prop, ...) \ 73 static ::tensorflow::InitOnStartupMarker const spmd_##name = \ 74 InitOnStartupMarker{} \ 75 << SparseExpanderRegistry::Global()->RegisterSparseExpansionFn( \ 76 mlir::op ::getOperationName().str(), \ 77 std::make_unique<prop>(__VA_ARGS__)) 78 79 } // namespace dtensor 80 } // namespace tensorflow 81 82 #endif // TENSORFLOW_DTENSOR_MLIR_SPARSE_EXPANDER_H_ 83