compiler.py

#!/usr/bin/env python3

import re
import os
import sys

class MDSLCompiler:
    def __init__(self):
        self.variables = {}
        self.action_queue = []
        self.cpp_code = []
        self.added_implementations = set() # Keep track of added functions
        
    def parse_file(self, filename):
        try:
            with open(filename, 'r') as f:
                lines = f.readlines()
        except FileNotFoundError:
            print(f"Error: Input file '{filename}' not found.")
            sys.exit(1)
        except Exception as e:
            print(f"Error reading input file '{filename}': {e}")
            sys.exit(1)

        for i, line in enumerate(lines):
            line_num = i + 1
            line = line.strip()
            if not line or line.startswith('#'):
                continue
                
            self.parse_line(line, line_num)
        
        return self.action_queue
    
    def parse_line(self, line, line_num):
        # Remove comments
        if '#' in line:
            line = line[:line.index('#')].strip()
        
        # Skip empty lines
        if not line:
            return
        
        # Define regex patterns for clarity
        load_pattern = r'image\s+(\w+)\s*=\s*load\s*\(\s*(?:"([^"]+)")?\s*\)\s*;'
        save_pattern = r'save\s*\(\s*(\w+)\s*,\s*"([^"]+)"\s*\)\s*;'
        dilate_pattern_new = r'image\s+(\w+)\s*=\s*dilate\s*\(\s*src\s*=\s*(\w+)\s*,\s*ksize\s*=\s*(\d+)\s*\)\s*;'
        dilate_pattern_reassign = r'(\w+)\s*=\s*dilate\s*\(\s*src\s*=\s*(\w+)\s*,\s*ksize\s*=\s*(\d+)\s*\)\s*;'
        erode_pattern_new = r'image\s+(\w+)\s*=\s*erode\s*\(\s*src\s*=\s*(\w+)\s*,\s*ksize\s*=\s*(\d+)\s*\)\s*;'
        erode_pattern_reassign = r'(\w+)\s*=\s*erode\s*\(\s*src\s*=\s*(\w+)\s*,\s*ksize\s*=\s*(\d+)\s*\)\s*;'
        math_pattern = r'(\w+)\s*=\s*(\w+)\s*([+-])\s*(\w+)\s*;' # Allow assignment to existing var

        # Try matching patterns sequentially
        load_match = re.match(load_pattern, line)
        if load_match:
            var_name, image_path = load_match.groups()
            if var_name in self.variables:
                print(f"Warning: Line {line_num}: Variable '{var_name}' redefined.")
            self.variables[var_name] = {'type': 'image', 'source': 'load'}
            self.action_queue.append({
                'action': 'load',
                'dest': var_name,
                'path': image_path
            })
            return
        
        save_match = re.match(save_pattern, line)
        if save_match:
            var_name, save_suffix = save_match.groups()
            if var_name not in self.variables:
                 print(f"Error: Line {line_num}: Variable '{var_name}' used before assignment in save operation.")
                 # Decide whether to exit or just warn
                 # sys.exit(1)
            self.action_queue.append({
                'action': 'save',
                'src': var_name,
                'suffix': save_suffix
            })
            return
        
        # Check both patterns for dilate (new var declaration and reassignment)
        dilate_match_new = re.match(dilate_pattern_new, line)
        dilate_match_reassign = re.match(dilate_pattern_reassign, line)
        
        if dilate_match_new:
            dest_var, src_var, ksize_str = dilate_match_new.groups()
            ksize = int(ksize_str)
            if src_var not in self.variables:
                 print(f"Error: Line {line_num}: Source variable '{src_var}' not defined for dilate operation.")
                 # sys.exit(1)
            if dest_var in self.variables:
                print(f"Warning: Line {line_num}: Variable '{dest_var}' redefined.")
            self.variables[dest_var] = {'type': 'image', 'source': 'dilate'}
            self.action_queue.append({
                'action': 'dilate',
                'dest': dest_var,
                'src': src_var,
                'ksize': ksize
            })
            return
        elif dilate_match_reassign:
            dest_var, src_var, ksize_str = dilate_match_reassign.groups()
            ksize = int(ksize_str)
            if src_var not in self.variables:
                 print(f"Error: Line {line_num}: Source variable '{src_var}' not defined for dilate operation.")
                 # sys.exit(1)
            if dest_var not in self.variables:
                 print(f"Warning: Line {line_num}: Destination variable '{dest_var}' not previously defined. Creating it now.")
                 self.variables[dest_var] = {'type': 'image', 'source': 'dilate'}
            self.action_queue.append({
                'action': 'dilate',
                'dest': dest_var,
                'src': src_var,
                'ksize': ksize
            })
            return

        # Check both patterns for erode (new var declaration and reassignment)
        erode_match_new = re.match(erode_pattern_new, line)
        erode_match_reassign = re.match(erode_pattern_reassign, line)
        
        if erode_match_new:
            dest_var, src_var, ksize_str = erode_match_new.groups()
            ksize = int(ksize_str)
            if src_var not in self.variables:
                 print(f"Error: Line {line_num}: Source variable '{src_var}' not defined for erode operation.")
                 # sys.exit(1)
            if dest_var in self.variables:
                print(f"Warning: Line {line_num}: Variable '{dest_var}' redefined.")
            self.variables[dest_var] = {'type': 'image', 'source': 'erode'}
            self.action_queue.append({
                'action': 'erode',
                'dest': dest_var,
                'src': src_var,
                'ksize': ksize
            })
            return
        elif erode_match_reassign:
            dest_var, src_var, ksize_str = erode_match_reassign.groups()
            ksize = int(ksize_str)
            if src_var not in self.variables:
                 print(f"Error: Line {line_num}: Source variable '{src_var}' not defined for erode operation.")
                 # sys.exit(1)
            if dest_var not in self.variables:
                 print(f"Warning: Line {line_num}: Destination variable '{dest_var}' not previously defined. Creating it now.")
                 self.variables[dest_var] = {'type': 'image', 'source': 'erode'}
            self.action_queue.append({
                'action': 'erode',
                'dest': dest_var,
                'src': src_var,
                'ksize': ksize
            })
            return

        math_match = re.match(math_pattern, line)
        if math_match:
            dest_var, left_var, op, right_var = math_match.groups()
            if left_var not in self.variables:
                 print(f"Error: Line {line_num}: Left operand '{left_var}' not defined for math operation.")
                 # sys.exit(1)
            if right_var not in self.variables:
                 print(f"Error: Line {line_num}: Right operand '{right_var}' not defined for math operation.")
                 # sys.exit(1)
            # Allow overwriting existing variables with math results
            if dest_var not in self.variables:
                 print(f"Warning: Line {line_num}: Destination variable '{dest_var}' not previously defined. Assuming image type.")
                 self.variables[dest_var] = {'type': 'image', 'source': 'math'}
            elif self.variables[dest_var]['type'] != 'image':
                 print(f"Warning: Line {line_num}: Overwriting non-image variable '{dest_var}' with image result.")
                 self.variables[dest_var]['type'] = 'image' # Update type if necessary

            self.action_queue.append({
                'action': 'math',
                'dest': dest_var,
                'left': left_var,
                'right': right_var,
                'op': op
            })
            return
        
        # If we get here, we couldn't parse the line
        print(f"Error: Line {line_num}: Could not parse line: {line}")
        # Consider exiting on parse error: sys.exit(1)
    
    def generate_cpp_code(self):
        self.cpp_code = []
        self.added_implementations = set() # Reset for generation

        # Add standard includes once
        self.cpp_code.append("// Generated by MDSLCompiler")
        self.cpp_code.append("#include <opencv2/opencv.hpp>")
        self.cpp_code.append("#include <iostream>")
        self.cpp_code.append("#include <string>")
        self.cpp_code.append("#include <vector>") # Needed by some kernels potentially
        self.cpp_code.append("#include <algorithm>") # Needed by kernels
        self.cpp_code.append("#include <immintrin.h>") # Needed by kernels
        # Only include omp if kernels might use it - check kernel files or add flag later
        self.cpp_code.append("#include <omp.h>")
        self.cpp_code.append("")
        self.cpp_code.append("using namespace cv;")
        self.cpp_code.append("using namespace std;")
        self.cpp_code.append("")

        # --- Add Kernel Implementations ---
        self.cpp_code.append("// --- Included Kernel Implementations ---")
        for action in self.action_queue:
            if action['action'] == 'dilate':
                ksize = action['ksize']
                func_name = f"dilate_{ksize}x{ksize}_avx512"
                self.add_implementation(func_name)
            elif action['action'] == 'erode':
                ksize = action['ksize']
                func_name = f"erode_{ksize}x{ksize}_avx512"
                self.add_implementation(func_name)
        self.cpp_code.append("// --- End Kernel Implementations ---")
        self.cpp_code.append("")


        # --- Main Function ---
        self.cpp_code.append("int main(int argc, char* argv[]) {")
        
        # Check command line arguments
        """
            if (argc < 2) {
        std::cerr << "Usage: " << argv[0] << " <input_image_path>" << std::endl;
        return -1;
    }
    
    // Get input path from command line argument
        std::string inputPath = argv[1];
        """
        self.cpp_code.append("    if (argc < 2) {")
        self.cpp_code.append("        cerr << \"Usage: \" << argv[0] << \" <input_image_path>\" << endl;")
        self.cpp_code.append("        return -1;")
        self.cpp_code.append("    }")
        self.cpp_code.append("    string input_path = argv[1];")
        

        # Variable declarations
        self.cpp_code.append("    // Variable declarations")
        declared_vars = set()
        for action in self.action_queue:
            # Declare destination variables if not already declared
            if 'dest' in action and action['dest'] not in declared_vars:
                 # Ensure all variables involved are declared, potentially checking type later if needed
                 self.cpp_code.append(f"    Mat {action['dest']};")
                 declared_vars.add(action['dest'])
            # Also ensure source variables referenced before assignment are declared (though parser should catch this ideally)
            if 'src' in action and action['src'] not in declared_vars:
                 print(f"Internal Warning: Source variable '{action['src']}' used before guaranteed declaration point. Declaring.")
                 self.cpp_code.append(f"    Mat {action['src']};")
                 declared_vars.add(action['src'])
            if 'left' in action and action['left'] not in declared_vars:
                 print(f"Internal Warning: Left operand '{action['left']}' used before guaranteed declaration point. Declaring.")
                 self.cpp_code.append(f"    Mat {action['left']};")
                 declared_vars.add(action['left'])
            if 'right' in action and action['right'] not in declared_vars:
                 print(f"Internal Warning: Right operand '{action['right']}' used before guaranteed declaration point. Declaring.")
                 self.cpp_code.append(f"    Mat {action['right']};")
                 declared_vars.add(action['right'])


        self.cpp_code.append("")
        self.cpp_code.append("    // Processing steps")

        # Process actions
        for action in self.action_queue:
            if action['action'] == 'load':
                # Use command line argument for image path
                self.cpp_code.append(f"    {action['dest']} = imread(input_path, IMREAD_GRAYSCALE);")
                self.cpp_code.append(f"    if ({action['dest']}.empty()) {{")
                self.cpp_code.append(f"        cerr << \"Error: Could not load image from \" << input_path << std::endl;")
                self.cpp_code.append(f"        return -1;")
                self.cpp_code.append(f"    }}")

            elif action['action'] == 'save':
                # Create output filename with suffix
                """
                inputPath.substr(0, inputPath.find_last_of('.')) + "-avx-512-processed.png"
                cv::imwrite(outputPath, tmp);
                """
                self.cpp_code.append(f"    string output_path = input_path.substr(0, input_path.find_last_of('.')) + \"{action['suffix']}\";")
                self.cpp_code.append(f"    imwrite(output_path, {action['src']});")

            elif action['action'] == 'dilate':
                ksize = action['ksize']
                func_name = f"dilate_{ksize}x{ksize}_avx512"
                self.cpp_code.append(f"    {func_name}({action['src']}, {action['dest']});")


            elif action['action'] == 'erode':
                ksize = action['ksize']
                func_name = f"erode_{ksize}x{ksize}_avx512"
                self.cpp_code.append(f"    {func_name}({action['src']}, {action['dest']});")

            elif action['action'] == 'math':
                if action['op'] == '+':
                    self.cpp_code.append(f"    add({action['left']}, {action['right']}, {action['dest']});") # Use OpenCV functions for potentially better safety/optimization
                elif action['op'] == '-':
                    self.cpp_code.append(f"    subtract({action['left']}, {action['right']}, {action['dest']});")


        self.cpp_code.append("")
        self.cpp_code.append("    return 0;")
        self.cpp_code.append("}")
        self.cpp_code.append("") # Ensure newline at end of file


        # Note: Implementations are added *before* main now
        # The section below is removed as add_implementation handles it

        return "\n".join(self.cpp_code)
    
    def add_implementation(self, func_name):
        # Only add each unique implementation once
        if func_name in self.added_implementations:
            return

        filename = f"{func_name}.cpp"
        print(f"Looking for implementation file: {filename}") # Debug print

        try:
            # Directly incorporate the user's optimized AVX-512 code
            with open(filename, 'r') as f:
                code = f.read()

            self.cpp_code.append(f"// --- Start Implementation for {func_name} (from {filename}) ---")
            # Basic check to avoid redundant includes/using if possible, though anonymous namespaces help
            # This crude check might remove necessary lines if the kernel file is structured differently
            filtered_code = []
            for line in code.splitlines():
                 stripped_line = line.strip()
                 # Avoid including headers again if already included globally
                 if stripped_line.startswith("#include <opencv2") or \
                    stripped_line.startswith("#include <iostream") or \
                    stripped_line.startswith("#include <string") or \
                    stripped_line.startswith("#include <vector") or \
                    stripped_line.startswith("#include <algorithm") or \
                    stripped_line.startswith("#include <immintrin.h") or \
                    stripped_line.startswith("#include <omp.h"):
                     continue
                 # Avoid adding 'using namespace' again
                 if stripped_line == "using namespace cv;" or stripped_line == "using namespace std;":
                     continue
                 filtered_code.append(line)

            self.cpp_code.append("\n".join(filtered_code))
            self.cpp_code.append(f"// --- End Implementation for {func_name} ---")
            self.cpp_code.append("") # Add a newline for separation
            self.added_implementations.add(func_name)
            print(f"Successfully included code from {filename}")

        except FileNotFoundError:
            print(f"Error: Required implementation file '{filename}' not found for function '{func_name}'.")
            print("Please ensure the corresponding C++ file with the AVX-512 implementation exists in the same directory.")
            sys.exit(1) # Stop compilation if a required kernel is missing
        except Exception as e:
            print(f"Error reading implementation file '{filename}': {e}")
            sys.exit(1)


def main():
    if len(sys.argv) < 2:
        print("Usage: python compiler.py input.mdsl [output.cpp]")
        sys.exit(1) # Exit if not enough args

    input_file = sys.argv[1]
    output_file = sys.argv[2] if len(sys.argv) > 2 else "output.cpp"

    if not input_file.endswith(".mdsl"):
        print(f"Warning: Input file '{input_file}' does not have a .mdsl extension.")

    compiler = MDSLCompiler()

    try:
        compiler.parse_file(input_file)
        cpp_code = compiler.generate_cpp_code() # This now raises FileNotFoundError if kernels are missing

        with open(output_file, 'w') as f:
            f.write(cpp_code)

        print(f"Successfully compiled {input_file} to {output_file}")

    except FileNotFoundError as e:
         # Error is already printed in add_implementation, just exit cleanly
         # print(f"Compilation failed: {e}") # Redundant message
         sys.exit(1)
    except Exception as e:
         print(f"An unexpected error occurred during compilation: {e}")
         sys.exit(1)


if __name__ == "__main__":
    main()