[mlir][linalg] Adapt yaml codegen to support scalar parameters.

[gysit]

Summary

The patch updates the C++ yaml code generation to support scalar operands as added in https://reviews.llvm.org/D104220.

[rsuderman]

Mind rebasing this on head? Presubmits are failing due to not being on a commited CL.

[gysit]

Rebase.

[stellaraccident]

Some nits and notes for the future. I like that this enables scalars and is much simpler than where the capture work was headed.

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml

113–115

Need to verify with the rest of this patch that this change in dims is intentional.

mlir/test/python/dialects/linalg/opsrun.py

134

Note for a follow-on patch: I believe this test file should be moved to the integration test directory.

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

116

Why switch this to an std::vector? (I know there are valid reasons to do so in this module, such as with type recursion, but I'm not seeing the rationale for this case)

152–153

I believe you removed the temporary option above. Remove from this (now stale) comment?

[gysit]

Address comments.

[gysit]

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml

113–115

I swapped the shape since the input I is read transposed (the indexing map is also swapped). It has no direct effect here since we only test the yaml to c++ translation.

mlir/test/python/dialects/linalg/opsrun.py

134

@stellaraccident Would you move it under a new integration test directory in the python subfolder (e.g. mlir/python/test/integration) or in the existing integration test directory (e.g. mlir/test/integration). I think having all python code in the same subfolder (mlir/python/) seems preferable?

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

116

When using SmallVector<ScalarAssign> I ran into a problem with a static assert checking the size (on Windows only). Setting a fixed size resolved the problem and may overall be more consistent. I just thought std::vector is more on the safe side since struct is big (close to 256 Byte). But assuming the parent struct is most likely heap allocated I think using inline storage is fine.

[stellaraccident]

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

116

Thanks and sg. It is fine to leave it as vector given this. Was just wondering the rationale for the change.

[gysit]

Address comment.

mlir/docs/Tools/LinalgOpDsl.md

 tree, MLIR must be build with Python bindings enabled
 (`-DMLIR_ENALBE_BINDINGS_PYTHON=ON`). Then add the `python` directory in the
 build tree to your `PYTHONPATH` environment variable (i.e.
 `export PYTHONPATH=$PWD/build/python`). Optionally, use an installed MLIR
 package, if available, to avoid building.
 
 ```shell
 # Dump the `core_named_ops.py` module as YAML.
-python -m python -m mlir.tools.linalg_opdsl.dump_oplib .ops.core_named_ops
+python -m mlir.dialects.linalg.opdsl.dump_oplib .ops.core_named_ops
 ```
 
 The tool is meant for use during both development and runtime, but not as
 a build tool of the core compiler: in order to export static named op
 definitions to be built as part of the compiler, the corresponding Linalg
 dialect YAML file must be updated and reviewed. TODO: Develop a script to
 automate op updates to these files.
 

mlir/include/mlir/Dialect/Linalg/IR/LinalgNamedStructuredOps.yaml

 --- !LinalgOpConfig
 metadata: !LinalgOpMetadata
   name: matmul
   cpp_class_name: MatmulOp
   doc: |-
     Performs a matrix multiplication of two 2D inputs.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
     them to the same data type as the accumulator/output.
   implements:
   - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
-  - !<LinalgTensorDef>
+  - !LinalgOperandDefConfig
     name: A
     usage: input
     shape: affine_map<()[s0, s1, s2] -> (s0, s2)>
-    element_type_var: T1
-  - !<LinalgTensorDef>
+    type_var: T1
+  - !LinalgOperandDefConfig
     name: B
     usage: input
     shape: affine_map<()[s0, s1, s2] -> (s2, s1)>
-    element_type_var: T2
-  - !<LinalgTensorDef>
+    type_var: T2
+  - !LinalgOperandDefConfig
     name: C
     usage: output
     shape: affine_map<()[s0, s1, s2] -> (s0, s1)>
-    element_type_var: U
+    type_var: U
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d2)>
     - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d2, d1)>
     - affine_map<(d0, d1, d2)[s0, s1, s2] -> (d0, d1)>
   iterator_types:
   - parallel
   - parallel
doc: |-
     Performs a batched matrix multiplication of two 3D inputs.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
     them to the same data type as the accumulator/output.
   implements:
   - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
-  - !<LinalgTensorDef>
+  - !LinalgOperandDefConfig
     name: A
     usage: input
     shape: affine_map<()[s0, s1, s2, s3] -> (s0, s1, s3)>
