Adding ClassOp, FieldOp, GetFieldOp to allow for a transfrom from func to class

Author: Jaddyen

mlir/include/mlir/Dialect/EmitC/IR/EmitC.td

@@ -1572,4 +1572,120 @@ def EmitC_SwitchOp : EmitC_Op<"switch", [RecursiveMemoryEffects,
   let hasVerifier = 1;
 }
 
+def EmitC_ClassOp : EmitC_Op<"class", [AutomaticAllocationScope,
+                                       IsolatedFromAbove, OpAsmOpInterface]> {
+  let summary =
+      "Represents a C++ class definition, encapsulating fields and methods.";
+
+  let description = [{
+    The `emitc.class` operation defines a C++ class, acting as a container
+    for its data fields (`emitc.variable`) and methods (`emitc.func`).
+    It creates a distinct scope, isolating its contents from the surrounding
+    MLIR region, similar to how C++ classes encapsulate their internals.
+
+    Example:
+    ```mlir
+    emitc.class @MyModelClass {
+        emitc.field @another_feature : !emitc.lvalue<!emitc.ptr<f32>>
+        emitc.field @some_feature : !emitc.lvalue<!emitc.ptr<f32>>
+        emitc.field @output_0 : !emitc.lvalue<!emitc.ptr<f32>>
+
+        emitc.func @main() attributes {tf.entry_function = {inputs = "serving_default_another_feature:0,serving_default_some_feature:0", outputs = "PartitionedCall:0"}, tf_saved_model.exported_names = ["serving_default"]} {
+          %c0 = "emitc.constant"() <{value = 0 : index}> : () -> !emitc.size_t
+
+          %some_ptr = emitc.get_field %self : @MyModelClass, @some_feature -> !emitc.ptr<f32>
+          %another_ptr = emitc.get_field %self : @MyModelClass, @another_feature -> !emitc.ptr<f32>
+          %output_ptr = emitc.get_field %self : @MyModelClass, @output_0 -> !emitc.ptr<f32>
+
+          %v1 = subscript %some_ptr[%c0] : (!emitc.ptr<f32>, !emitc.size_t) -> !emitc.lvalue<f32>
+          %v1_val = load %v1 : !emitc.lvalue<f32> -> f32
+
+          %v2 = subscript %another_ptr[%c0] : (!emitc.ptr<f32>, !emitc.size_t) -> !emitc.lvalue<f32>
+          %v2_val = load %v2 : !emitc.lvalue<f32> -> f32
+
+          %v3_val = add %v1_val, %v2_val : (f32, f32) -> f32
+
+          %output_lvalue = subscript %output_ptr[%c0] : (!emitc.ptr<f32>, !emitc.size_t) -> !emitc.lvalue<f32>
+          assign %v3_val, %output_lvalue : (f32, !emitc.lvalue<f32>) -> ()
+
+          return
+        }
+      }
+    }
+
+    ```
+  }];
+
+  let arguments = (ins SymbolNameAttr:$sym_name);
+
+  let regions = (region AnyRegion:$body);
+
+  let builders = [];
+
+  let extraClassDeclaration = [{
+    // Returns the body block containing class members and methods.
+    Block &getBlock();
+  }];
+
+  let hasCustomAssemblyFormat = 1;
+
+  let assemblyFormat = "`class` $sym_name attr-dict-with-keyword $body";
+}
+
+def EmitC_FieldOp : EmitC_Op<"field", [Symbol]> {
+  let summary = "A field within a class";
+  let description = [{
+    The `emitc.field` operation declares a named field within an `emitc.class`
+    operation. The field's type must be an EmitC type. An optional initial value can be provided.
+
+    Example with initial values:
+
+    ```mlir
+    emitc.class @MyModelClass {
+      emitc.field @another_feature : !emitc.lvalue<!emitc.ptr<f32>> = #emitc.value<0.0> : !emitc.f32
+      emitc.field @some_feature : !emitc.lvalue<!emitc.ptr<f32>> = #emitc.value<1.0> : !emitc.f32
+      emitc.field @output_0 : !emitc.lvalue<!emitc.ptr<f32>>
+    }
+    ```
+    Example without initial value:
+    ```mlir
+    emitc.class @MyModelClass {
+      emitc.field @another_feature : !emitc.lvalue<!emitc.ptr<f32>>
+    }
+    ```
+  }];
+
+  let arguments = (ins SymbolNameAttr:$sym_name, TypeAttr:$type,
+      OptionalAttr<AnyAttr>:$initial_value);
+
+  let assemblyFormat = "$sym_name `:` $type (`=` $initial_value^)? attr-dict";
+
+  let hasVerifier = 1;
+}
+
+def EmitC_GetFieldOp
+    : EmitC_Op<"get_field", [Pure, DeclareOpInterfaceMethods<
+                                       SymbolUserOpInterface>]> {
+  let summary = "Obtain access to a field within a class instance";
+  let description = [{
+     The `emitc.get_field` operation retrieves the lvalue of a
+     named field from a given class instance.
+
+     Example:
+
+     ```mlir
+     %some_ptr = emitc.get_field %self : @MyModelClass, @some_feature -> !emitc.ptr<f32>
+     %another_ptr = emitc.get_field %self : @MyModelClass, @another_feature -> !emitc.ptr<f32>
+     %output_ptr = emitc.get_field %self : @MyModelClass, @output_0 -> !emitc.ptr<f32>
+     ```
+  }];
+
+  let arguments = (ins AnyTypeOf<[EmitC_LValueType, EmitC_PointerType]>:$base,
+      FlatSymbolRefAttr:$class_name, FlatSymbolRefAttr:$field_name);
+
+  let results = (outs AnyTypeOf<[EmitC_LValueType, EmitC_PointerType]>:$result);
+  let assemblyFormat = "$base `:` type($base) $class_name `,` $field_name `->` "
+                       "type($result) attr-dict";
+}
+
 #endif // MLIR_DIALECT_EMITC_IR_EMITC

