diff --git a/mlir/python/mlir/dialects/linalg/opdsl/lang/config.py b/mlir/python/mlir/dialects/linalg/opdsl/lang/config.py --- a/mlir/python/mlir/dialects/linalg/opdsl/lang/config.py +++ b/mlir/python/mlir/dialects/linalg/opdsl/lang/config.py @@ -265,8 +265,8 @@ for index in collected_indices: if index.dim_def.dimname not in self.affine_state.all_dims: raise ValueError( - f"The dimension {index.dim.dimname} is not part of the iteration " - f"domain {self.affine_state.all_dims}") + f"The dimension {index.dim_def.dimname} is not part of the " + f"iteration domain {self.affine_state.all_dims}") index.resolve_dimension_name(self.affine_state) # Generate the scalar assignments (used to build a body). diff --git a/mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir b/mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir --- a/mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir +++ b/mlir/test/Dialect/Linalg/generalize-named-polymorphic-ops.mlir @@ -250,45 +250,28 @@ // ----- -func @generalize_fill_rng_2d_f32(%min: f64, %max: f64, %seed: i32, %O: tensor<16x32xf32>) -> tensor<16x32xf32> { +func @generalize_index(%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-DAG: ^{{.*}}(%[[MIN:.+]]: f64, %[[MAX:.+]]: f64, %[[SEED:.+]]: i32, %[[O:.+]]: f32 +// CHECK-LABEL: @generalize_index // CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index // CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index // CHECK-DAG: %[[IDX0_CAST:.+]] = arith.index_cast %[[IDX0]] : index to i32 // CHECK-DAG: %[[IDX1_CAST:.+]] = arith.index_cast %[[IDX1]] : index to i32 -// CHECK-DAG: %[[VAL0:.+]] = arith.addi %[[IDX0_CAST]], %[[SEED]] : i32 -// CHECK-DAG: %[[CST0:.+]] = arith.constant 1103515245 : i32 -// CHECK-DAG: %[[CST1:.+]] = arith.constant 12345 : i32 -// CHECK-DAG: %[[VAL1:.+]] = arith.muli %[[VAL0]], %[[CST0]] : i32 -// CHECK-DAG: %[[VAL2:.+]] = arith.addi %[[VAL1]], %[[CST1]] : i32 -// Skip random number computation for the second index. -// CHECK-DAG: %[[DIFF:.+]] = arith.subf %[[MAX]], %[[MIN]] : f64 -// CHECK-DAG: %[[CST2:.+]] = arith.constant 2.3283063999999999E-10 : f64 -// CHECK-DAG: %[[FACT:.+]] = arith.mulf %[[DIFF]], %[[CST2]] : f64 -// CHECK-DAG: %[[VAL4:.+]] = arith.mulf %{{.+}}, %[[FACT]] : f64 -// CHECK-DAG: %[[VAL5:.+]] = arith.addf %[[VAL4]], %[[MIN]] : f64 -// CHECK-DAG: %[[VAL6:.+]] = arith.truncf %[[VAL5]] : f64 to f32 -// CHECK-NEXT: linalg.yield %[[VAL6]] : f32 -// CHECK-NEXT: -> tensor<16x32xf32> // ----- -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> +func @generalize_const(%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_i32 -// CHECK: ^{{.*}}(%[[MIN:.+]]: f64, %[[MAX:.+]]: f64, %[[SEED:.+]]: i32, %[[O:.+]]: i32 -// Verifies floating point to integer cast. -// CHECK: %[[VAL6:.+]] = arith.fptosi %{{.+}} : f64 to i32 -// CHECK-NEXT: linalg.yield %[[VAL6]] : i32 -// CHECK-NEXT: -> tensor<16x32xi32> +// CHECK-LABEL: @generalize_const +// CHECK-DAG: %[[CST0:.+]] = arith.constant 1103515245 : i32 +// CHECK-DAG: %[[CST1:.+]] = arith.constant 12345 : i32 +// CHECK-DAG: %[[CST2:.+]] = arith.constant 2.3283063999999999E-10 : f64 // ----- diff --git a/mlir/test/python/dialects/linalg/opdsl/emit_misc.py b/mlir/test/python/dialects/linalg/opdsl/emit_misc.py --- a/mlir/test/python/dialects/linalg/opdsl/emit_misc.py +++ b/mlir/test/python/dialects/linalg/opdsl/emit_misc.py @@ -8,7 +8,7 @@ from mlir.dialects.linalg.opdsl.lang import * # This tests miscellaneous features of the emitter that are not tested by the -# matmul, convolution, or, pooling tests. The features include: +# fill, matmul, convolution, or, pooling tests. The features include: # - constant defined in the body # - fix/predefined types # - exponential functions @@ -16,28 +16,24 @@ @linalg_structured_op -def fill_rng_poly( - min=ScalarDef(F64), - max=ScalarDef(F64), - seed=ScalarDef(I32), - O=TensorDef(T, S.M, S.N, output=True)): - multiplier = TypeFn.cast_signed(I32, const(1103515245)) - increment = TypeFn.cast_signed(I32, const(12345)) - rand1 = (TypeFn.cast_signed(I32, index(D.m)) + seed) * multiplier + increment - rand2 = (TypeFn.cast_signed(I32, index(D.n)) + rand1) * multiplier + increment - inv_range = TypeFn.cast_signed(F64, const(2.3283064e-10)) - offset = TypeFn.cast_signed(F64, const(2147483647)) - scaling = (max - min) * inv_range - O[D.m, D.n] = TypeFn.cast_signed( - T, (offset + TypeFn.cast_signed(F64, rand2)) * scaling + min) +def test_const(O=TensorDef(F32, S.M, S.N, output=True)): + O[D.m, D.n] = TypeFn.cast_unsigned(F32, const(42)) + TypeFn.cast_unsigned( + F32, const(2.3283064e-10)) @linalg_structured_op -def soft_plus_poly( - I=TensorDef(T, S.M, S.N), O=TensorDef(U, S.M, S.N, output=True)): - O[D.m, D.n] = UnaryFn.log( - TypeFn.cast_signed(U, const(1.0)) + - TypeFn.cast_signed(U, UnaryFn.exp(I[D.m, D.n]))) +def test_index(O=TensorDef(I32, S.M, S.N, output=True)): + O[D.m, D.n] = TypeFn.cast_signed(I32, index(D.m)) + TypeFn.cast_signed( + I32, index(D.n)) + + +@linalg_structured_op +def elemwise_unary_poly( + I=TensorDef(T), + O=TensorDef(U, output=True), + fun=UnaryFnAttrDef(default=UnaryFn.exp), + cast=TypeFnAttrDef(default=TypeFn.cast_signed)): + O[None] = fun(cast(U, I[None])) @linalg_structured_op(op_name="custom_op_name") @@ -48,42 +44,50 @@ with Context() as ctx, Location.unknown(): module = Module.create() f32 = F32Type.get() - f64 = F64Type.get() i32 = IntegerType.get_signless(32) with InsertionPoint(module.body): - # CHECK-LABEL: @test_i32_fill_rng - # CHECK: ^{{.*}}(%[[MIN:.+]]: f64, %[[MAX:.+]]: f64, %[[SEED:.+]]: i32, %{{.*}} + # CHECK-LABEL: @test_f32_const + # CHECK-DAG: %[[CST0:.+]] = arith.constant 42 : i64 + # CHECK-DAG: %[[CST0_CAST:.+]] = arith.uitofp %[[CST0]] : i64 to f32 + # CHECK-DAG: %[[CST1:.+]] = arith.constant 2.3283063999999999E-10 : f64 + # CHECK-DAG: %[[CST1_CAST:.+]] = arith.truncf %[[CST1]] : f64 to f32 + # CHECK-DAG: %[[SUM:.+]] = arith.addf %[[CST0_CAST]], %[[CST1_CAST]] : f32 + # CHECK-NEXT: linalg.yield %[[SUM]] : f32 + @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), f32)) + def test_f32_const(init_result): + return test_const(outs=[init_result]) + + # CHECK-LABEL: @test_i32_index # CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index + # CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index # CHECK-DAG: %[[IDX0_CAST:.+]] = arith.index_cast %[[IDX0]] : index to i32 - # CHECK-DAG: %[[RND0:.+]] = arith.addi %[[IDX0_CAST]], %[[SEED]] : i32 - # CHECK-DAG: %[[CST0:.+]] = arith.constant 1103515245 : i64 - # CHECK-DAG: %[[CST0_CAST:.+]] = arith.trunci %[[CST0]] : i64 to i32 - # Skip the remaining random number computation and match the scaling logic. - # CHECK-DAG: %[[DIFF:.+]] = arith.subf %[[MAX]], %[[MIN]] : f64 - # CHECK-DAG: %[[CST3:.+]] = arith.constant 2.3283063999999999E-10 : f64 - # CHECK-DAG: %[[FACT:.+]] = arith.mulf %[[DIFF]], %[[CST3]] : f64 - # CHECK-DAG: %[[RND4:.+]] = arith.mulf %{{.+}}, %[[FACT]] : f64 - # CHECK-DAG: %[[RND5:.+]] = arith.addf %[[RND4]], %[[MIN]] : f64 - # CHECK-DAG: %{{.*}} = arith.fptosi %[[RND5]] : f64 to i32 - @builtin.FuncOp.from_py_func(f64, f64, i32, - RankedTensorType.get((4, 16), i32)) - def test_i32_fill_rng(min, max, seed, init_result): - return fill_rng_poly(min, max, seed, outs=[init_result]) - - # CHECK-LABEL: @test_f32_soft_plus + # CHECK-DAG: %[[IDX1_CAST:.+]] = arith.index_cast %[[IDX1]] : index to i32 + # CHECK-DAG: %[[SUM:.+]] = arith.addi %[[IDX0_CAST]], %[[IDX1_CAST]] : i32 + # CHECK-NEXT: linalg.yield %[[SUM]] : i32 + @builtin.FuncOp.from_py_func(RankedTensorType.get((4, 16), i32)) + def test_i32_index(init_result): + return test_index(outs=[init_result]) + + # CHECK-LABEL: @test_f32_elemwise_exp # CHECK: ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32) - # CHECK-NEXT: %[[C1:.+]] = arith.constant 1.000000e+00 : f64 - # CHECK-NEXT: %[[C1_CAST:.+]] = arith.truncf %[[C1]] : f64 to f32 # CHECK-NEXT: %[[EXP:.+]] = math.exp %[[IN]] : f32 - # CHECK-NEXT: %[[SUM:.+]] = arith.addf %[[C1_CAST]], %[[EXP]] : f32 - # CHECK-NEXT: %[[LOG:.+]] = math.log %[[SUM]] : f32 + # CHECK-NEXT: linalg.yield %[[EXP]] : f32 + # CHECK-NEXT: -> tensor<4x16xf32> + @builtin.FuncOp.from_py_func( + RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32)) + def test_f32_elemwise_exp(input, init_result): + return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.exp) + + # CHECK-LABEL: @test_f32_elemwise_log + # CHECK: ^{{.*}}(%[[IN:.+]]: f32, %[[OUT:.+]]: f32) + # CHECK-NEXT: %[[LOG:.+]] = math.log %[[IN]] : f32 # CHECK-NEXT: linalg.yield %[[LOG]] : f32 # CHECK-NEXT: -> tensor<4x16xf32> @builtin.FuncOp.from_py_func( RankedTensorType.get((4, 16), f32), RankedTensorType.get((4, 16), f32)) - def test_f32_soft_plus(input, init_result): - return soft_plus_poly(input, outs=[init_result]) + def test_f32_elemwise_log(input, init_result): + return elemwise_unary_poly(input, outs=[init_result], fun=UnaryFn.log) # Just check that we don't assert out on name mismatch. # CHECK-LABEL: @test_non_default_op_name