-    element_type_var: T1
-  - !<LinalgTensorDef>
+    type_var: T1
+  - !LinalgOperandDefConfig
     name: B
     usage: input
     shape: affine_map<()[s0, s1, s2, s3] -> (s0, s3, s2)>
-    element_type_var: T2
-  - !<LinalgTensorDef>
+    type_var: T2
+  - !LinalgOperandDefConfig
     name: C
     usage: output
     shape: affine_map<()[s0, s1, s2, s3] -> (s0, s1, s2)>
-    element_type_var: U
+    type_var: U
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0, d1, d3)>
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0, d3, d2)>
     - affine_map<(d0, d1, d2, d3)[s0, s1, s2, s3] -> (d0, d1, d2)>
   iterator_types:
   - parallel
   - parallel
doc: |-
     Performs a matrix-vector multiplication.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
     them to the same data type as the accumulator/output.
   implements:
   - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
-  - !<LinalgTensorDef>
+  - !LinalgOperandDefConfig
     name: A
     usage: input
     shape: affine_map<()[s0, s1] -> (s0, s1)>
-    element_type_var: T1
-  - !<LinalgTensorDef>
+    type_var: T1
+  - !LinalgOperandDefConfig
     name: y
     usage: input
     shape: affine_map<()[s0, s1] -> (s1)>
-    element_type_var: T2
-  - !<LinalgTensorDef>
+    type_var: T2
+  - !LinalgOperandDefConfig
     name: x
     usage: output
     shape: affine_map<()[s0, s1] -> (s0)>
-    element_type_var: U
+    type_var: U
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
     - affine_map<(d0, d1)[s0, s1] -> (d1)>
     - affine_map<(d0, d1)[s0, s1] -> (d0)>
   iterator_types:
   - parallel
   - reduction
doc: |-
     Performs a vector-matrix multiplication.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
     them to the same data type as the accumulator/output.
   implements:
   - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
-  - !<LinalgTensorDef>
+  - !LinalgOperandDefConfig
     name: y
     usage: input
     shape: affine_map<()[s0, s1] -> (s1)>
-    element_type_var: T1
-  - !<LinalgTensorDef>
+    type_var: T1
+  - !LinalgOperandDefConfig
     name: A
     usage: input
     shape: affine_map<()[s0, s1] -> (s1, s0)>
-    element_type_var: T2
-  - !<LinalgTensorDef>
+    type_var: T2
+  - !LinalgOperandDefConfig
     name: x
     usage: output
     shape: affine_map<()[s0, s1] -> (s0)>
-    element_type_var: U
+    type_var: U
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1] -> (d1)>
     - affine_map<(d0, d1)[s0, s1] -> (d1, d0)>
     - affine_map<(d0, d1)[s0, s1] -> (d0)>
   iterator_types:
   - parallel
   - reduction
doc: |-
     Performs a dot product of two vectors to a scalar result.
 
     Numeric casting is performed on the operands to the inner multiply, promoting
     them to the same data type as the accumulator/output.
   implements:
   - LinalgContractionOpInterface
 structured_op: !LinalgStructuredOpConfig
   args:
-  - !<LinalgTensorDef>
+  - !LinalgOperandDefConfig
     name: A
     usage: input
     shape: affine_map<()[s0] -> (s0)>
-    element_type_var: T1
-  - !<LinalgTensorDef>
+    type_var: T1
+  - !LinalgOperandDefConfig
     name: B
     usage: input
     shape: affine_map<()[s0] -> (s0)>
-    element_type_var: T2
-  - !<LinalgTensorDef>
+    type_var: T2
+  - !LinalgOperandDefConfig
     name: C
     usage: output
     shape: affine_map<()[s0] -> ()>
-    element_type_var: U
+    type_var: U
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0)[s0] -> (d0)>
     - affine_map<(d0)[s0] -> (d0)>
     - affine_map<(d0)[s0] -> ()>
   iterator_types:
   - reduction
   assignments:
doc: |-
 
     The operation generations pseudo random numbers using a linear congruential
     generator. It provides no guarantees regarding the distribution of the
     generated random numbers. Instead of generating the random numbers
     sequentially, it instantiates one random number generator per data element
     and runs them in parallel. The seed operand and the indices of the data
     element seed the random number generation. The min and max operands limit
     the range of the generated random numbers.
-
-    Note: The captures are hard-coded till there is capture support on the C++
-    side.
 structured_op: !LinalgStructuredOpConfig
   args:
-  - !<LinalgTensorDef>
+  - !LinalgOperandDefConfig
+    name: min
+    usage: input
+    type_var: F64
+  - !LinalgOperandDefConfig
+    name: max
+    usage: input
+    type_var: F64
+  - !LinalgOperandDefConfig
+    name: seed
+    usage: input
+    type_var: I32
+  - !LinalgOperandDefConfig
     name: O
     usage: output
     shape: affine_map<()[s0, s1] -> (s0, s1)>
-    element_type_var: T
+    type_var: T
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
+    - affine_map<(d0, d1)[s0, s1] -> ()>
+    - affine_map<(d0, d1)[s0, s1] -> ()>
+    - affine_map<(d0, d1)[s0, s1] -> ()>
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
   iterator_types:
   - parallel
   - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
value: !ScalarExpression
                                                 operands:
                                                 - !ScalarExpression
                                                   symbolic_cast:
                                                     type_var: I32
                                                     operands:
                                                     - !ScalarExpression
                                                       scalar_index: 0
                                                 - !ScalarExpression
-                                                  symbolic_cast:
-                                                    type_var: I32
-                                                    operands:
-                                                    - !ScalarExpression
-                                                      scalar_const: '42 : i64'
+                                                  scalar_arg: seed
                                             - !ScalarExpression
                                               symbolic_cast:
                                                 type_var: I32
                                                 operands:
                                                 - !ScalarExpression
                                                   scalar_const: '1103515245 : i64'
                                         - !ScalarExpression
                                           symbolic_cast:
value: !ScalarExpression
                   scalar_apply:
                     fn_name: mul
                     operands:
                     - !ScalarExpression
                       scalar_apply:
                         fn_name: sub
                         operands:
                         - !ScalarExpression
-                          symbolic_cast:
-                            type_var: F64
-                            operands:
-                            - !ScalarExpression
-                              scalar_const: '1000 : i64'
-                        - !ScalarExpression
-                          symbolic_cast:
-                            type_var: F64
-                            operands:
-                            - !ScalarExpression
-                              scalar_const: '-1000 : i64'
+                          scalar_arg: max
+                          scalar_arg: min
                     - !ScalarExpression
                       symbolic_cast:
                         type_var: F64
                         operands:
                         - !ScalarExpression
                           scalar_const: '2.3283063999999999E-10 : f64'
             - !ScalarExpression
-              symbolic_cast:
-                type_var: F64
-                operands:
-                - !ScalarExpression
-                  scalar_const: '-1000 : i64'
+              scalar_arg: min

mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir

 // CHECK:      ^{{.*}}(%[[A_ARG:.+]]: i32, %[[B_ARG:.+]]: i32, %[[C_ARG:.+]]: i32)
 // CHECK-NEXT:   %[[MUL:.+]] = muli %[[A_ARG]], %[[B_ARG]] : i32
 // CHECK-NEXT:   %[[ADD:.+]] = addi %[[C_ARG]], %[[MUL]] : i32
 // CHECK-NEXT:   linalg.yield %[[ADD]] : i32
 // CHECK-NEXT: -> tensor<16x32xi32>
 
 // -----
 