mlir/include/mlir/Dialect/EmitC/Transforms/Passes.h

@@ -15,6 +15,7 @@ namespace mlir {
 namespace emitc {
 
 #define GEN_PASS_DECL_FORMEXPRESSIONSPASS
+#define GEN_PASS_DECL_CONVERTFUNCTOCLASSPASS
 #include "mlir/Dialect/EmitC/Transforms/Passes.h.inc"
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Dialect/EmitC/Transforms/Passes.td

@@ -20,4 +20,15 @@ def FormExpressionsPass : Pass<"form-expressions"> {
   let dependentDialects = ["emitc::EmitCDialect"];
 }
 
+def ConvertFuncToClassPass
+    : Pass<"convert-emitc-func-to-class", "mlir::emitc::FuncOp"> {
+  let summary = "Convert functions to classes, using arguments as fields.";
+  let description = [{
+    This pass transforms `emitc.func` operations into `emitc.class` operations.
+    Function arguments become fields of the class, and the function body is moved
+    to a new `execute` method within the class.
+  }];
+  let dependentDialects = ["emitc::EmitCDialect"];
+}
+
 #endif // MLIR_DIALECT_EMITC_TRANSFORMS_PASSES

mlir/include/mlir/Dialect/EmitC/Transforms/Transforms.h

@@ -28,6 +28,12 @@ ExpressionOp createExpression(Operation *op, OpBuilder &builder);
 /// Populates `patterns` with expression-related patterns.
 void populateExpressionPatterns(RewritePatternSet &patterns);
 
+//===----------------------------------------------------------------------===//
+// Convert Func to Class Transform
+//===----------------------------------------------------------------------===//
+
+ClassOp createClass(FuncOp funcOp, OpBuilder &builder);
+
 } // namespace emitc
 } // namespace mlir
 

mlir/lib/Dialect/EmitC/IR/EmitC.cpp

@@ -1400,6 +1400,79 @@ void FileOp::build(OpBuilder &builder, OperationState &state, StringRef id) {
       builder.getNamedAttr("id", builder.getStringAttr(id)));
 }
 
