[TORCH] Transformer encoder decomposition

- Add a dedicated DecomposeTransformerEncoder pass to expand encoder ops into primitive Torch patterns.
- Extend shared lowering helpers (ReduceOpVariants.cpp, Utils.h)
  so the new pass can reuse reduction utilities during decomposition.
- Register the pass in the Torch Transform pipeline so it runs as part of the decomposition flow.
- Expand e2e coverage with new transformer encoder tests to validate the lowering path.

Signed-off-by: Cathal Corbett <[email protected]>
Change-Id: I6bcda53569cf7b06df4cb97c624bbf512d8fecb7

Author: catcor01

include/torch-mlir/Dialect/Torch/Transforms/Passes.h

@@ -12,6 +12,7 @@
 
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Pass/Pass.h"
+#include "llvm/ADT/StringSet.h"
 
 #include <memory>
 
@@ -157,6 +158,10 @@ static const char kTorchOpPrefix[] = R"(torch.)";
 void populateRestructureNonConstantAxesPattern(RewritePatternSet &patterns,
                                                MLIRContext *context);
 
+void populateTransformerEncoderPatterns(
+    RewritePatternSet &patterns, const llvm::StringSet<> &legalOpsSet);
+
+
 std::unique_ptr<OperationPass<func::FuncOp>>
 createRestructureNonConstantAxesPass();
 

lib/Dialect/Torch/Transforms/CMakeLists.txt

@@ -1,6 +1,7 @@
 add_mlir_library(TorchMLIRTorchPasses
   AdjustCallingConventions.cpp
   DecomposeComplexOps.cpp
+  DecomposeTransformerEncoder.cpp
   DropAbstractInterpCalculations.cpp
   EraseModuleInitializer.cpp
   FuseQuantizedOps.cpp

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -13084,6 +13084,8 @@ class DecomposeComplexOpsPass
     legalOpsSet.clear();
     legalOpsSet.insert(legalOps.begin(), legalOps.end());
 
+    populateTransformerEncoderPatterns(patterns, legalOpsSet);
+
     addPatternIfTargetOpIsIllegal<DecomposeAten_WeightNormInterfaceOp>(
         patterns);
     addPatternIfTargetOpIsIllegal<DecomposeAtenSoftmaxIntOp>(patterns);