-func @generalize_fill_rng_2d_f32(%O: tensor<16x32xf32>) -> tensor<16x32xf32> {
-  %0 = linalg.fill_rng_2d outs(%O : tensor<16x32xf32>) -> tensor<16x32xf32>
+func @generalize_fill_rng_2d_f32(%min: f64, %max: f64, %seed: i32, %O: tensor<16x32xf32>) -> tensor<16x32xf32> {
+  %0 = linalg.fill_rng_2d ins(%min, %max, %seed: f64, f64, i32) outs(%O : tensor<16x32xf32>) -> tensor<16x32xf32>
   return %0: tensor<16x32xf32>
 }
 
 // CHECK-LABEL: @generalize_fill_rng_2d_f32
-// CHECK-SAME: (%[[O:.+]]: tensor<16x32xf32>)
-// CHECK-DAG:    %[[MIN:.+]] = constant -1000 : i64
-// CHECK-DAG:    %[[MAX:.+]] = constant 1000 : i64
-// CHECK-DAG:    %[[SEED:.+]] = constant 42 : i32
+// CHECK-DAG:  ^{{.*}}(%[[MIN:.+]]: f64, %[[MAX:.+]]: f64, %[[SEED:.+]]: i32, %[[O:.+]]: f32
 // CHECK-DAG:    %[[IDX0:.+]] = linalg.index 0 : index
 // CHECK-DAG:    %[[IDX1:.+]] = linalg.index 1 : index
 // CHECK-DAG:    %[[IDX0_CAST:.+]] = index_cast %[[IDX0]] : index to i32
 // CHECK-DAG:    %[[IDX1_CAST:.+]] = index_cast %[[IDX1]] : index to i32
 // CHECK-DAG:    %[[VAL0:.+]] = addi %[[IDX0_CAST]], %[[SEED]] : i32
 // CHECK-DAG:    %[[CST0:.+]] = constant 1103515245 : i32
 // CHECK-DAG:    %[[CST1:.+]] = constant 12345 : i32
 // CHECK-DAG:    %[[VAL1:.+]] = muli %[[VAL0]], %[[CST0]] : i32
 // CHECK-DAG:    %[[VAL2:.+]] = addi %[[VAL1]], %[[CST1]] : i32
 // Skip random number computation for the second index.
-// CHECK-DAG:    %[[MIN_CAST1:.+]] = sitofp %[[MIN]] : i64 to f64
-// CHECK-DAG:    %[[MAX_CAST:.+]] = sitofp %[[MAX]] : i64 to f64
-// CHECK-DAG:    %[[DIFF:.+]] = subf %[[MAX_CAST]], %[[MIN_CAST1]] : f64
+// CHECK-DAG:    %[[DIFF:.+]] = subf %[[MAX]], %[[MIN]] : f64
 // CHECK-DAG:    %[[CST2:.+]] = constant 2.3283063999999999E-10 : f64
 // CHECK-DAG:    %[[FACT:.+]] = mulf %[[DIFF]], %[[CST2]] : f64
 // CHECK-DAG:    %[[VAL4:.+]] = mulf %{{.+}}, %[[FACT]] : f64
-// CHECK-DAG:    %[[MIN_CAST2:.+]] = sitofp %[[MIN]] : i64 to f64
-// CHECK-DAG:    %[[VAL5:.+]] = addf %[[VAL4]], %[[MIN_CAST2]] : f64
+// CHECK-DAG:    %[[VAL5:.+]] = addf %[[VAL4]], %[[MIN]] : f64
 // CHECK-DAG:    %[[VAL6:.+]] = fptrunc %[[VAL5]] : f64 to f32
 // CHECK-NEXT:   linalg.yield %[[VAL6]] : f32
 // CHECK-NEXT: -> tensor<16x32xf32>
 
 // -----
 
-func @generalize_fill_rng_2d_i32(%O: tensor<16x32xi32>) -> tensor<16x32xi32> {
-  %0 = linalg.fill_rng_2d outs(%O : tensor<16x32xi32>) -> tensor<16x32xi32>
+func @generalize_fill_rng_2d_i32(%min: f64, %max: f64, %seed: i32, %O: tensor<16x32xi32>) -> tensor<16x32xi32> {
+  %0 = linalg.fill_rng_2d ins(%min, %max, %seed: f64, f64, i32) outs(%O : tensor<16x32xi32>) -> tensor<16x32xi32>
   return %0: tensor<16x32xi32>
 }
 
 // CHECK-LABEL: @generalize_fill_rng_2d_i32
-// CHECK-SAME: (%[[O:.+]]: tensor<16x32xi32>)
+// CHECK: ^{{.*}}(%[[MIN:.+]]: f64, %[[MAX:.+]]: f64, %[[SEED:.+]]: i32, %[[O:.+]]: i32
 // Verifies floating point to integer cast.
 // CHECK:        %[[VAL6:.+]] = fptosi %{{.+}} : f64 to i32
 // CHECK-NEXT:   linalg.yield %[[VAL6]] : i32
 // CHECK-NEXT: -> tensor<16x32xi32>
 
 // -----
 // Verifies floating point to integer cast.
 func @generalize_matmul_tensor_f32_f32_i16(%A : tensor<16x8xf32>, %B: tensor<8x32xf32>, %C: tensor<16x32xi16>) -> tensor<16x32xi16> {

mlir/test/mlir-linalg-ods-gen/test-linalg-ods-yaml-gen.yaml

metadata: !LinalgOpMetadata
   name: test1
   cpp_class_name: Test1Op
   doc: |-
     Title.
 
     Detailed description.
 structured_op: !LinalgStructuredOpConfig
   args:
-  - !<LinalgTensorDef>
+  - !LinalgOperandDefConfig
     name: O
     usage: output
     shape: affine_map<()[s0, s1] -> (s0, s1)>
-    element_type_var: T
+    type_var: T
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
   iterator_types:
   - parallel
   - parallel
   assignments:
   - !ScalarAssign

 # ODS-LABEL:  def Test1Op : LinalgStructuredBase_Op<"test1"
 
 #       ODS:  let summary = [{ Title. }];
 #  ODS-NEXT:  let description = [{
 #  ODS-NEXT:    Detailed description.
 #  ODS-NEXT:  }];
 
 #       ODS:  let arguments =
-#  ODS-NEXT:    Variadic<AnyShaped>:$inputs,
+#  ODS-NEXT:    Variadic<AnyType>:$inputs,
 #  ODS-NEXT:    Variadic<AnyShaped>:$outputs
 
 #       ODS:  let builders =
 #       ODS:    $_state.addOperands(inputs);
 #  ODS-NEXT:    $_state.addOperands(outputs);
 #  ODS-NEXT:    $_state.addAttribute(
 #  ODS-NEXT:      "operand_segment_sizes",
 #  ODS-NEXT:      $_builder.getI32VectorAttr({
metadata: !LinalgOpMetadata
   name: test2
   cpp_class_name: Test2Op
   doc: |-
     Title.
 
     Detailed description.
 structured_op: !LinalgStructuredOpConfig
   args:
-  - !<LinalgTensorDef>
+  - !LinalgOperandDefConfig
+    name: value
+    usage: input
+    type_var: T
+  - !LinalgOperandDefConfig
     name: I
     usage: input
-    shape: affine_map<()[s0, s1] -> (s0, s1)>
-    element_type_var: T
-  - !<LinalgTensorDef>
+    shape: affine_map<()[s0, s1] -> (s1, s0)>
+    type_var: T
+  - !LinalgOperandDefConfig
     name: O
     usage: output
     shape: affine_map<()[s0, s1] -> (s0, s1)>
-    element_type_var: T
+    type_var: T
   indexing_maps: !LinalgIndexingMapsConfig
     static_indexing_maps:
+    - affine_map<(d0, d1)[s0, s1] -> ()>
     - affine_map<(d0, d1)[s0, s1] -> (d1, d0)>
     - affine_map<(d0, d1)[s0, s1] -> (d0, d1)>
   iterator_types:
   - parallel
   - parallel
   assignments:
   - !ScalarAssign
     arg: O
     value: !ScalarExpression
+      scalar_apply:
+        fn_name: add
+        operands:
+        - !ScalarExpression
+          scalar_arg: value
+        - !ScalarExpression
           scalar_arg: I
 
 # IMPL-LABEL:  Test2Op::iterator_types()
 #  IMPL-NEXT:  { getParallelIteratorTypeName(), getParallelIteratorTypeName() }
 
 #       IMPL:  Test2Op::indexing_maps()
+#       IMPL:  "affine_map<(d0, d1)[s0, s1] -> ()>"
 #       IMPL:  "affine_map<(d0, d1)[s0, s1] -> (d1, d0)>"
 #       IMPL:  "affine_map<(d0, d1)[s0, s1] -> (d0, d1)>"
 
 #       IMPL:  void Test2Op::regionBuilder(
 #       IMPL:    ImplicitLocOpBuilder &b, Block &block, ValueRange captures)
-#       IMPL:  yields.push_back(block.getArgument(0));
+
+#       IMPL:   = helper.applyfn__add(block.getArgument(0), block.getArgument(1));

mlir/test/python/dialects/linalg/opsrun.py

with Context() as ctx, Location.unknown():
 
     log("RESULT: ", res[0])
     # CHECK: RESULT: 32.0
 
 
 test_matmul_generic()
 
 
+def test_fill_builtin():
+  with Context() as ctx, Location.unknown():
+    module = Module.create()
+    f64 = F64Type.get()
+    i32 = IntegerType.get_signless(32)
+    with InsertionPoint(module.body):
+
+      @builtin.FuncOp.from_py_func(f64, f64, i32, MemRefType.get((4, 16), i32))
+      def fill_on_buffers(min, max, seed, out):
+        linalg.fill_rng_2d(min, max, seed, outs=[out])
+
+    execution_engine = ExecutionEngine(transform(module, fill_boiler))
+
+    # TODO: FFI-based solution to allow testing and printing with python code.
+    # Prepare arguments: one result i32.
+    # Arguments must be passed as pointers.
+    c_int_p = ctypes.c_int * 1
+    res = c_int_p(-1)
+    execution_engine.invoke("main", res)
+
+    log("RESULT: ", res[0])
+    # CHECK: RESULT: -480
+
+
+test_fill_builtin()
+
+
 def test_fill_generic():
   with Context() as ctx, Location.unknown():
     module = Module.create()
     f64 = F64Type.get()
     i32 = IntegerType.get_signless(32)
     with InsertionPoint(module.body):
 
       @builtin.FuncOp.from_py_func(f64, f64, i32, MemRefType.get((4, 16), i32))

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp

 };
 
 struct SerializedAffineMap {
   AffineMapAttr affineMapAttr;
 
   AffineMap affineMap() { return affineMapAttr.getValue(); }
 };
 
-enum class LinalgTensorUsageDef {
-  input,
-  output,
-  temporary,
-};
+enum class LinalgOperandDefUsage { input, output };
 
-struct LinalgTensorDef {
+struct LinalgOperandDef {
   std::string name;
-  LinalgTensorUsageDef usage;
-  SerializedAffineMap shape;
-  std::string elementTypeVar;
+  LinalgOperandDefUsage usage;
+  Optional<SerializedAffineMap> shape;
+  std::string typeVar;
 };
 
 enum class LinalgIteratorTypeDef {
   parallel,
   reduction,
 };
 
 struct LinalgIndexingMapsConfig {

 };
 
 struct ScalarAssign {
   std::string arg;
   ScalarExpression value;
 };
 
 struct LinalgStructuredOpConfig {
-  SmallVector<LinalgTensorDef> args;
+  SmallVector<LinalgOperandDef> args;
   LinalgIndexingMapsConfig indexingMaps;
   SmallVector<LinalgIteratorTypeDef> iteratorTypes;
-  SmallVector<ScalarAssign, 2> assignments;
+  std::vector<ScalarAssign> assignments;
 };
 
 struct LinalgOpConfig {
   Optional<LinalgOpMetadata> metadata;
   Optional<LinalgStructuredOpConfig> structuredOp;
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Mapping traits.
 //===----------------------------------------------------------------------===//
 
-LLVM_YAML_IS_SEQUENCE_VECTOR(LinalgTensorDef)
+LLVM_YAML_IS_SEQUENCE_VECTOR(LinalgOperandDef)
 LLVM_YAML_IS_SEQUENCE_VECTOR(SerializedAffineMap)
 LLVM_YAML_IS_SEQUENCE_VECTOR(LinalgIteratorTypeDef)
 LLVM_YAML_IS_SEQUENCE_VECTOR(ScalarAssign)
 LLVM_YAML_IS_SEQUENCE_VECTOR(ScalarExpression)
 LLVM_YAML_IS_DOCUMENT_LIST_VECTOR(LinalgOpConfig)
 
 namespace llvm {
 namespace yaml {
 
 /// Top-level type containing op metadata and one of a concrete op type.
 /// Currently, the only defined op type is `structured_op` (maps to
 /// `LinalgStructuredOpConfig`).
 template <>
 struct MappingTraits<LinalgOpConfig> {
   static void mapping(IO &io, LinalgOpConfig &info) {
     io.mapOptional("metadata", info.metadata);
     io.mapOptional("structured_op", info.structuredOp);
   }
 };
 
 /// A structured op models (at most) a single contraction by modeling
-///   - A list of named arguments (`LinalgTensorDef`), which can be inputs,
-///     outputs, or temporaries.
+///   - A list of named arguments (`LinalgOperandDef`), which can be inputs or
+///     outputs.
 ///   - List of indexing maps (see `LinalgIndexingMaps`).
 ///   - Iterator types (see `LinalgIteratorTypeDef`).
 ///   - List of scalar level assignment (see `ScalarAssign`).
 template <>
 struct MappingTraits<LinalgStructuredOpConfig> {
   static void mapping(IO &io, LinalgStructuredOpConfig &info) {
     io.mapRequired("args", info.args);
     io.mapRequired("indexing_maps", info.indexingMaps);
     io.mapRequired("iterator_types", info.iteratorTypes);
     io.mapRequired("assignments", info.assignments);
   }
 };
 
-/// Maps a named tensor-argument to an operation, consisting of:
+/// Maps a named tensor- or scalar-argument to an operation, consisting of:
 ///   - `name`: Must be unique within the operation.
 ///   - `usage`: How the argument is used (input, output, etc).
-///   - `shape`: An AffineMap from all op symbols to the specific shape
-///     of this argument. Each shape must be normalized over the same list of
-///     symbols and have no dimension inputs.
-///   - `element_type_var`: The symbolic type variable that binds to the scalar
-///     element type of this TensorDef.
+///   - `shape`: An optional AffineMap from all op symbols to the shape of the
+///     argument. Only tensor-arguments have a shape. Each shape must be
+///     normalized over the same list of symbols and have no dimension inputs.
+///   - `type_var`: The symbolic type variable that binds to the element or self
+///     type of the tensor- or scalar-argument, respectively.
 template <>
-struct MappingTraits<LinalgTensorDef> {
-  static void mapping(IO &io, LinalgTensorDef &info) {
+struct MappingTraits<LinalgOperandDef> {
+  static void mapping(IO &io, LinalgOperandDef &info) {
     io.mapRequired("name", info.name);
     io.mapRequired("usage", info.usage);
-    io.mapRequired("shape", info.shape);
-    io.mapRequired("element_type_var", info.elementTypeVar);
+    io.mapOptional("shape", info.shape);
+    io.mapRequired("type_var", info.typeVar);
   }
 };
 
 /// Usage enum for a named argument.
 template <>
-struct ScalarEnumerationTraits<LinalgTensorUsageDef> {
-  static void enumeration(IO &io, LinalgTensorUsageDef &value) {
-    io.enumCase(value, "input", LinalgTensorUsageDef::input);
-    io.enumCase(value, "output", LinalgTensorUsageDef::output);
-    io.enumCase(value, "temporary", LinalgTensorUsageDef::temporary);
+struct ScalarEnumerationTraits<LinalgOperandDefUsage> {
+  static void enumeration(IO &io, LinalgOperandDefUsage &value) {
+    io.enumCase(value, "input", LinalgOperandDefUsage::input);
+    io.enumCase(value, "output", LinalgOperandDefUsage::output);
   }
 };
 
 /// Iterator type enum.
 template <>
 struct ScalarEnumerationTraits<LinalgIteratorTypeDef> {
   static void enumeration(IO &io, LinalgIteratorTypeDef &value) {
     io.enumCase(value, "parallel", LinalgIteratorTypeDef::parallel);

 template <>
 struct MappingTraits<LinalgIndexingMapsConfig> {
   static void mapping(IO &io, LinalgIndexingMapsConfig &info) {
     io.mapOptional("static_indexing_maps", info.staticIndexingMaps);
   }
 };
 
 /// Models an assignment to a named output.
-///   - The `arg` name must match a named output or temporary.
+///   - The `arg` name must match a named output.
 ///   - The `value` is a scalar expression for computing the value to
 ///     assign (see `ScalarExpression`).
 template <>
 struct MappingTraits<ScalarAssign> {
   static void mapping(IO &io, ScalarAssign &info) {
     io.mapRequired("arg", info.arg);
     io.mapRequired("value", info.value);
   }
static std::string interleaveToString(Container &container,
   std::string result;
   llvm::raw_string_ostream ss(result);
   llvm::interleave(container, ss, separator);
   ss.flush();
   return result;
 }
 
 static Optional<int>
-findTensorDefArgIndex(StringRef name, SmallVectorImpl<LinalgTensorDef> &args) {
+findTensorDefArgIndex(StringRef name, SmallVectorImpl<LinalgOperandDef> &args) {
   for (auto it : llvm::enumerate(args)) {
     if (it.value().name == name)
       return it.index();
   }
   return None;
 }
 
 // Try to map the TypeVar to a predefined or an argument type.
 static Optional<std::string>
-findTypeValue(StringRef typeVar, SmallVectorImpl<LinalgTensorDef> &args) {
+findTypeValue(StringRef typeVar, SmallVectorImpl<LinalgOperandDef> &args) {
   // Handle all predefined types.
   if (typeVar == "I32")
     return std::string("helper.getIntegerType(32)");
   if (typeVar == "I64")
     return std::string("helper.getIntegerType(64)");
   if (typeVar == "F32")
     return std::string("helper.getFloat32Type()");
   if (typeVar == "F64")
     return std::string("helper.getFloat64Type()");
 
   // Search all argument types.
   for (auto it : llvm::enumerate(args)) {
-    if (it.value().elementTypeVar == typeVar)
+    if (it.value().typeVar == typeVar)
       return llvm::formatv("block.getArgument({0}).getType()", it.index())
           .str();
   }
 
   return None;
 }
 
-static ScalarAssign *
-findAssignment(StringRef name, SmallVectorImpl<ScalarAssign> &assignments) {
+static ScalarAssign *findAssignment(StringRef name,
+                                    std::vector<ScalarAssign> &assignments) {
   for (auto &assign : assignments) {
     if (assign.arg == name)
       return &assign;
   }
   return nullptr;
 }
 
 //===----------------------------------------------------------------------===//

 
 def {0} : LinalgStructuredBase_Op<"{1}", !listconcat([
   AttrSizedOperandSegments,
   DeclareOpInterfaceMethods<MemoryEffectsOpInterface>,
   SingleBlockImplicitTerminator<"YieldOp">],
   /*extraInterfaces=*/[{2}])> {
     {3}
     let arguments = (ins
-      Variadic<AnyShaped>:$inputs,
+      Variadic<AnyType>:$inputs,
       Variadic<AnyShaped>:$outputs{4}
     );
     let results = (outs Variadic<AnyRankedTensor>:$result_tensors);
     let regions = (region AnyRegion:$region);
 
     let skipDefaultBuilders = 1;
     let builders = [
       OpBuilder<
       (ins "ValueRange":$inputs, "ValueRange":$outputs),
       [{{
         $_state.addOperands(inputs);
         $_state.addOperands(outputs);
         $_state.addAttribute(
           "operand_segment_sizes",
           $_builder.getI32VectorAttr({{
             static_cast<int32_t>(inputs.size()),
             static_cast<int32_t>(outputs.size())}));
         createAndFillStructuredOpRegion<{0}>(
           $_builder,
           $_state,
           TypeRange(inputs),
-          TypeRange(outputs)/*, TODO: support captures*/);
+          TypeRange(outputs));
       }]>,
       OpBuilder<
       (ins "TypeRange":$resultTensorTypes, "ValueRange":$inputs,
             "ValueRange":$outputs),
       [{{
         $_state.addOperands(inputs);
         $_state.addOperands(outputs);
         $_state.addTypes(resultTensorTypes);
         $_state.addAttribute(
           "operand_segment_sizes",
           $_builder.getI32VectorAttr({{
             static_cast<int32_t>(inputs.size()),
             static_cast<int32_t>(outputs.size())}));
         createAndFillStructuredOpRegion<{0}>(
           $_builder,
           $_state,
           TypeRange(inputs),
-          TypeRange(outputs)/*, TODO: support captures*/);
+          TypeRange(outputs));
       }]>,
       OpBuilder<
       (ins "TypeRange":$resultTensorTypes, "ValueRange":$operands,
             CArg<"ArrayRef<NamedAttribute>", "{{}">:$attributes),
       [{{
         $_state.addOperands(operands);
         $_state.addAttributes(attributes);
         $_state.addTypes(resultTensorTypes);
         (void)$_state.addRegion();
       }]>
       {6}
     ];
     let printer = [{{ return ::printNamedStructuredOp(p, *this); }];
     let parser = [{{
-      return ::parseNamedStructuredOp<{0}>(parser, result/*TODO:, captures*/);
+      return ::parseNamedStructuredOp<{0}>(parser, result);
     }];
     let hasFolder = 1;
 
     let extraClassDeclaration = structuredOpsBaseDecls # [{{
       // Auto-generated.
       ArrayAttr iterator_types();
       ArrayAttr indexing_maps();
       static void regionBuilder(
void {0}::regionBuilder(
   helper.yieldOutputs(yields);
 }
 )FMT";
     auto &args = opConfig.structuredOp->args;
     auto &assignments = opConfig.structuredOp->assignments;
     size_t generatedAssignmentCount = 0;
     int localCounter = 0;
     SmallVector<std::string> stmts;
-    for (LinalgTensorDef &arg : args) {
-      if (arg.usage != LinalgTensorUsageDef::output &&
-          arg.usage != LinalgTensorUsageDef::temporary)
+    for (LinalgOperandDef &arg : args) {
+      if (arg.usage != LinalgOperandDefUsage::output)
         continue;
 
       // Find the assignment that correlates with the argument.
       ScalarAssign *assignment = findAssignment(arg.name, assignments);
       if (!assignment)
         return emitError(genContext.getLoc())
                << "no assignment found for output argument " << arg.name;
       ++generatedAssignmentCount;