tensor-dialect: changed to pass mac tests, moved bufferize pass

Author: liamsemeria

src/xtc/backends/mlir/MlirCompiler.py

@@ -20,7 +20,7 @@
 from xtc.backends.mlir.MlirCompilerPasses import (
     MlirProgramInsertTransformPass,
     MlirProgramApplyTransformPass,
-    MlirProgramApplyPasses,
+    apply_bufferization_passes,
 )
 
 from xtc.backends.mlir.MlirTarget import (
@@ -151,18 +151,11 @@ def mlir_apply_transform_pass(self) -> None:
             self.dump_ir("IR Dump After transform")
 
     def mlir_apply_tensor_lowering_pass(self) -> None:
-        apply_transform_pass = MlirProgramApplyPasses(
-            mlir_program=self._mlir_program,
-        )
         if self._config.print_bufferization_ir:
             self.dump_ir("IR Dump Before Tensor Lowering")
-        apply_transform_pass.run(
-            [
-                "eliminate-empty-tensors",  # causes ops to write directly to out buffer
-                "one-shot-bufferize{bufferize-function-boundaries=1 function-boundary-type-conversion=identity-layout-map buffer-alignment=256}",
-                "func.func(promote-buffers-to-stack)",
-            ]
-        )
+
+        apply_bufferization_passes(self._mlir_program)
+
         if self._config.print_bufferization_ir:
             self.dump_ir("IR Dump After Tensor Lowering")
 

src/xtc/backends/mlir/MlirCompilerPasses.py

@@ -26,6 +26,7 @@
     OpResult,
 )
 from mlir.passmanager import PassManager
+import platform
 
 # Import SDist if available
 try:
@@ -549,3 +550,22 @@ def run(self, pass_names: list[str]) -> None:
         for name in pass_names:
             pm.add(name)  # type: ignore # no attribute add
         pm.run(self._mlir_program.mlir_module.operation)
+
+
+def apply_bufferization_passes(mlir_program: RawMlirProgram):
+    apply_passes = MlirProgramApplyPasses(mlir_program)
+    bufferize_options = [
+        "bufferize-function-boundaries=1",
+        "function-boundary-type-conversion=identity-layout-map",
+        "buffer-alignment=256",
+    ]
+    # needed for now because macos mlir version needs to be updated
+    if platform.system() != "Darwin":
+        bufferize_options.append("buffer-alignment=256")
+    apply_passes.run(
+        [
+            "eliminate-empty-tensors",  # causes ops to write directly to out buffer
+            f"one-shot-bufferize{{{' '.join(bufferize_options)}}}",
+            "func.func(promote-buffers-to-stack)",
+        ]
+    )