+//===----------------------------------------------------------------------===//
+// FieldOp
+//===----------------------------------------------------------------------===//
+LogicalResult FieldOp::verify() {
+  if (!isSupportedEmitCType(getType())) {
+    return emitOpError("expected valid emitc type");
+  }
+
+  if (!getInitialValue().has_value()) {
+    return success();
+  }
+
+  Attribute initValue = getInitialValue().value();
+  // Check that the type of the initial value is compatible with the type of
+  // the global variable.
+  if (auto elementsAttr = llvm::dyn_cast<ElementsAttr>(initValue)) {
+    auto initialValueType = elementsAttr.getType();
+    if (!initialValueType) {
+      return emitOpError("initial value attribute must have a type");
+    }
+    auto fieldType = getType();
+    if (initialValueType != fieldType) {
+      if (auto lvalueType = dyn_cast<LValueType>(fieldType)) {
+        auto innerFieldType = lvalueType.getValueType();
+        if (innerFieldType != initialValueType) {
+          return emitOpError("initial value type ")
+                 << initialValueType << " is not compatible with field type "
+                 << fieldType << " its inner type " << innerFieldType;
+        }
+
+      } else {
+        return emitOpError("initial value type ")
+               << initialValueType << " is not compatible with field type "
+               << fieldType;
+      }
+    }
+  }
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// GetFieldOp
+//===----------------------------------------------------------------------===//
+LogicalResult GetFieldOp::verifySymbolUses(SymbolTableCollection &symbolTable) {
+  auto classNameAttr = getClassNameAttr();
+  auto fieldNameAttr = getFieldNameAttr();
+  if (!classNameAttr || !fieldNameAttr) {
+    return emitError("class and field name attributes are mandatory");
+  }
+  StringRef className = classNameAttr.getValue();
+  StringRef fieldName = fieldNameAttr.getValue();
+
+  auto fieldOp =
+      symbolTable.lookupNearestSymbolFrom<FieldOp>(*this, getFieldNameAttr());
+
+  if (!fieldOp) {
+    return emitOpError("field '")
+           << fieldName << "' not found in class '" << className << "'";
+  }
+
+  Type getFieldResultType = getResult().getType();
+  Type fieldType = fieldOp.getType();
+
+  if (fieldType != getFieldResultType) {
+    return emitOpError("result type ")
+           << getFieldResultType << " does not match field '" << fieldName
+           << "' type " << fieldType;
+  }
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/EmitC/Transforms/CMakeLists.txt

@@ -2,6 +2,7 @@ add_mlir_dialect_library(MLIREmitCTransforms
   Transforms.cpp
   FormExpressions.cpp
   TypeConversions.cpp
+  ConvertFuncToClass.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/EmitC/Transforms

mlir/lib/Dialect/EmitC/Transforms/ConvertFuncToClass.cpp

@@ -0,0 +1,70 @@
+//===- ConvertFuncToClass.cpp - Convert functions to classes -------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/EmitC/IR/EmitC.h"
+#include "mlir/Dialect/EmitC/Transforms/Passes.h"
+#include "mlir/Dialect/EmitC/Transforms/Transforms.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace mlir {
+namespace emitc {
+
+#define GEN_PASS_DEF_CONVERTFUNCTOCLASSPASS
+#include "mlir/Dialect/EmitC/Transforms/Passes.h.inc"
+
+namespace {
+
+struct ConvertFuncToClassPass
+    : public impl::ConvertFuncToClassPassBase<ConvertFuncToClassPass> {
+  void runOnOperation() override {
+    emitc::FuncOp funcOp = getOperation();
+    MLIRContext *context = funcOp->getContext();
+
+    // Wrap each C operator op with an expression op.
+    OpBuilder builder(context);
+    createClass(funcOp, builder);
+
+    // // Create the new function inside the class
+    // auto funcType = FunctionType::get(funcOp.getContext(),
+    // funcOp.getFunctionType().getInputs(),
+    // funcOp.getFunctionType().getResults()); auto newFuncOp =
+    // builder.create<emitc::FuncOp>(
+    //     funcOp.getLoc(),builder.getStringAttr("execute"), funcType );
+
+    // builder.createBlock(&newFuncOp.getBody());
+    // builder.setInsertionPointToStart(&newFuncOp.getBody().front());
+
+    // // 7. Remap original arguments to field pointers
+    // IRMapping mapper;
+
+    // // 8. move or clone operations from original function
+    // for (Operation &opToClone :
+    // llvm::make_early_inc_range(funcOp.getBody().front())) {
+    //     if (isa<emitc::ConstantOp>(opToClone) ||
+    //         isa<emitc::SubscriptOp>(opToClone) ||
+    //         isa<emitc::LoadOp>(opToClone) ||
+    //         isa<emitc::AddOp>(opToClone) ||
+    //         isa<emitc::AssignOp>(opToClone) ||
+    //         isa<emitc::ReturnOp>(opToClone )) {
+    //     builder.clone(opToClone, mapper);
+    //     } else  {
+    //     opToClone.emitOpError("Unsupported operation found");
+    //     }
+    // }
+    // if (funcOp->use_empty()) funcOp->erase();
+  }
+};
+
+} // namespace
+
+} // namespace emitc
+} // namespace mlir