[OPTIMIZER] Improved layout simplification pass so it handles swizzled layouts better (#789)

Note: uncommented `test_gemm`, since backend has an issue with swizzling. This will get uncommented in a subsequent PR.

python/src/triton.cc

@@ -1185,6 +1185,10 @@ void init_triton_ir(py::module &&m) {
            [](mlir::PassManager &self) {
              self.addPass(mlir::createTritonGPUCombineOpsPass());
            })
+      .def("add_triton_gpu_swizzle_pass",
+           [](mlir::PassManager &self) {
+             self.addPass(mlir::createTritonGPUSwizzlePass());
+           })
       .def("add_triton_gpu_to_llvm",
            [](mlir::PassManager &self) {
              self.addPass(mlir::triton::createConvertTritonGPUToLLVMPass());