Update llzk-pod-to-scalar to split pods within arrays to multiple arrays

changelogs/unreleased/th__make_pod_to_scalar_split_arrays.yaml

@@ -0,0 +1,2 @@
+changed:
+  - Update `llzk-pod-to-scalar` to split pods within arrays by creating multiple arrays.

include/llzk/Dialect/POD/Transforms/TransformationPasses.td

@@ -15,7 +15,15 @@ include "llzk/Pass/PassBase.td"
 def PodToScalarPass : LLZKPass<"llzk-pod-to-scalar"> {
   let summary = "Replace PODs with scalar values";
   let description = [{
-    Replace `pod.type` values with the proper number of scalar values
+    Scalarize `pod.type` values by splitting POD-typed struct members into
+    multiple scalar members, splitting POD-typed array elements into parallel
+    arrays, then rewriting affected member accesses plus function signatures,
+    calls, and returns, and finally running POD-specific SROA + mem2reg cleanup
+    so the remaining POD storage is promoted to SSA values.
+
+    If it is necessary to scalarize both PODs and arrays, run this pass before
+    running the `-llzk-array-to-scalar` pass because that pass will not scalarize
+    array types that are within a POD type.
   }];
 }
 

include/llzk/Util/TypeHelper.h

@@ -171,6 +171,14 @@ namespace llzk {
 
 bool isDynamic(mlir::IntegerAttr intAttr);
 
+/// Flatten any array-valued element type into the dimensions of `outerArrTy`.
+///
+/// This is used when an LLZK array logically resolves to a higher-rank array even though array
+/// element types cannot themselves be arrays. The returned type keeps `outerArrTy`'s leading
+/// dimensions, appends any nested dimensions from `elementType`, and uses the innermost non-array
+/// element type as the final element type.
+array::ArrayType flattenArrayElementType(array::ArrayType outerArrTy, mlir::Type elementType);
+
 /// Compute the cardinality (i.e. number of scalar constraints) for an EmitEqualityOp type since the
 /// op can be used to constrain two same-size arrays.
 uint64_t computeEmitEqCardinality(mlir::Type type);

lib/Dialect/Polymorphic/Transforms/FlatteningPass.cpp

@@ -969,18 +969,6 @@ LogicalResult run(ModuleOp modOp, ConversionTracker &tracker) {
 
 namespace Step1B_InstantiateFunctions {
 
-/// Flatten nested array instantiations by appending any dimensions contributed by the converted
-/// element type onto the outer array. This allows wildcard element types to resolve to
-/// higher-rank arrays even though LLZK array element types cannot themselves be arrays.
-static ArrayType flattenInstantiatedArrayType(ArrayType inputTy, Type convertedElemTy) {
-  SmallVector<Attribute> mergedDims(inputTy.getDimensionSizes());
-  while (ArrayType nestedArrTy = llvm::dyn_cast<ArrayType>(convertedElemTy)) {
-    llvm::append_range(mergedDims, nestedArrTy.getDimensionSizes());
-    convertedElemTy = nestedArrTy.getElementType();
-  }
-  return ArrayType::get(convertedElemTy, mergedDims);
-}
-
 /// TypeConverter for function instantiation that replaces TypeVarType and symbolic
 /// ArrayType/StructType parameters with their concrete values determined by unification.
 class FuncInstTypeConverter : public TypeConverter {
@@ -1020,7 +1008,7 @@ class FuncInstTypeConverter : public TypeConverter {
       if (!changed && newElemTy == inputTy.getElementType()) {
         return inputTy;
       }
-      return flattenInstantiatedArrayType(
+      return flattenArrayElementType(
           inputTy.cloneWith(inputTy.getElementType(), updated), newElemTy
       );
     });

lib/Transforms/LLZKTransformationPassPipelines.cpp

@@ -49,14 +49,14 @@ void buildFullStructInliningPipelineImpl(
   }
   pm.addPass(polymorphic::createFlatteningPass(flattening));
 
-  // Run array-to-scalar first because it can split arrays within a pod
-  // but pod-to-scalar cannot split pods within an array.
-  if (arrayToScalar) {
-    pm.addPass(array::createArrayToScalarPass());
-  }
+  // Run pod-to-scalar first because it is able to split `pod.type` used as array element type
+  // (into parallel arrays) so it should be able to fully remove all `pod.type` usages.
   if (podToScalar) {
     pm.addPass(pod::createPodToScalarPass());
   }
+  if (arrayToScalar) {
+    pm.addPass(array::createArrayToScalarPass());
+  }
   // Canonicalize to remove known-condition `scf.if` regions so struct inlining
   // can link "@compute" calls to struct members.
   pm.addPass(mlir::createCanonicalizerPass());

lib/Util/TypeHelper.cpp

@@ -564,6 +564,15 @@ bool hasAffineMapAttr(Type type) {
 
 bool isDynamic(IntegerAttr intAttr) { return ShapedType::isDynamic(fromAPInt(intAttr.getValue())); }
 
+ArrayType flattenArrayElementType(ArrayType outerArrTy, Type elementType) {
+  SmallVector<Attribute> mergedDims(outerArrTy.getDimensionSizes());
+  while (ArrayType nestedArrTy = llvm::dyn_cast<ArrayType>(elementType)) {
+    llvm::append_range(mergedDims, nestedArrTy.getDimensionSizes());
+    elementType = nestedArrTy.getElementType();
+  }
+  return ArrayType::get(elementType, mergedDims);
+}
+
 uint64_t computeEmitEqCardinality(Type type) {
   struct Impl : LLZKTypeSwitch<Impl, uint64_t> {
     uint64_t caseBool(IntegerType) { return 1; }

test/Transforms/PodToScalar/array_extract_insert.llzk

@@ -0,0 +1,30 @@
+// RUN: llzk-opt -split-input-file -llzk-pod-to-scalar %s | FileCheck --enable-var-scope %s
+
+!Pair = !pod.type<[@lhs: index, @rhs: !felt.type]>
+!PairRow = !array.type<2 x !Pair>
+!PairMatrix = !array.type<2,2 x !Pair>
+module attributes {llzk.lang} {
+  function.def @extract_then_insert(%src: !PairMatrix) -> !PairMatrix {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %row = array.extract %src[%c0] : !PairMatrix
+    %dst = array.new : !PairMatrix
+    array.insert %dst[%c1] = %row : !PairMatrix, !PairRow
+    function.return %dst : !PairMatrix
+  }
+}
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    function.def @extract_then_insert
+// CHECK-SAME:    (%[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<2,2 x index>, %[[VAL_1:[0-9a-zA-Z_\.]+]]: !array.type<2,2 x !felt.type>)
+// CHECK-SAME:    -> (!array.type<2,2 x index>, !array.type<2,2 x !felt.type>) {
+// CHECK-NEXT:      %[[VAL_2:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index
+// CHECK-NEXT:      %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.constant 1 : index
+// CHECK-NEXT:      %[[VAL_4:[0-9a-zA-Z_\.]+]] = array.extract %[[VAL_0]]{{\[}}%[[VAL_2]]] : <2,2 x index>
+// CHECK-NEXT:      %[[VAL_5:[0-9a-zA-Z_\.]+]] = array.extract %[[VAL_1]]{{\[}}%[[VAL_2]]] : <2,2 x !felt.type>
+// CHECK-NEXT:      %[[VAL_6:[0-9a-zA-Z_\.]+]] = array.new : <2,2 x index>
+// CHECK-NEXT:      %[[VAL_7:[0-9a-zA-Z_\.]+]] = array.new : <2,2 x !felt.type>
+// CHECK-NEXT:      array.insert %[[VAL_6]]{{\[}}%[[VAL_3]]] = %[[VAL_4]] : <2,2 x index>, <2 x index>
+// CHECK-NEXT:      array.insert %[[VAL_7]]{{\[}}%[[VAL_3]]] = %[[VAL_5]] : <2,2 x !felt.type>, <2 x !felt.type>
+// CHECK-NEXT:      function.return %[[VAL_6]], %[[VAL_7]] : !array.type<2,2 x index>, !array.type<2,2 x !felt.type>
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }

test/Transforms/PodToScalar/array_leaf_in_pod_array.llzk

@@ -0,0 +1,78 @@
+// RUN: llzk-opt -split-input-file -llzk-pod-to-scalar %s | FileCheck --enable-var-scope %s
+
+!Elem = !pod.type<[@vals: !array.type<3 x index>, @tag: index]>
+!ElemArray = !array.type<2 x !Elem>
+module attributes {llzk.lang} {
+  function.def @pass_through(%arg: !ElemArray) -> !ElemArray {
+    function.return %arg : !ElemArray
+  }
+}
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    function.def @pass_through(
+// CHECK-SAME:    %[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<2,3 x index>,
+// CHECK-SAME:    %[[VAL_1:[0-9a-zA-Z_\.]+]]: !array.type<2 x index>
+// CHECK-SAME:    ) -> (!array.type<2,3 x index>, !array.type<2 x index>) {
+// CHECK-NEXT:      function.return %[[VAL_0]], %[[VAL_1]] : !array.type<2,3 x index>, !array.type<2 x index>
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+
+!Elem = !pod.type<[@vals: !array.type<3 x index>]>
+!ElemArray = !array.type<2 x !Elem>
+module attributes {llzk.lang} {
+  function.def @pack(
+      %a0: index, %a1: index, %a2: index, %b0: index, %b1: index, %b2: index
+  ) -> !ElemArray {
+    %vals0 = array.new %a0, %a1, %a2 : !array.type<3 x index>
+    %vals1 = array.new %b0, %b1, %b2 : !array.type<3 x index>
+    %lhs = pod.new : !Elem
+    pod.write %lhs[@vals] = %vals0 : !Elem, !array.type<3 x index>
+    %rhs = pod.new : !Elem
+    pod.write %rhs[@vals] = %vals1 : !Elem, !array.type<3 x index>
+    %arr = array.new %lhs, %rhs : !ElemArray
+    function.return %arr : !ElemArray
+  }
+}
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    function.def @pack(
+// CHECK-SAME:    %[[VAL_0:[0-9a-zA-Z_\.]+]]: index, %[[VAL_1:[0-9a-zA-Z_\.]+]]: index,
+// CHECK-SAME:    %[[VAL_2:[0-9a-zA-Z_\.]+]]: index, %[[VAL_3:[0-9a-zA-Z_\.]+]]: index,
+// CHECK-SAME:    %[[VAL_4:[0-9a-zA-Z_\.]+]]: index, %[[VAL_5:[0-9a-zA-Z_\.]+]]: index
+// CHECK-SAME:    ) -> !array.type<2,3 x index> {
+// CHECK-NEXT:      %[[VAL_6:[0-9a-zA-Z_\.]+]] = array.new %[[VAL_0]], %[[VAL_1]], %[[VAL_2]] : <3 x index>
+// CHECK-NEXT:      %[[VAL_7:[0-9a-zA-Z_\.]+]] = array.new %[[VAL_3]], %[[VAL_4]], %[[VAL_5]] : <3 x index>
+// CHECK-NEXT:      %[[VAL_8:[0-9a-zA-Z_\.]+]] = array.new  : <2,3 x index>
+// CHECK-NEXT:      %[[VAL_9:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index
+// CHECK-NEXT:      array.insert %[[VAL_8]]{{\[}}%[[VAL_9]]] = %[[VAL_6]] : <2,3 x index>, <3 x index>
+// CHECK-NEXT:      %[[VAL_10:[0-9a-zA-Z_\.]+]] = arith.constant 1 : index
+// CHECK-NEXT:      array.insert %[[VAL_8]]{{\[}}%[[VAL_10]]] = %[[VAL_7]] : <2,3 x index>, <3 x index>
+// CHECK-NEXT:      function.return %[[VAL_8]] : !array.type<2,3 x index>
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+
+!Elem = !pod.type<[@vals: !array.type<3 x index>, @tag: index]>
+!ElemArray = !array.type<2 x !Elem>
+module attributes {llzk.lang} {
+  function.def @update_first(%arr: !ElemArray, %newVals: !array.type<3 x index>) -> !ElemArray {
+    %c0 = arith.constant 0 : index
+    %elem = array.read %arr[%c0] : !ElemArray, !Elem
+    pod.write %elem[@vals] = %newVals : !Elem, !array.type<3 x index>
+    array.write %arr[%c0] = %elem : !ElemArray, !Elem
+    function.return %arr : !ElemArray
+  }
+}
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    function.def @update_first(
+// CHECK-SAME:    %[[VAL_0:[0-9a-zA-Z_\.]+]]: !array.type<2,3 x index>,
+// CHECK-SAME:    %[[VAL_1:[0-9a-zA-Z_\.]+]]: !array.type<2 x index>,
+// CHECK-SAME:    %[[VAL_2:[0-9a-zA-Z_\.]+]]: !array.type<3 x index>
+// CHECK-SAME:    ) -> (!array.type<2,3 x index>, !array.type<2 x index>) {
+// CHECK-NEXT:      %[[VAL_3:[0-9a-zA-Z_\.]+]] = arith.constant 0 : index
+// CHECK-NEXT:      %[[VAL_4:[0-9a-zA-Z_\.]+]] = array.extract %[[VAL_0]]{{\[}}%[[VAL_3]]] : <2,3 x index>
+// CHECK-NEXT:      %[[VAL_5:[0-9a-zA-Z_\.]+]] = array.read %[[VAL_1]]{{\[}}%[[VAL_3]]] : <2 x index>, index
+// CHECK-NEXT:      array.insert %[[VAL_0]]{{\[}}%[[VAL_3]]] = %[[VAL_2]] : <2,3 x index>, <3 x index>
+// CHECK-NEXT:      array.write %[[VAL_1]]{{\[}}%[[VAL_3]]] = %[[VAL_5]] : <2 x index>, index
+// CHECK-NEXT:      function.return %[[VAL_0]], %[[VAL_1]] : !array.type<2,3 x index>, !array.type<2 x index>
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }

test/Transforms/PodToScalar/array_length.llzk

@@ -0,0 +1,129 @@
+// RUN: llzk-opt -split-input-file -llzk-pod-to-scalar %s | FileCheck --enable-var-scope %s
+
+!Pair = !pod.type<[@x: index, @y: !felt.type]>
+module attributes {llzk.lang} {
+  function.def @len_multi_leaf(%arr: !array.type<2 x !Pair>, %dim: index) -> index {
+    %len = array.len %arr, %dim : !array.type<2 x !Pair>
+    function.return %len : index
+  }
+}
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    function.def @len_multi_leaf(%[[ARR_X:[0-9a-zA-Z_\.]+]]: !array.type<2 x index>, %[[ARR_Y:[0-9a-zA-Z_\.]+]]: !array.type<2 x !felt.type>, %[[DIM:[0-9a-zA-Z_\.]+]]: index) -> index {
+// CHECK-NEXT:      %[[LEN:[0-9a-zA-Z_\.]+]] = array.len %[[ARR_X]], %[[DIM]] : <2 x index>
+// CHECK-NEXT:      function.return %[[LEN]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+
+// Tests: preserve the `array.len` semantics so a dynamic dimension selection
+// cannot observe dimensions that did not exist before pod flattening.
+!Elem = !pod.type<[@vals: !array.type<3 x index>]>
+module attributes {llzk.lang} {
+  function.def @len_leaf_array_dynamic_dim(%arr: !array.type<2 x !Elem>, %dim: index) -> index {
+    %len = array.len %arr, %dim : !array.type<2 x !Elem>
+    function.return %len : index
+  }
+}
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    function.def @len_leaf_array_dynamic_dim(%[[ARR:[0-9a-zA-Z_\.]+]]: !array.type<2,3 x index>, %[[DIM:[0-9a-zA-Z_\.]+]]: index) -> index {
+// CHECK-NEXT:      %[[SHAPE:[0-9a-zA-Z_\.]+]] = array.new  : <2 x index>
+// CHECK-NEXT:      %[[LEN:[0-9a-zA-Z_\.]+]] = array.len %[[SHAPE]], %[[DIM]] : <2 x index>
+// CHECK-NEXT:      function.return %[[LEN]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+
+module attributes {llzk.lang} {
+  function.def @len_empty_leaf_static_create(%dim: index) -> index {
+    %arr = array.new : !array.type<4 x !pod.type<[]>>
+    %len = array.len %arr, %dim : !array.type<4 x !pod.type<[]>>
+    function.return %len : index
+  }
+}
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    function.def @len_empty_leaf_static_create(%[[DIM:[0-9a-zA-Z_\.]+]]: index) -> index {
+// CHECK-NEXT:      %[[ARR:[0-9a-zA-Z_\.]+]] = array.new  : <4 x index>
+// CHECK-NEXT:      %[[LEN:[0-9a-zA-Z_\.]+]] = array.len %[[ARR]], %[[DIM]] : <4 x index>
+// CHECK-NEXT:      function.return %[[LEN]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+
+#map = affine_map<()[s0] -> (s0)>
+module attributes {llzk.lang} {
+  function.def @len_empty_leaf_array(%n: index, %dim: index) -> index {
+    %arr = array.new{()[%n]} : !array.type<#map x !pod.type<[]>>
+    %len = array.len %arr, %dim : !array.type<#map x !pod.type<[]>>
+    function.return %len : index
+  }
+}
+// CHECK: #[[$MAP:[0-9a-zA-Z_\.]+]] = affine_map<()[s0] -> (s0)>
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    function.def @len_empty_leaf_array(%[[N:[0-9a-zA-Z_\.]+]]: index, %[[DIM:[0-9a-zA-Z_\.]+]]: index) -> index {
+// CHECK-NEXT:      %[[ARR:[0-9a-zA-Z_\.]+]] = array.new{()[%[[N]]]} : <#[[$MAP]] x index>
+// CHECK-NEXT:      %[[LEN:[0-9a-zA-Z_\.]+]] = array.len %[[ARR]], %[[DIM]] : <#[[$MAP]] x index>
+// CHECK-NEXT:      function.return %[[LEN]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+
+module attributes {llzk.lang} {
+  function.def @len_empty_leaf_static_arg(%arr: !array.type<4 x !pod.type<[]>>, %dim: index) -> index {
+    %len = array.len %arr, %dim : !array.type<4 x !pod.type<[]>>
+    function.return %len : index
+  }
+}
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    function.def @len_empty_leaf_static_arg(%[[ARR:[0-9a-zA-Z_\.]+]]: !array.type<4 x index>, %[[DIM:[0-9a-zA-Z_\.]+]]: index) -> index {
+// CHECK-NEXT:      %[[LEN:[0-9a-zA-Z_\.]+]] = array.len %[[ARR]], %[[DIM]] : <4 x index>
+// CHECK-NEXT:      function.return %[[LEN]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+
+#map_arg = affine_map<()[s0] -> (s0)>
+module attributes {llzk.lang} {
+  function.def @len_empty_leaf_affine_arg(
+      %arr: !array.type<#map_arg x !pod.type<[]>>, %dim: index
+  ) -> index {
+    %len = array.len %arr, %dim : !array.type<#map_arg x !pod.type<[]>>
+    function.return %len : index
+  }
+}
+// CHECK: #[[$MAP:[0-9a-zA-Z_\.]+]] = affine_map<()[s0] -> (s0)>
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    function.def @len_empty_leaf_affine_arg(%[[ARR:[0-9a-zA-Z_\.]+]]: !array.type<#[[$MAP]] x index>, %[[DIM:[0-9a-zA-Z_\.]+]]: index) -> index {
+// CHECK-NEXT:      %[[LEN:[0-9a-zA-Z_\.]+]] = array.len %[[ARR]], %[[DIM]] : <#[[$MAP]] x index>
+// CHECK-NEXT:      function.return %[[LEN]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+
+#map_call = affine_map<()[s0] -> (s0)>
+module attributes {llzk.lang} {
+  function.def @len_empty_leaf_affine_sink(
+      %arr: !array.type<#map_call x !pod.type<[]>>, %dim: index
+  ) -> index {
+    %len = array.len %arr, %dim : !array.type<#map_call x !pod.type<[]>>
+    function.return %len : index
+  }
+
+  function.def @len_empty_leaf_affine_call(
+      %arr: !array.type<#map_call x !pod.type<[]>>, %dim: index
+  ) -> index {
+    %len = function.call @len_empty_leaf_affine_sink(%arr, %dim)
+        : (!array.type<#map_call x !pod.type<[]>>, index) -> index
+    function.return %len : index
+  }
+}
+// CHECK: #[[$MAP:[0-9a-zA-Z_\.]+]] = affine_map<()[s0] -> (s0)>
+// CHECK-LABEL: module attributes {llzk.lang} {
+// CHECK-NEXT:    function.def @len_empty_leaf_affine_sink(%[[ARR0:[0-9a-zA-Z_\.]+]]: !array.type<#[[$MAP]] x index>, %[[DIM0:[0-9a-zA-Z_\.]+]]: index) -> index {
+// CHECK-NEXT:      %[[LEN0:[0-9a-zA-Z_\.]+]] = array.len %[[ARR0]], %[[DIM0]] : <#[[$MAP]] x index>
+// CHECK-NEXT:      function.return %[[LEN0]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:    function.def @len_empty_leaf_affine_call(%[[ARR1:[0-9a-zA-Z_\.]+]]: !array.type<#[[$MAP]] x index>, %[[DIM1:[0-9a-zA-Z_\.]+]]: index) -> index {
+// CHECK-NEXT:      %[[LEN1:[0-9a-zA-Z_\.]+]] = function.call @len_empty_leaf_affine_sink(%[[ARR1]], %[[DIM1]]) : (!array.type<#[[$MAP]] x index>, index) -> index
+// CHECK-NEXT:      function.return %[[LEN1]] : index
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }